Spaces:
Sleeping
Sleeping
| import cv2 | |
| import random | |
| import numpy as np | |
| from PIL import Image | |
| def convert_OpenCV_to_PIL(image): | |
| return Image.fromarray(image[..., ::-1]) | |
| def convert_PIL_to_OpenCV(image): | |
| return np.asarray(image)[..., ::-1] | |
| class RandomResize: | |
| def __init__(self, min_image_size, max_image_size): | |
| self.min_image_size = min_image_size | |
| self.max_image_size = max_image_size | |
| self.modes = [Image.BICUBIC, Image.NEAREST] | |
| def __call__(self, image, mode=Image.BICUBIC): | |
| rand_image_size = random.randint(self.min_image_size, self.max_image_size) | |
| w, h = image.size | |
| if w < h: | |
| scale = rand_image_size / h | |
| else: | |
| scale = rand_image_size / w | |
| size = (int(round(w*scale)), int(round(h*scale))) | |
| if size[0] == w and size[1] == h: | |
| return image | |
| return image.resize(size, mode) | |
| class RandomResize_For_Segmentation: | |
| def __init__(self, min_image_size, max_image_size): | |
| self.min_image_size = min_image_size | |
| self.max_image_size = max_image_size | |
| self.modes = [Image.BICUBIC, Image.NEAREST] | |
| def __call__(self, data): | |
| image, mask = data['image'], data['mask'] | |
| rand_image_size = random.randint(self.min_image_size, self.max_image_size) | |
| w, h = image.size | |
| if w < h: | |
| scale = rand_image_size / h | |
| else: | |
| scale = rand_image_size / w | |
| size = (int(round(w*scale)), int(round(h*scale))) | |
| if size[0] == w and size[1] == h: | |
| pass | |
| else: | |
| data['image'] = image.resize(size, Image.BICUBIC) | |
| data['mask'] = mask.resize(size, Image.NEAREST) | |
| return data | |
| class RandomHorizontalFlip: | |
| def __init__(self): | |
| pass | |
| def __call__(self, image): | |
| if bool(random.getrandbits(1)): | |
| return image.transpose(Image.FLIP_LEFT_RIGHT) | |
| return image | |
| class RandomHorizontalFlip_For_Segmentation: | |
| def __init__(self): | |
| pass | |
| def __call__(self, data): | |
| image, mask = data['image'], data['mask'] | |
| if bool(random.getrandbits(1)): | |
| data['image'] = image.transpose(Image.FLIP_LEFT_RIGHT) | |
| data['mask'] = mask.transpose(Image.FLIP_LEFT_RIGHT) | |
| return data | |
| class Normalize: | |
| def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)): | |
| self.mean = mean | |
| self.std = std | |
| def __call__(self, image): | |
| image = np.asarray(image) | |
| norm_image = np.empty_like(image, np.float32) | |
| norm_image[..., 0] = (image[..., 0] / 255. - self.mean[0]) / self.std[0] | |
| norm_image[..., 1] = (image[..., 1] / 255. - self.mean[1]) / self.std[1] | |
| norm_image[..., 2] = (image[..., 2] / 255. - self.mean[2]) / self.std[2] | |
| return norm_image | |
| class Normalize_For_Segmentation: | |
| def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)): | |
| self.mean = mean | |
| self.std = std | |
| def __call__(self, data): | |
| image, mask = data['image'], data['mask'] | |
| image = np.asarray(image, dtype=np.float32) | |
| mask = np.asarray(mask, dtype=np.int64) | |
| norm_image = np.empty_like(image, np.float32) | |
| norm_image[..., 0] = (image[..., 0] / 255. - self.mean[0]) / self.std[0] | |
| norm_image[..., 1] = (image[..., 1] / 255. - self.mean[1]) / self.std[1] | |
| norm_image[..., 2] = (image[..., 2] / 255. - self.mean[2]) / self.std[2] | |
| data['image'] = norm_image | |
| data['mask'] = mask | |
| return data | |
| class Top_Left_Crop: | |
| def __init__(self, crop_size, channels=3): | |
| self.bg_value = 0 | |
| self.crop_size = crop_size | |
| self.crop_shape = (self.crop_size, self.crop_size, channels) | |
| def __call__(self, image): | |
| h, w, c = image.shape | |
| ch = min(self.crop_size, h) | |
| cw = min(self.crop_size, w) | |
| cropped_image = np.ones(self.crop_shape, image.dtype) * self.bg_value | |
| cropped_image[:ch, :cw] = image[:ch, :cw] | |
| return cropped_image | |
| class Top_Left_Crop_For_Segmentation: | |
| def __init__(self, crop_size, channels=3): | |
| self.bg_value = 0 | |
| self.crop_size = crop_size | |
| self.crop_shape = (self.crop_size, self.crop_size, channels) | |
| self.crop_shape_for_mask = (self.crop_size, self.crop_size) | |
| def __call__(self, data): | |
| image, mask = data['image'], data['mask'] | |
| h, w, c = image.shape | |
| ch = min(self.crop_size, h) | |
| cw = min(self.crop_size, w) | |
| cropped_image = np.ones(self.crop_shape, image.dtype) * self.bg_value | |
| cropped_image[:ch, :cw] = image[:ch, :cw] | |
| cropped_mask = np.ones(self.crop_shape_for_mask, mask.dtype) * 255 | |
| cropped_mask[:ch, :cw] = mask[:ch, :cw] | |
| data['image'] = cropped_image | |
| data['mask'] = cropped_mask | |
| return data | |
| class RandomCrop: | |
| def __init__(self, crop_size, channels=3, with_bbox=False): | |
| self.bg_value = 0 | |
| self.with_bbox = with_bbox | |
| self.crop_size = crop_size | |
| self.crop_shape = (self.crop_size, self.crop_size, channels) | |
| def get_random_crop_box(self, image): | |
| h, w, c = image.shape | |
| ch = min(self.crop_size, h) | |
| cw = min(self.crop_size, w) | |
| w_space = w - self.crop_size | |
| h_space = h - self.crop_size | |
| if w_space > 0: | |
| cont_left = 0 | |
| img_left = random.randrange(w_space + 1) | |
| else: | |
| cont_left = random.randrange(-w_space + 1) | |
| img_left = 0 | |
| if h_space > 0: | |
| cont_top = 0 | |
| img_top = random.randrange(h_space + 1) | |
| else: | |
| cont_top = random.randrange(-h_space + 1) | |
| img_top = 0 | |
| dst_bbox = { | |
| 'xmin' : cont_left, 'ymin' : cont_top, | |
| 'xmax' : cont_left+cw, 'ymax' : cont_top+ch | |
| } | |
| src_bbox = { | |
| 'xmin' : img_left, 'ymin' : img_top, | |
| 'xmax' : img_left+cw, 'ymax' : img_top+ch | |
| } | |
| return dst_bbox, src_bbox | |
| def __call__(self, image, bbox_dic=None): | |
| if bbox_dic is None: | |
| dst_bbox, src_bbox = self.get_random_crop_box(image) | |
| else: | |
| dst_bbox, src_bbox = bbox_dic['dst_bbox'], bbox_dic['src_bbox'] | |
| cropped_image = np.ones(self.crop_shape, image.dtype) * self.bg_value | |
| cropped_image[dst_bbox['ymin']:dst_bbox['ymax'], dst_bbox['xmin']:dst_bbox['xmax']] = \ | |
| image[src_bbox['ymin']:src_bbox['ymax'], src_bbox['xmin']:src_bbox['xmax']] | |
| if self.with_bbox: | |
| return cropped_image, {'dst_bbox':dst_bbox, 'src_bbox':src_bbox} | |
| else: | |
| return cropped_image | |
| class RandomCrop_For_Segmentation(RandomCrop): | |
| def __init__(self, crop_size): | |
| super().__init__(crop_size) | |
| self.crop_shape_for_mask = (self.crop_size, self.crop_size) | |
| def __call__(self, data): | |
| image, mask = data['image'], data['mask'] | |
| dst_bbox, src_bbox = self.get_random_crop_box(image) | |
| cropped_image = np.ones(self.crop_shape, image.dtype) * self.bg_value | |
| cropped_image[dst_bbox['ymin']:dst_bbox['ymax'], dst_bbox['xmin']:dst_bbox['xmax']] = \ | |
| image[src_bbox['ymin']:src_bbox['ymax'], src_bbox['xmin']:src_bbox['xmax']] | |
| cropped_mask = np.ones(self.crop_shape_for_mask, mask.dtype) * 255 | |
| cropped_mask[dst_bbox['ymin']:dst_bbox['ymax'], dst_bbox['xmin']:dst_bbox['xmax']] = \ | |
| mask[src_bbox['ymin']:src_bbox['ymax'], src_bbox['xmin']:src_bbox['xmax']] | |
| data['image'] = cropped_image | |
| data['mask'] = cropped_mask | |
| return data | |
| class Transpose: | |
| def __init__(self): | |
| pass | |
| def __call__(self, image): | |
| return image.transpose((2, 0, 1)) | |
| class Transpose_For_Segmentation: | |
| def __init__(self): | |
| pass | |
| def __call__(self, data): | |
| # h, w, c -> c, h, w | |
| data['image'] = data['image'].transpose((2, 0, 1)) | |
| return data | |
| class Resize_For_Mask: | |
| def __init__(self, size): | |
| self.size = (size, size) | |
| def __call__(self, data): | |
| mask = Image.fromarray(data['mask'].astype(np.uint8)) | |
| mask = mask.resize(self.size, Image.NEAREST) | |
| data['mask'] = np.asarray(mask, dtype=np.uint64) | |
| return data | |