Spaces:
Runtime error
Runtime error
File size: 1,625 Bytes
f949b3f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import cv2
import numpy as np
from albumentations.augmentations.geometric import functional as F
from albumentations.core.transforms_interface import DualTransform
__all__ = ["ProportionalMinScale"]
class ProportionalMinScale(DualTransform):
def __init__(
self,
width: int,
height: int,
interpolation: int = cv2.INTER_LINEAR,
always_apply: bool = False,
p: float = 1,
):
super(ProportionalMinScale, self).__init__(always_apply, p)
self.width = width
self.height = height
def apply(
self, img: np.ndarray, width: int = 256, height: int = 256, interpolation: int = cv2.INTER_LINEAR, **params):
h_img, w_img, _ = img.shape
min_side = np.min([h_img, w_img])
if (height/h_img)*w_img >= width:
if h_img == min_side:
return F.smallest_max_size(img, max_size=height, interpolation=interpolation)
else:
return F.longest_max_size(img, max_size=height, interpolation=interpolation)
if (width/w_img)*h_img >= height:
if w_img == min_side:
return F.smallest_max_size(img, max_size=width, interpolation=interpolation)
else:
return F.longest_max_size(img, max_size=width, interpolation=interpolation)
return F.longest_max_size(img, max_size=width, interpolation=interpolation)
def get_params(self):
return {"width": self.width, "height": self.height}
def get_transform_init_args_names(self):
return ("width", "height", "intepolation")
|