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