Spaces:
Build error
Build error
File size: 4,016 Bytes
c426e13 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import random
import cv2
import numpy as np
from albumentations import DualTransform, ImageOnlyTransform
from albumentations.augmentations.crops.functional import crop
#from albumentations.augmentations.functional import crop
def isotropically_resize_image(img, size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC):
h, w = img.shape[:2]
if max(w, h) == size:
return img
if w > h:
scale = size / w
h = h * scale
w = size
else:
scale = size / h
w = w * scale
h = size
interpolation = interpolation_up if scale > 1 else interpolation_down
resized = cv2.resize(img, (int(w), int(h)), interpolation=interpolation)
return resized
class IsotropicResize(DualTransform):
def __init__(self, max_side, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC,
always_apply=False, p=1):
super(IsotropicResize, self).__init__(always_apply, p)
self.max_side = max_side
self.interpolation_down = interpolation_down
self.interpolation_up = interpolation_up
def apply(self, img, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC, **params):
return isotropically_resize_image(img, size=self.max_side, interpolation_down=interpolation_down,
interpolation_up=interpolation_up)
def apply_to_mask(self, img, **params):
return self.apply(img, interpolation_down=cv2.INTER_NEAREST, interpolation_up=cv2.INTER_NEAREST, **params)
def get_transform_init_args_names(self):
return ("max_side", "interpolation_down", "interpolation_up")
class Resize4xAndBack(ImageOnlyTransform):
def __init__(self, always_apply=False, p=0.5):
super(Resize4xAndBack, self).__init__(always_apply, p)
def apply(self, img, **params):
h, w = img.shape[:2]
scale = random.choice([2, 4])
img = cv2.resize(img, (w // scale, h // scale), interpolation=cv2.INTER_AREA)
img = cv2.resize(img, (w, h),
interpolation=random.choice([cv2.INTER_CUBIC, cv2.INTER_LINEAR, cv2.INTER_NEAREST]))
return img
class RandomSizedCropNonEmptyMaskIfExists(DualTransform):
def __init__(self, min_max_height, w2h_ratio=[0.7, 1.3], always_apply=False, p=0.5):
super(RandomSizedCropNonEmptyMaskIfExists, self).__init__(always_apply, p)
self.min_max_height = min_max_height
self.w2h_ratio = w2h_ratio
def apply(self, img, x_min=0, x_max=0, y_min=0, y_max=0, **params):
cropped = crop(img, x_min, y_min, x_max, y_max)
return cropped
@property
def targets_as_params(self):
return ["mask"]
def get_params_dependent_on_targets(self, params):
mask = params["mask"]
mask_height, mask_width = mask.shape[:2]
crop_height = int(mask_height * random.uniform(self.min_max_height[0], self.min_max_height[1]))
w2h_ratio = random.uniform(*self.w2h_ratio)
crop_width = min(int(crop_height * w2h_ratio), mask_width - 1)
if mask.sum() == 0:
x_min = random.randint(0, mask_width - crop_width + 1)
y_min = random.randint(0, mask_height - crop_height + 1)
else:
mask = mask.sum(axis=-1) if mask.ndim == 3 else mask
non_zero_yx = np.argwhere(mask)
y, x = random.choice(non_zero_yx)
x_min = x - random.randint(0, crop_width - 1)
y_min = y - random.randint(0, crop_height - 1)
x_min = np.clip(x_min, 0, mask_width - crop_width)
y_min = np.clip(y_min, 0, mask_height - crop_height)
x_max = x_min + crop_height
y_max = y_min + crop_width
y_max = min(mask_height, y_max)
x_max = min(mask_width, x_max)
return {"x_min": x_min, "x_max": x_max, "y_min": y_min, "y_max": y_max}
def get_transform_init_args_names(self):
return "min_max_height", "height", "width", "w2h_ratio"
|