221 lines
8.4 KiB
Python
221 lines
8.4 KiB
Python
import json
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
from torch.utils.data import Dataset
|
|
from PIL import Image
|
|
import cv2
|
|
from .data_utils import *
|
|
cv2.setNumThreads(0)
|
|
cv2.ocl.setUseOpenCL(False)
|
|
import albumentations as A
|
|
|
|
|
|
class BaseDataset(Dataset):
|
|
def __init__(self):
|
|
image_mask_dict = {}
|
|
self.data = []
|
|
|
|
def __len__(self):
|
|
# We adjust the ratio of different dataset by setting the length.
|
|
pass
|
|
|
|
|
|
def aug_data_back(self, image):
|
|
transform = A.Compose([
|
|
A.ColorJitter(p=0.5, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5),
|
|
A.ChannelShuffle()
|
|
])
|
|
transformed = transform(image=image.astype(np.uint8))
|
|
transformed_image = transformed["image"]
|
|
return transformed_image
|
|
|
|
def aug_data_mask(self, image, mask):
|
|
transform = A.Compose([
|
|
A.HorizontalFlip(p=0.5),
|
|
A.RandomBrightnessContrast(p=0.5),
|
|
#A.Rotate(limit=20, border_mode=cv2.BORDER_CONSTANT, value=(0,0,0)),
|
|
])
|
|
|
|
transformed = transform(image=image.astype(np.uint8), mask = mask)
|
|
transformed_image = transformed["image"]
|
|
transformed_mask = transformed["mask"]
|
|
return transformed_image, transformed_mask
|
|
|
|
|
|
def check_region_size(self, image, yyxx, ratio, mode = 'max'):
|
|
pass_flag = True
|
|
H,W = image.shape[0], image.shape[1]
|
|
H,W = H * ratio, W * ratio
|
|
y1,y2,x1,x2 = yyxx
|
|
h,w = y2-y1,x2-x1
|
|
if mode == 'max':
|
|
if h > H or w > W:
|
|
pass_flag = False
|
|
elif mode == 'min':
|
|
if h < H or w < W:
|
|
pass_flag = False
|
|
return pass_flag
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
while(True):
|
|
try:
|
|
idx = np.random.randint(0, len(self.data)-1)
|
|
item = self.get_sample(idx)
|
|
return item
|
|
except:
|
|
idx = np.random.randint(0, len(self.data)-1)
|
|
|
|
def get_sample(self, idx):
|
|
# Implemented for each specific dataset
|
|
pass
|
|
|
|
def sample_timestep(self, max_step =1000):
|
|
if np.random.rand() < 0.3:
|
|
step = np.random.randint(0,max_step)
|
|
return np.array([step])
|
|
|
|
if self.dynamic == 1:
|
|
# coarse videos
|
|
step_start = max_step // 2
|
|
step_end = max_step
|
|
elif self.dynamic == 0:
|
|
# static images
|
|
step_start = 0
|
|
step_end = max_step // 2
|
|
else:
|
|
# fine multi-view images/videos/3Ds
|
|
step_start = 0
|
|
step_end = max_step
|
|
step = np.random.randint(step_start, step_end)
|
|
return np.array([step])
|
|
|
|
def check_mask_area(self, mask):
|
|
H,W = mask.shape[0], mask.shape[1]
|
|
ratio = mask.sum() / (H * W)
|
|
if ratio > 0.8 * 0.8 or ratio < 0.1 * 0.1:
|
|
return False
|
|
else:
|
|
return True
|
|
|
|
|
|
def process_pairs(self, ref_image, ref_mask, tar_image, tar_mask, max_ratio = 0.8):
|
|
assert mask_score(ref_mask) > 0.90
|
|
assert self.check_mask_area(ref_mask) == True
|
|
assert self.check_mask_area(tar_mask) == True
|
|
|
|
# ========= Reference ===========
|
|
'''
|
|
# similate the case that the mask for reference object is coarse. Seems useless :(
|
|
|
|
if np.random.uniform(0, 1) < 0.7:
|
|
ref_mask_clean = ref_mask.copy()
|
|
ref_mask_clean = np.stack([ref_mask_clean,ref_mask_clean,ref_mask_clean],-1)
|
|
ref_mask = perturb_mask(ref_mask, 0.6, 0.9)
|
|
|
|
# select a fake bg to avoid the background leakage
|
|
fake_target = tar_image.copy()
|
|
h,w = ref_image.shape[0], ref_image.shape[1]
|
|
fake_targe = cv2.resize(fake_target, (w,h))
|
|
fake_back = np.fliplr(np.flipud(fake_target))
|
|
fake_back = self.aug_data_back(fake_back)
|
|
ref_image = ref_mask_clean * ref_image + (1-ref_mask_clean) * fake_back
|
|
'''
|
|
|
|
# Get the outline Box of the reference image
|
|
ref_box_yyxx = get_bbox_from_mask(ref_mask)
|
|
assert self.check_region_size(ref_mask, ref_box_yyxx, ratio = 0.10, mode = 'min') == True
|
|
|
|
# Filtering background for the reference image
|
|
ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
|
|
masked_ref_image = ref_image * ref_mask_3 + np.ones_like(ref_image) * 255 * (1-ref_mask_3)
|
|
|
|
y1,y2,x1,x2 = ref_box_yyxx
|
|
masked_ref_image = masked_ref_image[y1:y2,x1:x2,:]
|
|
ref_mask = ref_mask[y1:y2,x1:x2]
|
|
|
|
ratio = np.random.randint(11, 15) / 10
|
|
masked_ref_image, ref_mask = expand_image_mask(masked_ref_image, ref_mask, ratio=ratio)
|
|
ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
|
|
|
|
# Padding reference image to square and resize to 224
|
|
masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False)
|
|
masked_ref_image = cv2.resize(masked_ref_image.astype(np.uint8), (224,224) ).astype(np.uint8)
|
|
|
|
ref_mask_3 = pad_to_square(ref_mask_3 * 255, pad_value = 0, random = False)
|
|
ref_mask_3 = cv2.resize(ref_mask_3.astype(np.uint8), (224,224) ).astype(np.uint8)
|
|
ref_mask = ref_mask_3[:,:,0]
|
|
|
|
# Augmenting reference image
|
|
#masked_ref_image_aug = self.aug_data(masked_ref_image)
|
|
|
|
# Getting for high-freqency map
|
|
masked_ref_image_compose, ref_mask_compose = self.aug_data_mask(masked_ref_image, ref_mask)
|
|
masked_ref_image_aug = masked_ref_image_compose.copy()
|
|
|
|
ref_mask_3 = np.stack([ref_mask_compose,ref_mask_compose,ref_mask_compose],-1)
|
|
ref_image_collage = sobel(masked_ref_image_compose, ref_mask_compose/255)
|
|
|
|
|
|
# ========= Training Target ===========
|
|
tar_box_yyxx = get_bbox_from_mask(tar_mask)
|
|
tar_box_yyxx = expand_bbox(tar_mask, tar_box_yyxx, ratio=[1.1,1.2]) #1.1 1.3
|
|
assert self.check_region_size(tar_mask, tar_box_yyxx, ratio = max_ratio, mode = 'max') == True
|
|
|
|
# Cropping around the target object
|
|
tar_box_yyxx_crop = expand_bbox(tar_image, tar_box_yyxx, ratio=[1.3, 3.0])
|
|
tar_box_yyxx_crop = box2squre(tar_image, tar_box_yyxx_crop) # crop box
|
|
y1,y2,x1,x2 = tar_box_yyxx_crop
|
|
cropped_target_image = tar_image[y1:y2,x1:x2,:]
|
|
cropped_tar_mask = tar_mask[y1:y2,x1:x2]
|
|
tar_box_yyxx = box_in_box(tar_box_yyxx, tar_box_yyxx_crop)
|
|
y1,y2,x1,x2 = tar_box_yyxx
|
|
|
|
# Prepairing collage image
|
|
ref_image_collage = cv2.resize(ref_image_collage.astype(np.uint8), (x2-x1, y2-y1))
|
|
ref_mask_compose = cv2.resize(ref_mask_compose.astype(np.uint8), (x2-x1, y2-y1))
|
|
ref_mask_compose = (ref_mask_compose > 128).astype(np.uint8)
|
|
|
|
collage = cropped_target_image.copy()
|
|
collage[y1:y2,x1:x2,:] = ref_image_collage
|
|
|
|
collage_mask = cropped_target_image.copy() * 0.0
|
|
collage_mask[y1:y2,x1:x2,:] = 1.0
|
|
|
|
if np.random.uniform(0, 1) < 0.7:
|
|
cropped_tar_mask = perturb_mask(cropped_tar_mask)
|
|
collage_mask = np.stack([cropped_tar_mask,cropped_tar_mask,cropped_tar_mask],-1)
|
|
|
|
H1, W1 = collage.shape[0], collage.shape[1]
|
|
|
|
cropped_target_image = pad_to_square(cropped_target_image, pad_value = 0, random = False).astype(np.uint8)
|
|
collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
|
|
collage_mask = pad_to_square(collage_mask, pad_value = 2, random = False).astype(np.uint8)
|
|
H2, W2 = collage.shape[0], collage.shape[1]
|
|
|
|
cropped_target_image = cv2.resize(cropped_target_image.astype(np.uint8), (512,512)).astype(np.float32)
|
|
collage = cv2.resize(collage.astype(np.uint8), (512,512)).astype(np.float32)
|
|
collage_mask = cv2.resize(collage_mask.astype(np.uint8), (512,512), interpolation = cv2.INTER_NEAREST).astype(np.float32)
|
|
collage_mask[collage_mask == 2] = -1
|
|
|
|
# Prepairing dataloader items
|
|
masked_ref_image_aug = masked_ref_image_aug / 255
|
|
cropped_target_image = cropped_target_image / 127.5 - 1.0
|
|
collage = collage / 127.5 - 1.0
|
|
collage = np.concatenate([collage, collage_mask[:,:,:1] ] , -1)
|
|
|
|
item = dict(
|
|
ref=masked_ref_image_aug.copy(),
|
|
jpg=cropped_target_image.copy(),
|
|
hint=collage.copy(),
|
|
extra_sizes=np.array([H1, W1, H2, W2]),
|
|
tar_box_yyxx_crop=np.array(tar_box_yyxx_crop)
|
|
)
|
|
return item
|
|
|
|
|
|
|
|
|
|
|