def get_transforms(*, data): if data == 'train': return Compose([ LongestMaxSize(CFG.size), PadIfNeeded(min_height=CFG.size, min_width=CFG.size, border_mode=1), # Resize(CFG.size,CFG.size), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ]) elif data == 'valid': return Compose([ LongestMaxSize(CFG.size), PadIfNeeded(min_height=CFG.size, min_width=CFG.size, border_mode=1), # Resize(CFG.size,CFG.size), Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ), ToTensorV2(), ])
def get_transform(model_name): if 'TAPNet' in model_name: # transform for sequences of images is very tricky # TODO: more transforms should be adopted for better results train_transform_ops = [ PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), Normalize(p=1), # optional transforms Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) ] else: train_transform_ops = [ VerticalFlip(p=0.5), HorizontalFlip(p=0.5), PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), Normalize(p=1), # optional transforms # Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) RandomCrop(height=args.input_height, width=args.input_width, p=1), ] valid_transform_ops = [ Normalize(p=1), PadIfNeeded(min_height=args.input_height, min_width=args.input_width, p=1), # optional transforms Resize(height=args.input_height, width=args.input_width, p=1), # CenterCrop(height=args.input_height, width=args.input_width, p=1) ] return Compose(train_transform_ops, p=1), Compose(valid_transform_ops, p=1)
def pad_and_crop(pad_size, crop_size=None, resize_applied=False): if crop_size: return Compose([ PadIfNeeded(pad_size, pad_size), RandomCrop(crop_size, crop_size) ], p=1.0) else: return PadIfNeeded(pad_size, pad_size)
def __init__(self, data_root='data', split_file='', size=(256, 256), fold=0, resize=False): self.data_root = data_root pkl_data = pickle.load(open(split_file, 'rb')) if fold == -1: self.path_list = pkl_data[0]['train'] self.path_list.extend(pkl_data[0]['val']) else: self.path_list = pkl_data[fold]['train'] self.len = len(self.path_list) if resize: self.transforms = Compose([Resize(height=size[0], width=size[1], interpolation=cv2.INTER_NEAREST), ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, p=0.7, border_mode=cv2.BORDER_CONSTANT, value=0), HorizontalFlip(p=0.5), OneOf([ElasticTransform(p=1, alpha=50, sigma=30, alpha_affine=30, border_mode=cv2.BORDER_CONSTANT, value=0), OpticalDistortion(p=1, distort_limit=0.5, shift_limit=0.1, border_mode=cv2.BORDER_CONSTANT, value=0)], p=0.5), RandomGamma(gamma_limit=(80, 120), p=0.5), GaussNoise(var_limit=(0.02, 0.1), mean=0, p=0.5) ]) else: self.transforms = Compose([LongestMaxSize(max_size=max(size)), PadIfNeeded(min_height=size[0], min_width=size[1], value=0, border_mode=cv2.BORDER_CONSTANT), ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, p=0.7, border_mode=cv2.BORDER_CONSTANT, value=0), HorizontalFlip(p=0.5), OneOf([ElasticTransform(p=1, alpha=50, sigma=30, alpha_affine=30, border_mode=cv2.BORDER_CONSTANT, value=0), OpticalDistortion(p=1, distort_limit=0.5, shift_limit=0.1, border_mode=cv2.BORDER_CONSTANT, value=0)], p=0.5), RandomGamma(gamma_limit=(80, 120), p=0.5), GaussNoise(var_limit=(0.02, 0.1), mean=0, p=0.5) ])
def pad_mask_image(self, mask, image, img_id, crop_shape): composed = Compose([ PadIfNeeded(crop_shape[0], crop_shape[1], p=1), RandomCrop(crop_shape[0], crop_shape[1], p=1) ], p=1) if np.sum(mask) != 0: s = 0 tries = 0 while s == 0: # crop = composed(crop_shape[0], crop_shape[1]) croped = composed(image=image, mask=mask) image_padded = croped['image'] mask_padded = croped['mask'] # print(mask_padded.shape) s = np.sum(mask_padded) tries += 1 if tries > 5: break else: croped = composed(image=image, mask=mask) image_padded = croped['image'] mask_padded = croped['mask'] return mask_padded, image_padded
def albumentations_transforms(p=1.0, is_train=False): # Mean and standard deviation of train dataset mean = np.array([0.4914, 0.4822, 0.4465]) std = np.array([0.2023, 0.1994, 0.2010]) transforms_list = [] # Use data aug only for train data if is_train: transforms_list.extend([ PadIfNeeded(min_height=40, min_width=40, border_mode=BORDER_CONSTANT, value=mean*255.0, p=1.0), OneOf([ RandomCrop(height=32, width=32, p=0.8), CenterCrop(height=32, width=32, p=0.2), ], p=1.0), HorizontalFlip(p=0.5), CoarseDropout(max_holes=1, max_height=8, max_width=8, min_height=8, min_width=8, fill_value=mean*255.0, p=0.75), ]) transforms_list.extend([ Normalize( mean=mean, std=std, max_pixel_value=255.0, p=1.0 ), ToTensor() ]) transforms = Compose(transforms_list, p=p) return lambda img:transforms(image=np.array(img))["image"]
def __init__(self, root_path, file_list, is_test=False, is_val=False, augment=False): self.is_test = is_test self.augment = augment self.root_path = root_path self.file_list = file_list self.pad = Compose([ PadIfNeeded(p=1, min_height=PAD_SIZE, min_width=PAD_SIZE), ToTensor(), ]) original_height, original_width = 101, 101 self.augmentation = Compose([ RandomSizedCrop(min_max_height=(50, 101), height=original_height, width=original_width, p=0.9), HorizontalFlip(p=0.5), GridDistortion(p=0.8), RandomContrast(p=0.8), RandomBrightness(p=0.8), RandomGamma(p=0.8) ])
def get_transforms(phase): list_transforms = [] if phase == "train": list_transforms.extend( [ OneOf([ RandomSizedCrop(min_max_height=(50, 101), height=original_height, width=original_width, p=0.5), PadIfNeeded(min_height=original_height, min_width=original_width, p=0.5)], p=1), VerticalFlip(p=0.5), # RandomRotate90(p=0.5), OneOf([ ElasticTransform(p=0.5, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03), GridDistortion(p=0.5), OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5) ], p=0.8), CLAHE(p=0.8), RandomBrightnessContrast(p=0.8), RandomGamma(p=0.8), ] ) list_transforms.extend( [ Resize(height=int(original_height/4), width=int(original_width/4), interpolation=cv2.INTER_NEAREST), Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), p=1), ToTensor(), ] ) list_trfms = Compose(list_transforms) return list_trfms
def __init__(self): self.mean = np.array([0.4914, 0.4822, 0.4465]) self.std = np.array([0.2023, 0.1994, 0.2010]) self.transforms_elist = [ PadIfNeeded(min_height=36, min_width=36, value=self.mean * 255.0), RandomCrop(height=32, width=32, p=1.0), HorizontalFlip(p=0.5), # RandomBrightnessContrast(), # Rotate(limit=7), Cutout(num_holes=1, max_h_size=8, max_w_size=8, fill_value=self.mean * 255.0, p=0.5), ] self.transforms_test = [ Resize(32, 32), ] self.transforms_main = [ Normalize(mean=self.mean, std=self.std, max_pixel_value=255.0, p=1.0), ToTensor(), ]
def __init__(self, means, stddev, settype): self.settype = settype self.means = np.array(means) self.stddev = np.array(stddev) if self.settype == 'train': print("Train set") self.albumentation_transform = Compose([ PadIfNeeded(min_height=72, min_width=72, border_mode=1, value=list(255 * self.means), p=1.0), # RandomBrightnessContrast(always_apply=False, p=0.5, brightness_limit=(-0.40, 0.82), contrast_limit=(-0.40, 0.82), brightness_by_max=True), RandomCrop(height=64, width=64, always_apply=True, p=1.0), HorizontalFlip(always_apply=False, p=0.5), Rotate(limit=15, interpolation=1, border_mode=4, value=None, mask_value=None, always_apply=False, p=0.5), # Cutout(always_apply=True, p=1.0, num_holes=1, max_h_size=8, max_w_size=8, fill_value=list(255 * self.means)), GaussNoise(always_apply=False, p=1.0, var_limit=(60, 100)), CoarseDropout(max_holes=1, max_height=16, max_width=16, min_holes=1, min_height=8, min_width=8, fill_value=list(255 * self.means), always_apply=False, p=1.0), Normalize( mean = list(self.means), std = list(self.stddev), ), ToTensor() ]) elif self.settype == 'test': print("Test set") self.albumentation_transform = Compose([ Normalize( mean = list(self.means), std = list(self.stddev), ), ToTensor() ])
def __init__(self, width, height): self.aug = Compose([ LongestMaxSize(max_size=width if width > height else height), PadIfNeeded(min_height=height, min_width=width, border_mode=cv2.BORDER_CONSTANT) ])
def test_image(model, image_path): img_transforms = transforms.Compose([ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) size_transform = Compose([ PadIfNeeded(736, 1280) ]) crop = CenterCrop(720, 1280) img = cv2.imread(image_path + '_blur_err.png') img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_s = size_transform(image=img)['image'] img_tensor = torch.from_numpy(np.transpose(img_s / 255, (2, 0, 1)).astype('float32')) img_tensor = img_transforms(img_tensor) with torch.no_grad(): img_tensor = Variable(img_tensor.unsqueeze(0).cuda()) result_image = model(img_tensor) result_image = result_image[0].cpu().float().numpy() result_image = (np.transpose(result_image, (1, 2, 0)) + 1) / 2.0 * 255.0 result_image = crop(image=result_image)['image'] result_image = result_image.astype('uint8') # gt_image = get_gt_image(image_path) gt_image = cv2.cvtColor(cv2.imread(image_path + '_ref.png'), cv2.COLOR_BGR2RGB) _, file = os.path.split(image_path) psnr = PSNR(result_image, gt_image) pilFake = Image.fromarray(result_image) pilReal = Image.fromarray(gt_image) ssim = SSIM(pilFake).cw_ssim_value(pilReal) sample_img_names = set(["010221", "024071", "033451", "051271", "060201", "070041", "090541", "100841", "101031", "113201"]) if file[-3:] == '001' or file in sample_img_names: print('test_{}: PSNR = {} dB, SSIM = {}'.format(file, psnr, ssim)) result_image = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR) cv2.imwrite(os.path.join('./test', 'test_'+image_path[-6:]+'.png'), result_image) return psnr, ssim
def aug_pad(prob=1.0, img_size=224): return Compose([ PadIfNeeded(min_height=img_size, min_width=img_size, border_mode=cv2.BORDER_CONSTANT) ], p=prob)
def __init__(self, dir_root, dir_image='images', dir_mask='masks', is_train=True, val_rate=0.2, transform=None, img_size=(128, 128)): self.dir_root = dir_root self.dir_image = dir_image self.dir_mask = dir_mask self.is_train = is_train self.val_rate = val_rate self.transform = transform self.file_list = [] self.img_size = img_size self.aug = PadIfNeeded(p=1, min_height=128, min_width=128) for path, _, files in os.walk(os.path.join(dir_root, dir_image)): for file_ in files: self.file_list.append(file_) if self.dir_mask: cut = ceil(len(self) * val_rate) if self.is_train: self.file_list = self.file_list[:-cut] else: self.file_list = self.file_list[-cut:]
def hard_transform(image_size: int = 256, p: float = 0.5, **kwargs): """Hard augmentations (on training)""" _add_transform_default_params(kwargs) transforms = Compose([ ShiftScaleRotate( shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=cv2.BORDER_REFLECT, p=p, ), IAAPerspective(scale=(0.02, 0.05), p=p), OneOf([ HueSaturationValue(p=p), ToGray(p=p), RGBShift(p=p), ChannelShuffle(p=p), ]), RandomBrightnessContrast(brightness_limit=0.5, contrast_limit=0.5, p=p), RandomGamma(p=p), CLAHE(p=p), JpegCompression(quality_lower=50, p=p), PadIfNeeded(image_size, image_size, border_mode=cv2.BORDER_CONSTANT), ], **kwargs) return transforms
def augment(self, image, mask): aug = Compose([ OneOf([ RandomSizedCrop(min_max_height=(50, 101), height=self.out_size, width=self.out_size, p=0.5), PadIfNeeded( min_height=self.out_size, min_width=self.out_size, p=0.5) ], p=1), VerticalFlip(p=0.5), RandomRotate90(p=0.5), OneOf([ ElasticTransform(p=0.5, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03), GridDistortion(p=0.5), OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5) ], p=0.8), CLAHE(p=0.8), RandomBrightnessContrast(p=0.8), RandomGamma(p=0.8) ]) augmented = aug(image=image, mask=mask) image_heavy = augmented['image'] mask_heavy = augmented['mask'] return image_heavy, mask_heavy
def Auger(self): List_Transforms = [] if self.is_train: List_Transforms.extend([ HueSaturationValue(10, 10, 10, p=0.3), HorizontalFlip(0.3), VerticalFlip(p=0.3), # May be it not work,will rescale [0,255] -> [0.0,1.0] ToFloat(always_apply=True), ShiftScaleRotate( shift_limit=0.1, # no resizing scale_limit=0.1, rotate_limit=3, # rotate p=0.5, border_mode=cv2.BORDER_REFLECT), PadIfNeeded(self.padshape, self.padshape), ]) List_Transforms.extend([ # [0.12110683835022196, 0.1308642819666743, 0.14265566800591103] #Normalize(mean=(0.397657144,0.351649219,0.305031406),std=(0.12110683835022196, 0.1308642819666743, 0.14265566800591103)), RandomCrop(self.shape, self.shape), ToTensor(), ]) TR = Compose(List_Transforms) return TR
def gettraintransforms(self, mean, std, p=1): # Train Phase transformations albumentations_transform = Compose([ # RandomRotate90(), PadIfNeeded(72, 72, border_mode=cv2.BORDER_REFLECT, always_apply=True), RandomCrop(64, 64, True), Flip(), GaussNoise(p=0.8, mean=mean), OneOf([ MotionBlur(p=0.4), MedianBlur(blur_limit=3, p=0.2), Blur(blur_limit=3, p=0.2), ], p=0.4), ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.6), OneOf([ OpticalDistortion(p=0.8), GridDistortion(p=0.4), ], p=0.6), HueSaturationValue(hue_shift_limit=20, sat_shift_limit=0.1, val_shift_limit=0.1, p=0.6), CoarseDropout(always_apply=True, max_holes=1, min_holes=1, max_height=16, max_width=16, fill_value=(255 * .6), min_height=16, min_width=16), Normalize(mean=mean, std=std, always_apply=True), pytorch.ToTensorV2(always_apply=True), ], p=p) return albumentations_transform;
def __build_augmentator(self): return Compose( [ ShiftScaleRotate( shift_limit=0.0625, scale_limit=0.0, rotate_limit=0, p=0.3), OneOf([ RandomScale(scale_limit=0.05, interpolation=1, p=0.5), Rotate(limit=7, interpolation=1, border_mode=cv2.BORDER_CONSTANT, value=0, p=0.5) ], p=0.5), PadIfNeeded(always_apply=True, min_width=self.width, min_height=self.height), RandomCrop(width=self.width, height=self.height), OneOf( [ VerticalFlip(), # HorizontalFlip(p=0.2), ], p=0.5), # OneOf([ # RandomBrightness(limit=0.2, always_apply=False, p=0.5), # RandomContrast(), # RandomGamma() # ], p=0.7), ], p=self.p)
def img_transform(p=1): return Compose([ Resize(202, 202, interpolation=cv2.INTER_NEAREST), PadIfNeeded(256, 256), Normalize(p=1) ], p=p)
def __init__(self, train_val_set, batch_size, pred_teacher=None, temperature=None): """ constructor :param train_val_set: wether to train on 90% of the training set and use the rest as validation or use the train and test set :param batch_size: the abtch size to use :param pred_teacher: the array storing the predictions of the teacher, if None, no KD is used :param temperature: the temperature to use for training, must be specified if pred_teacher is not None """ (self.x_train, self.y_train), (self.x_test, self.y_test) = cifar10.load_data() if train_val_set: self.x_train, self.x_test, self.y_train, self.y_test = train_test_split(self.x_train, self.y_train, test_size=0.1, stratify=self.y_train, random_state=17) self.train_val_set = train_val_set self.batch_size = batch_size self.pred_teacher = pred_teacher self.temperature = temperature self.images_offsets = np.arange(self.x_train.shape[0]) # remember this so that teacher predictions are # correctly offsetted # normalizes using train_set data channel_wise_mean = np.reshape(np.array([125.3, 123.0, 113.9]), (1, 1, 1, -1)) channel_wise_std = np.reshape(np.array([63.0, 62.1, 66.7]), (1, 1, 1, -1)) self.x_train = (self.x_train - channel_wise_mean) / channel_wise_std self.x_test = (self.x_test - channel_wise_mean) / channel_wise_std # Convert class vectors to binary class matrices. self.y_train = to_categorical(self.y_train, num_classes=10) self.y_test = to_categorical(self.y_test, num_classes=10) # data augmentation self.pads = Compose([PadIfNeeded(min_height=40, min_width=40, border_mode=cv2.BORDER_REFLECT_101, p=1.0)])
def train_pipeline(cache, mask_db, path): image, mask = read_image_and_mask_cached(cache, mask_db, (101, 101), path) args = Compose([ LabelMaskBorder(), HorizontalFlip(p=0.5), OneOf([ ShiftScaleRotate(rotate_limit=15, border_mode=cv2.BORDER_REPLICATE), RandomSizedCrop(min_max_height=(70, 100), height=101, width=101) ], p=0.2), GaussNoise(p=0.2), OneOf([ RandomBrightness(limit=0.4), RandomGamma(), ], p=0.5), OneOf([Blur(), MedianBlur(), MotionBlur()], p=0.2), OneOf([ ElasticTransform(alpha=10, sigma=10, alpha_affine=10), GridDistortion() ], p=0.2), Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), PadIfNeeded(128, 128, cv2.BORDER_REPLICATE), ChannelsFirst() ])(image=image, mask=mask) return args['image'], args.get('mask')
def albumentations_transforms(p=1.0, is_train=False): # Mean and standard deviation of train dataset mean = np.array([0.4914, 0.4822, 0.4465]) std = np.array([0.2023, 0.1994, 0.2010]) transforms_list = [] # Use data aug only for train data if is_train: transforms_list.extend([ PadIfNeeded(min_height=72, min_width=72, p=1.0), RandomCrop(height=64, width=64, p=1.0), HorizontalFlip(p=0.25), Rotate(limit=15, p=0.25), RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=0.25), #CoarseDropout(max_holes=1, max_height=32, max_width=32, min_height=8, #min_width=8, fill_value=mean*255.0, p=0.5), ]) transforms_list.extend([ Normalize(mean=mean, std=std, max_pixel_value=255.0, p=1.0), ToTensor() ]) data_transforms = Compose(transforms_list, p=p) return lambda img: data_transforms(image=np.array(img))["image"]
def hard_aug(original_height=128, original_width=128, k=4): aug = Compose([ OneOf([ RandomSizedCrop( min_max_height=(original_height // k, original_height), height=original_height, width=original_width, p=0.5), PadIfNeeded( min_height=original_height, min_width=original_width, p=0.5) ], p=1), VerticalFlip(p=0.5), HorizontalFlip(p=0.5), RandomRotate90(p=0.5), Transpose(p=0.5), OneOf([ ElasticTransform( p=0.5, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03), GridDistortion(p=0.5), OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5) ], p=0.8), CLAHE(p=0.8), RandomBrightnessContrast(p=0.8), RandomGamma(p=0.8) ]) return aug
def inference(args): with open(args.config) as cfg: config = yaml.load(cfg) model = get_generator(config['model']) model.load_state_dict(torch.load(args.weights)['model']) model = model.cuda() os.makedirs(os.path.dirname(args.output), exist_ok=True) img_transforms = transforms.Compose( [transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]) size_transform = Compose([PadIfNeeded(736, 1280)]) crop = CenterCrop(720, 1280) img = cv2.imread(args.input) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_s = size_transform(image=img)['image'] img_tensor = torch.from_numpy( np.transpose(img_s / 255, (2, 0, 1)).astype('float32')) img_tensor = img_transforms(img_tensor) with torch.no_grad(): img_tensor = Variable(img_tensor.unsqueeze(0).cuda()) result_image = model(img_tensor) result_image = result_image[0].cpu().float().numpy() result_image = (np.transpose(result_image, (1, 2, 0)) + 1) / 2.0 * 255.0 result_image = crop(image=result_image)['image'] result_image = result_image.astype('uint8') cv2.imwrite(args.output, cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR))
def __getitem__(self, index): img_path = self.image_paths[index] mask_path = self.label_paths[index] if img_path.endswith(".npy"): img = np.load(img_path) else: img = cv2.imread(img_path) if mask_path.endswith(".npy"): mask = np.load(mask_path) else: mask = cv2.imread(mask_path, 0) if self.augmentation: task = [ HorizontalFlip(p=0.5), VerticalFlip(p=0.5), RandomGamma(), RandomBrightnessContrast(p=0.5), PadIfNeeded(self.img_size, self.img_size), ShiftScaleRotate(scale_limit=0.5, p=0.5), #Normalize(mean=[0.210, 0.210, 0.210], std=[0.196, 0.196, 0.196], always_apply=True) ] aug = Compose(task) aug_data = aug(image=img, mask=mask) img, mask = aug_data["image"], aug_data["mask"] img = self._normalize(img) img = cv2.resize(img, (self.img_size, self.img_size)) mask = cv2.resize(mask, (self.img_size, self.img_size)) mask = mask // 255.0 if img.ndim < 3: img = np.expand_dims(img, 0) else: img = np.transpose(img, axes=[2, 0, 1]) return torch.from_numpy(img), torch.from_numpy(mask)
def test_image(model, save_path, image_path): img_transforms = transforms.Compose( [transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]) size_transform = Compose([PadIfNeeded(736, 1280)]) crop = CenterCrop(720, 1280) img = cv2.imread(image_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_s = size_transform(image=img)['image'] img_tensor = torch.from_numpy( np.transpose(img_s / 255, (2, 0, 1)).astype('float32')) img_tensor = img_transforms(img_tensor) with torch.no_grad(): img_tensor = Variable(img_tensor.unsqueeze(0).cuda()) result_image = model(img_tensor) result_image = result_image[0].cpu().float().numpy() result_image = (np.transpose(result_image, (1, 2, 0)) + 1) / 2.0 * 255.0 result_image = crop(image=result_image)['image'] result_image = result_image.astype('uint8') gt_image = get_gt_image(image_path) lap = estimate_blur(result_image) lap_sharp = estimate_blur(gt_image) lap_blur = estimate_blur(img) _, filename = os.path.split(image_path) psnr = PSNR(result_image, gt_image) pilFake = Image.fromarray(result_image) pilReal = Image.fromarray(gt_image) ssim = SSIM(pilFake).cw_ssim_value(pilReal) #result_image = np.hstack((img_s, result_image, gt_image)) #cv2.imwrite(os.path.join(save_path, filename), cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)) return psnr, ssim, lap, lap_blur, lap_sharp
def __getitem__(self, index): path = os.path.join(self.data_dir, self.filename[index] + self.suffix) label = None if self.ids is None else self.ids[index] img = Image.open(path) img = img.convert("RGB") if self.transform is not None: img = self.transform(img) else: if self.augumentation: img = np.asarray(img) task = [ HorizontalFlip(p=0.5), VerticalFlip(p=0.5), RandomGamma(), RandomBrightnessContrast(p=0.5), PadIfNeeded(self.img_size, self.img_size), ShiftScaleRotate(scale_limit=0.5, p=0.5) ] aug = Compose(task) aug_data = aug(image=img) img = aug_data["image"] #img = cv2.resize(img,(self.img_size,self.img_size)) img = Image.fromarray(img) img = self.transform(img) if label is not None: return img, label else: return img
def create_train_transforms(size=300): return Compose([ ImageCompression(quality_lower=60, quality_upper=100, p=0.5), GaussNoise(p=0.1), GaussianBlur(blur_limit=3, p=0.05), HorizontalFlip(), OneOf([ IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC), IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_LINEAR), IsotropicResize(max_side=size, interpolation_down=cv2.INTER_LINEAR, interpolation_up=cv2.INTER_LINEAR), ], p=1), PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT), OneOf([RandomBrightnessContrast(), FancyPCA(), HueSaturationValue()], p=0.7), ToGray(p=0.2), ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5), ])
def _read_image_valid(self, id): # get the image from the validation folder img = cv2.imread(os.path.join(args.images_dir, '{}.png'.format(id)), cv2.IMREAD_COLOR) mask = cv2.imread(os.path.join(args.images_dir, '{}.png'.format(id)), cv2.IMREAD_GRAYSCALE) # convert the image into an numpy array img = np.array(img, np.float32) # Resize the image img = cv2.resize(img, (args.resize_size, args.resize_size)) mask = cv2.resize(mask, (args.resize_size, args.resize_size)) # Apply padding if needed augmentation = PadIfNeeded(min_height=self.input_shape[0], min_width=self.input_shape[1], p=1.0, border_mode=4) data = {'image': img, 'mask': mask} augmented = augmentation(**data) img, mask = augmented['image'], augmented['mask'] img = np.array(img, np.float32) #img = self.preprocess(img) mask = np.array(mask / 255., np.float32) if len(mask.shape) < 3: mask = np.expand_dims(mask, axis=2) return img, mask