Spaces:
Paused
Paused
| 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 * | |
| from .base import BaseDataset | |
| class SaliencyDataset(BaseDataset): | |
| def __init__(self, MSRA_root, TR_root, TE_root, HFlickr_root): | |
| image_mask_dict = {} | |
| # ====== MSRA-10k ====== | |
| file_lst = os.listdir(MSRA_root) | |
| image_lst = [MSRA_root+i for i in file_lst if '.jpg' in i] | |
| for i in image_lst: | |
| mask_path = i.replace('.jpg','.png') | |
| image_mask_dict[i] = mask_path | |
| # ===== DUT-TR ======== | |
| file_lst = os.listdir(TR_root) | |
| image_lst = [TR_root+i for i in file_lst if '.jpg' in i] | |
| for i in image_lst: | |
| mask_path = i.replace('.jpg','.png').replace('DUTS-TR-Image','DUTS-TR-Mask') | |
| image_mask_dict[i] = mask_path | |
| # ===== DUT-TE ======== | |
| file_lst = os.listdir(TE_root) | |
| image_lst = [TE_root+i for i in file_lst if '.jpg' in i] | |
| for i in image_lst: | |
| mask_path = i.replace('.jpg','.png').replace('DUTS-TE-Image','DUTS-TE-Mask') | |
| image_mask_dict[i] = mask_path | |
| # ===== HFlickr ======= | |
| file_lst = os.listdir(HFlickr_root) | |
| mask_list = [HFlickr_root+i for i in file_lst if '.png' in i] | |
| for i in file_lst: | |
| image_name = i.split('_')[0] +'.jpg' | |
| image_path = HFlickr_root.replace('masks', 'real_images') + image_name | |
| mask_path = HFlickr_root + i | |
| image_mask_dict[image_path] = mask_path | |
| self.image_mask_dict = image_mask_dict | |
| self.data = list(self.image_mask_dict.keys() ) | |
| self.size = (512,512) | |
| self.clip_size = (224,224) | |
| self.dynamic = 0 | |
| def __len__(self): | |
| return 20000 | |
| 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 get_sample(self, idx): | |
| # ==== get pairs ===== | |
| image_path = self.data[idx] | |
| mask_path = self.image_mask_dict[image_path] | |
| instances_mask = cv2.imread(mask_path) | |
| if len(instances_mask.shape) == 3: | |
| instances_mask = instances_mask[:,:,0] | |
| instances_mask = (instances_mask > 128).astype(np.uint8) | |
| # ====================== | |
| ref_image = cv2.imread(image_path) | |
| ref_image = cv2.cvtColor(ref_image.copy(), cv2.COLOR_BGR2RGB) | |
| tar_image = ref_image | |
| ref_mask = instances_mask | |
| tar_mask = instances_mask | |
| item_with_collage = self.process_pairs(ref_image, ref_mask, tar_image, tar_mask) | |
| sampled_time_steps = self.sample_timestep() | |
| item_with_collage['time_steps'] = sampled_time_steps | |
| return item_with_collage | |