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