Ejemplo n.º 1
0
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(),
        ])
Ejemplo n.º 2
0
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)
Ejemplo n.º 3
0
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
Ejemplo n.º 6
0
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"]
Ejemplo n.º 7
0
 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)
     ])
Ejemplo n.º 8
0
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
Ejemplo n.º 9
0
    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(),
        ]
Ejemplo n.º 10
0
 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()
       ])
Ejemplo n.º 11
0
 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)
     ])
Ejemplo n.º 12
0
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
Ejemplo n.º 13
0
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)
Ejemplo n.º 14
0
    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:]
Ejemplo n.º 15
0
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
Ejemplo n.º 16
0
 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
Ejemplo n.º 17
0
    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;
Ejemplo n.º 19
0
 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)
Ejemplo n.º 20
0
def img_transform(p=1):
    return Compose([
        Resize(202, 202, interpolation=cv2.INTER_NEAREST),
        PadIfNeeded(256, 256),
        Normalize(p=1)
    ],
                   p=p)
Ejemplo n.º 21
0
    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)])
Ejemplo n.º 22
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')
Ejemplo n.º 23
0
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"]
Ejemplo n.º 24
0
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
Ejemplo n.º 25
0
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))
Ejemplo n.º 26
0
 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)
Ejemplo n.º 27
0
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
Ejemplo n.º 28
0
 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),
    ])
Ejemplo n.º 30
0
    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