import torch
import cv2
import numpy as np
import torchvision
import os
import random
from utils.misc import prepare_cooridinate_input, customRandomCrop
from datasets.build_INR_dataset import Implicit2DGenerator
import albumentations
from albumentations import Resize, RandomResizedCrop, HorizontalFlip
from torch.utils.data import DataLoader
class dataset_generator(torch.utils.data.Dataset):
def __init__(self, dataset_txt, alb_transforms, torch_transforms, opt, area_keep_thresh=0.2, mode='Train'):
super().__init__()
self.opt = opt
self.root_path = opt.dataset_path
self.mode = mode
self.alb_transforms = alb_transforms
self.torch_transforms = torch_transforms
self.kp_t = area_keep_thresh
with open(dataset_txt, 'r') as f:
self.dataset_samples = [os.path.join(self.root_path, x.strip()) for x in f.readlines()]
self.INR_dataset = Implicit2DGenerator(opt, self.mode)
def __len__(self):
return len(self.dataset_samples)
def __getitem__(self, idx):
composite_image = self.dataset_samples[idx]
if self.opt.hr_train:
if self.opt.isFullRes:
"Since in dataset preprocessing, we resize the image in HAdobe5k to a lower resolution for " \
"quick loading, we need to change the path here to that of the original resolution of HAdobe5k " \
"if `opt.isFullRes` is set to True."
composite_image = composite_image.replace("HAdobe5k", "HAdobe5kori")
real_image = '_'.join(composite_image.split('_')[:2]).replace("composite_images", "real_images") + '.jpg'
mask = '_'.join(composite_image.split('_')[:-1]).replace("composite_images", "masks") + '.png'
composite_image = cv2.imread(composite_image)
composite_image = cv2.cvtColor(composite_image, cv2.COLOR_BGR2RGB)
real_image = cv2.imread(real_image)
real_image = cv2.cvtColor(real_image, cv2.COLOR_BGR2RGB)
mask = cv2.imread(mask)
mask = mask[:, :, 0].astype(np.float32) / 255.
"""
If set `opt.hr_train` to True:
Apply multi resolution crop for HR image train. Specifically, for 1024/2048 `input_size` (not fullres),
the training phase is first to RandomResizeCrop 1024/2048 `input_size`, then to random crop a `base_size`
patch to feed in multiINR process. For inference, just resize it.
While for fullres, the RandomResizeCrop is removed and just do a random crop. For inference, just keep the size.
BTW, we implement LR and HR mixing train. I.e., the following `random.random() < 0.5`
"""
if self.opt.hr_train:
if self.mode == 'Train' and self.opt.isFullRes:
if random.random() < 0.5: # LR mix training
mixTransform = albumentations.Compose(
[
RandomResizedCrop(self.opt.base_size, self.opt.base_size, scale=(0.5, 1.0)),
HorizontalFlip()],
additional_targets={'real_image': 'image', 'object_mask': 'image'}
)
origin_fg_ratio = mask.sum() / (mask.shape[0] * mask.shape[1])
origin_bg_ratio = 1 - origin_fg_ratio
"Ensure fg and bg not disappear after transformation"
valid_augmentation = False
transform_out = None
time = 0
while not valid_augmentation:
time += 1
# There are some extreme ratio pics, this code is to avoid being hindered by them.
if time == 20:
tmp_transform = albumentations.Compose(
[Resize(self.opt.base_size, self.opt.base_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = True
else:
transform_out = mixTransform(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = check_augmented_sample(transform_out['object_mask'],
origin_fg_ratio,
origin_bg_ratio,
self.kp_t)
composite_image = transform_out['image']
real_image = transform_out['real_image']
mask = transform_out['object_mask']
else: # Padding to ensure that the original resolution can be divided by 4. This is for pixel-aligned crop.
if real_image.shape[0] < 256:
bottom_pad = 256 - real_image.shape[0]
else:
bottom_pad = (4 - real_image.shape[0] % 4) % 4
if real_image.shape[1] < 256:
right_pad = 256 - real_image.shape[1]
else:
right_pad = (4 - real_image.shape[1] % 4) % 4
composite_image = cv2.copyMakeBorder(composite_image, 0, bottom_pad, 0, right_pad,
cv2.BORDER_REPLICATE)
real_image = cv2.copyMakeBorder(real_image, 0, bottom_pad, 0, right_pad, cv2.BORDER_REPLICATE)
mask = cv2.copyMakeBorder(mask, 0, bottom_pad, 0, right_pad, cv2.BORDER_REPLICATE)
origin_fg_ratio = mask.sum() / (mask.shape[0] * mask.shape[1])
origin_bg_ratio = 1 - origin_fg_ratio
"Ensure fg and bg not disappear after transformation"
valid_augmentation = False
transform_out = None
time = 0
if self.opt.hr_train:
if self.mode == 'Train':
if not self.opt.isFullRes:
if random.random() < 0.5: # LR mix training
mixTransform = albumentations.Compose(
[
RandomResizedCrop(self.opt.base_size, self.opt.base_size, scale=(0.5, 1.0)),
HorizontalFlip()],
additional_targets={'real_image': 'image', 'object_mask': 'image'}
)
while not valid_augmentation:
time += 1
# There are some extreme ratio pics, this code is to avoid being hindered by them.
if time == 20:
tmp_transform = albumentations.Compose(
[Resize(self.opt.base_size, self.opt.base_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = True
else:
transform_out = mixTransform(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = check_augmented_sample(transform_out['object_mask'],
origin_fg_ratio,
origin_bg_ratio,
self.kp_t)
else:
while not valid_augmentation:
time += 1
# There are some extreme ratio pics, this code is to avoid being hindered by them.
if time == 20:
tmp_transform = albumentations.Compose(
[Resize(self.opt.input_size, self.opt.input_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = True
else:
transform_out = self.alb_transforms(image=composite_image, real_image=real_image,
object_mask=mask)
valid_augmentation = check_augmented_sample(transform_out['object_mask'],
origin_fg_ratio,
origin_bg_ratio,
self.kp_t)
composite_image = transform_out['image']
real_image = transform_out['real_image']
mask = transform_out['object_mask']
origin_fg_ratio = mask.sum() / (mask.shape[0] * mask.shape[1])
full_coord = prepare_cooridinate_input(mask).transpose(1, 2, 0)
tmp_transform = albumentations.Compose([Resize(self.opt.base_size, self.opt.base_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image, object_mask=mask)
compos_list = [self.torch_transforms(transform_out['image'])]
real_list = [self.torch_transforms(transform_out['real_image'])]
mask_list = [
torchvision.transforms.ToTensor()(transform_out['object_mask'][..., np.newaxis].astype(np.float32))]
coord_map_list = []
valid_augmentation = False
while not valid_augmentation:
# RSC strategy. To crop different resolutions.
transform_out, c_h, c_w = customRandomCrop([composite_image, real_image, mask, full_coord],
self.opt.base_size, self.opt.base_size)
valid_augmentation = check_hr_crop_sample(transform_out[2], origin_fg_ratio)
compos_list.append(self.torch_transforms(transform_out[0]))
real_list.append(self.torch_transforms(transform_out[1]))
mask_list.append(
torchvision.transforms.ToTensor()(transform_out[2][..., np.newaxis].astype(np.float32)))
coord_map_list.append(torchvision.transforms.ToTensor()(transform_out[3]))
coord_map_list.append(torchvision.transforms.ToTensor()(transform_out[3]))
for n in range(2):
tmp_comp = cv2.resize(composite_image, (
composite_image.shape[1] // 2 ** (n + 1), composite_image.shape[0] // 2 ** (n + 1)))
tmp_real = cv2.resize(real_image,
(real_image.shape[1] // 2 ** (n + 1), real_image.shape[0] // 2 ** (n + 1)))
tmp_mask = cv2.resize(mask, (mask.shape[1] // 2 ** (n + 1), mask.shape[0] // 2 ** (n + 1)))
tmp_coord = prepare_cooridinate_input(tmp_mask).transpose(1, 2, 0)
transform_out, c_h, c_w = customRandomCrop([tmp_comp, tmp_real, tmp_mask, tmp_coord],
self.opt.base_size // 2 ** (n + 1),
self.opt.base_size // 2 ** (n + 1), c_h, c_w)
compos_list.append(self.torch_transforms(transform_out[0]))
real_list.append(self.torch_transforms(transform_out[1]))
mask_list.append(
torchvision.transforms.ToTensor()(transform_out[2][..., np.newaxis].astype(np.float32)))
coord_map_list.append(torchvision.transforms.ToTensor()(transform_out[3]))
out_comp = compos_list
out_real = real_list
out_mask = mask_list
out_coord = coord_map_list
fg_INR_coordinates, bg_INR_coordinates, fg_INR_RGB, fg_transfer_INR_RGB, bg_INR_RGB = self.INR_dataset.generator(
self.torch_transforms, transform_out[0], transform_out[1], mask)
return {
'file_path': self.dataset_samples[idx],
'category': self.dataset_samples[idx].split("\\")[-1].split("/")[0],
'composite_image': out_comp,
'real_image': out_real,
'mask': out_mask,
'coordinate_map': out_coord,
'composite_image0': out_comp[0],
'real_image0': out_real[0],
'mask0': out_mask[0],
'coordinate_map0': out_coord[0],
'composite_image1': out_comp[1],
'real_image1': out_real[1],
'mask1': out_mask[1],
'coordinate_map1': out_coord[1],
'composite_image2': out_comp[2],
'real_image2': out_real[2],
'mask2': out_mask[2],
'coordinate_map2': out_coord[2],
'composite_image3': out_comp[3],
'real_image3': out_real[3],
'mask3': out_mask[3],
'coordinate_map3': out_coord[3],
'fg_INR_coordinates': fg_INR_coordinates,
'bg_INR_coordinates': bg_INR_coordinates,
'fg_INR_RGB': fg_INR_RGB,
'fg_transfer_INR_RGB': fg_transfer_INR_RGB,
'bg_INR_RGB': bg_INR_RGB
}
else:
if not self.opt.isFullRes:
tmp_transform = albumentations.Compose([Resize(self.opt.input_size, self.opt.input_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image, object_mask=mask)
coordinate_map = prepare_cooridinate_input(transform_out['object_mask'])
"Generate INR dataset."
mask = (torchvision.transforms.ToTensor()(
transform_out['object_mask']).squeeze() > 100 / 255.).view(-1)
mask = np.bool_(mask.numpy())
fg_INR_coordinates, bg_INR_coordinates, fg_INR_RGB, fg_transfer_INR_RGB, bg_INR_RGB = self.INR_dataset.generator(
self.torch_transforms, transform_out['image'], transform_out['real_image'], mask)
return {
'file_path': self.dataset_samples[idx],
'category': self.dataset_samples[idx].split("\\")[-1].split("/")[0],
'composite_image': self.torch_transforms(transform_out['image']),
'real_image': self.torch_transforms(transform_out['real_image']),
'mask': transform_out['object_mask'][np.newaxis, ...].astype(np.float32),
# Can automatically transfer to Tensor.
'coordinate_map': coordinate_map,
'fg_INR_coordinates': fg_INR_coordinates,
'bg_INR_coordinates': bg_INR_coordinates,
'fg_INR_RGB': fg_INR_RGB,
'fg_transfer_INR_RGB': fg_transfer_INR_RGB,
'bg_INR_RGB': bg_INR_RGB
}
else:
coordinate_map = prepare_cooridinate_input(mask)
"Generate INR dataset."
mask_tmp = (torchvision.transforms.ToTensor()(mask).squeeze() > 100 / 255.).view(-1)
mask_tmp = np.bool_(mask_tmp.numpy())
fg_INR_coordinates, bg_INR_coordinates, fg_INR_RGB, fg_transfer_INR_RGB, bg_INR_RGB = self.INR_dataset.generator(
self.torch_transforms, composite_image, real_image, mask_tmp)
return {
'file_path': self.dataset_samples[idx],
'category': self.dataset_samples[idx].split("\\")[-1].split("/")[0],
'composite_image': self.torch_transforms(composite_image),
'real_image': self.torch_transforms(real_image),
'mask': mask[np.newaxis, ...].astype(np.float32),
# Can automatically transfer to Tensor.
'coordinate_map': coordinate_map,
'fg_INR_coordinates': fg_INR_coordinates,
'bg_INR_coordinates': bg_INR_coordinates,
'fg_INR_RGB': fg_INR_RGB,
'fg_transfer_INR_RGB': fg_transfer_INR_RGB,
'bg_INR_RGB': bg_INR_RGB
}
while not valid_augmentation:
time += 1
# There are some extreme ratio pics, this code is to avoid being hindered by them.
if time == 20:
tmp_transform = albumentations.Compose([Resize(self.opt.input_size, self.opt.input_size)],
additional_targets={'real_image': 'image',
'object_mask': 'image'})
transform_out = tmp_transform(image=composite_image, real_image=real_image, object_mask=mask)
valid_augmentation = True
else:
transform_out = self.alb_transforms(image=composite_image, real_image=real_image, object_mask=mask)
valid_augmentation = check_augmented_sample(transform_out['object_mask'], origin_fg_ratio,
origin_bg_ratio,
self.kp_t)
coordinate_map = prepare_cooridinate_input(transform_out['object_mask'])
"Generate INR dataset."
mask = (torchvision.transforms.ToTensor()(transform_out['object_mask']).squeeze() > 100 / 255.).view(-1)
mask = np.bool_(mask.numpy())
fg_INR_coordinates, bg_INR_coordinates, fg_INR_RGB, fg_transfer_INR_RGB, bg_INR_RGB = self.INR_dataset.generator(
self.torch_transforms, transform_out['image'], transform_out['real_image'], mask)
return {
'file_path': self.dataset_samples[idx],
'category': self.dataset_samples[idx].split("\\")[-1].split("/")[0],
'composite_image': self.torch_transforms(transform_out['image']),
'real_image': self.torch_transforms(transform_out['real_image']),
'mask': transform_out['object_mask'][np.newaxis, ...].astype(np.float32),
# Can automatically transfer to Tensor.
'coordinate_map': coordinate_map,
'fg_INR_coordinates': fg_INR_coordinates,
'bg_INR_coordinates': bg_INR_coordinates,
'fg_INR_RGB': fg_INR_RGB,
'fg_transfer_INR_RGB': fg_transfer_INR_RGB,
'bg_INR_RGB': bg_INR_RGB
}
def check_augmented_sample(mask, origin_fg_ratio, origin_bg_ratio, area_keep_thresh):
current_fg_ratio = mask.sum() / (mask.shape[0] * mask.shape[1])
current_bg_ratio = 1 - current_fg_ratio
if current_fg_ratio < origin_fg_ratio * area_keep_thresh or current_bg_ratio < origin_bg_ratio * area_keep_thresh:
return False
return True
def check_hr_crop_sample(mask, origin_fg_ratio):
current_fg_ratio = mask.sum() / (mask.shape[0] * mask.shape[1])
if current_fg_ratio < 0.8 * origin_fg_ratio:
return False
return True