import os import random import cv2 import numpy as np import torch from torch.utils.data import Dataset import json DATASET_NAMES = [ 'BIPED', 'BIPED-B2', 'BIPED-B3', 'BIPED-B5', 'BIPED-B6', 'BSDS', # 5 'BRIND', # 6 'ICEDA', #7 'BSDS300', 'CID', #9 'DCD', 'MDBD', #11 'PASCAL', 'NYUD', #13 'BIPBRI', 'UDED', # 15 just for testing 'DMRIR', 'CLASSIC' ] # 8 # [108, 109.451,112.230,137.86] BIPED_mean = [103.939,116.779,123.68,137.86] def dataset_info(dataset_name, is_linux=True): if is_linux: config = { 'UDED': { 'img_height': 512, # 321 'img_width': 512, # 481 'train_list': None, 'test_list': 'test_pair.lst', 'data_dir': '/root/workspace/datasets/UDED', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86]# [104.007, 116.669, 122.679, 137.86] }, #[98.939,111.779,117.68,137.86] 'BSDS': { 'img_height': 512, #321 'img_width': 512, #481 'train_list': 'train_pair.lst', 'test_list': 'test_pair.lst', 'data_dir': '/root/workspace/datasets/BSDS', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'BRIND': { 'img_height': 512, # 321 'img_width': 512, # 481 'train_list': 'train_pair_all.lst', # all train_pair_all.lst # less train_pair.lst 'test_list': 'test_pair.lst', 'data_dir': '/root/workspace/datasets/BRIND', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'ICEDA': { 'img_height': 1024, # 321 'img_width': 1408, # 481 'train_list': None, 'test_list': 'test_pair.lst', 'data_dir': '/root/workspace/datasets/ICEDA', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'BSDS300': { 'img_height': 512, #321 'img_width': 512, #481 'test_list': 'test_pair.lst', 'train_list': None, 'data_dir': '/root/workspace/datasets/BSDS300', # NIR 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'PASCAL': { 'img_height': 416, # 375 'img_width': 512, #500 'test_list': 'test_pair.lst', 'train_list': None, 'data_dir': '/root/datasets/PASCAL', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'CID': { 'img_height': 512, 'img_width': 512, 'test_list': 'test_pair.lst', 'train_list': None, 'data_dir': '/root/datasets/CID', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'NYUD': { 'img_height': 448,#425 'img_width': 560,#560 'test_list': 'test_pair.lst', 'train_list': None, 'data_dir': '/root/datasets/NYUD', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'MDBD': { 'img_height': 720, 'img_width': 1280, 'test_list': 'test_pair.lst', 'train_list': 'train_pair.lst', 'data_dir': '/root/workspace/datasets/MDBD', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'BIPED': { 'img_height': 720, #720 # 1088 'img_width': 1280, # 1280 5 1920 'test_list': 'test_pair.lst', 'train_list': 'train_pair0.lst', # Base augmentation # 'train_list': 'train_pairB3.lst', # another augmentation # 'train_list': 'train_pairB5.lst', # Last augmentation 'data_dir': '/root/workspace/datasets/BIPED', # mean_rgb 'yita': 0.5, 'mean':BIPED_mean # }, 'CLASSIC': { 'img_height': 512,# 'img_width': 512,# 512 'test_list': None, 'train_list': None, 'data_dir': 'data', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'BIPED-B2': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B3': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B5': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B6': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'DCD': { 'img_height': 352, #240 'img_width': 480,# 360 'test_list': 'test_pair.lst', 'train_list': None, 'data_dir': '/opt/dataset/DCD', # mean_rgb 'yita': 0.2, 'mean': [104.007, 116.669, 122.679, 137.86] } } else: config = { 'UDED': { 'img_height': 512, # 321 'img_width': 512, # 481 'train_list': None, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/UDED', # mean_rgb 'yita': 0.5, 'mean':[104.007, 116.669, 122.679, 137.86] # [183.939,196.779,203.68,137.86] # [104.007, 116.669, 122.679, 137.86] }, 'BSDS': {'img_height': 480, # 321 'img_width': 480, # 481 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/BSDS', # mean_rgb 'yita': 0.5, 'mean':[103.939, 116.669, 122.679, 137.86] }, # [103.939, 116.669, 122.679, 137.86] #[159.510, 159.451,162.230,137.86] 'BRIND': { 'img_height': 512, # 321 'img_width': 512, # 481 'train_list': 'train_pair_all.lst', # all train_pair_all.lst # less train_pair.lst 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/BRIND', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'ICEDA': { 'img_height': 1024, # 321 'img_width': 1408, # 481 'train_list': None, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/ICEDA', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86] }, 'BSDS300': {'img_height': 512, # 321 'img_width': 512, # 481 'test_list': 'test_pair.lst', 'data_dir': 'C:/Users/xavysp/dataset/BSDS300', # NIR 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'PASCAL': {'img_height': 375, 'img_width': 500, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/PASCAL', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'CID': {'img_height': 512, 'img_width': 512, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/CID', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'NYUD': {'img_height': 425, 'img_width': 560, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/NYUD', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'MDBD': {'img_height': 720, 'img_width': 1280, 'test_list': 'test_pair.lst', 'train_list': 'train_pair.lst', 'data_dir': 'C:/dataset/MDBD', # mean_rgb 'yita': 0.3, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'BIPED': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_pair0.lst', # 'train_list': 'train_rgb.lst', 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B2': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B3': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B5': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'BIPED-B6': {'img_height': 720, # 720 'img_width': 1280, # 1280 'test_list': 'test_pair.lst', 'train_list': 'train_rgb.lst', 'data_dir': 'C:/Users/xavysp/dataset/BIPED', # WIN: '../.../dataset/BIPED/edges' 'yita': 0.5, 'mean':BIPED_mean}, 'CLASSIC': {'img_height': 512, 'img_width': 512, 'test_list': None, 'train_list': None, 'data_dir': 'teed_tmp', # mean_rgb 'yita': 0.5, 'mean': [104.007, 116.669, 122.679, 137.86]}, 'DCD': {'img_height': 240, 'img_width': 360, 'test_list': 'test_pair.lst', 'data_dir': 'C:/dataset/DCD', # mean_rgb 'yita': 0.2, 'mean': [104.007, 116.669, 122.679, 137.86]} } return config[dataset_name] class TestDataset(Dataset): def __init__(self, data_root, test_data, img_height, img_width, test_list=None, arg=None ): if test_data not in DATASET_NAMES: raise ValueError(f"Unsupported dataset: {test_data}") self.data_root = data_root self.test_data = test_data self.test_list = test_list self.args = arg self.up_scale = arg.up_scale self.mean_bgr = arg.mean_test if len(arg.mean_test) == 3 else arg.mean_test[:3] self.img_height = img_height self.img_width = img_width self.data_index = self._build_index() def _build_index(self): sample_indices = [] if self.test_data == "CLASSIC": # for single image testing images_path = os.listdir(self.data_root) labels_path = None sample_indices = [images_path, labels_path] else: # image and label paths are located in a list file if not self.test_list: raise ValueError( f"Test list not provided for dataset: {self.test_data}") list_name = os.path.join(self.data_root, self.test_list) if self.test_data.upper() in ['BIPED', 'BRIND','UDED','ICEDA']: with open(list_name,encoding='utf-8') as f: files = json.load(f) for pair in files: tmp_img = pair[0] tmp_gt = pair[1] sample_indices.append( (os.path.join(self.data_root, tmp_img), os.path.join(self.data_root, tmp_gt),)) else: with open(list_name, 'r') as f: files = f.readlines() files = [line.strip() for line in files] pairs = [line.split() for line in files] for pair in pairs: tmp_img = pair[0] tmp_gt = pair[1] sample_indices.append( (os.path.join(self.data_root, tmp_img), os.path.join(self.data_root, tmp_gt),)) return sample_indices def __len__(self): return len(self.data_index[0]) if self.test_data.upper() == 'CLASSIC' else len(self.data_index) def __getitem__(self, idx): # get data sample # image_path, label_path = self.data_index[idx] if self.data_index[1] is None: image_path = self.data_index[0][idx] if len(self.data_index[0]) > 1 else self.data_index[0][idx - 1] else: image_path = self.data_index[idx][0] label_path = None if self.test_data == "CLASSIC" else self.data_index[idx][1] img_name = os.path.basename(image_path) # print(img_name) file_name = os.path.splitext(img_name)[0] + ".png" # base dir if self.test_data.upper() == 'BIPED': img_dir = os.path.join(self.data_root, 'imgs', 'test') gt_dir = os.path.join(self.data_root, 'edge_maps', 'test') elif self.test_data.upper() == 'CLASSIC': img_dir = self.data_root gt_dir = None else: img_dir = self.data_root gt_dir = self.data_root # load data image = cv2.imdecode(np.fromfile(os.path.join(img_dir, image_path), np.uint8), cv2.IMREAD_COLOR) if not self.test_data == "CLASSIC": label = cv2.imread(os.path.join( gt_dir, label_path), cv2.IMREAD_COLOR) else: label = None im_shape = [image.shape[0], image.shape[1]] image, label = self.transform(img=image, gt=label) return dict(images=image, labels=label, file_names=file_name, image_shape=im_shape) def transform(self, img, gt): # gt[gt< 51] = 0 # test without gt discrimination # up scale test image if self.up_scale: # For TEED BIPBRIlight Upscale img = cv2.resize(img,(0,0),fx=1.3,fy=1.3) if img.shape[0] < 512 or img.shape[1] < 512: #TEED BIPED standard proposal if you want speed up the test, comment this block img = cv2.resize(img, (0, 0), fx=1.5, fy=1.5) # else: # img = cv2.resize(img, (0, 0), fx=1.1, fy=1.1) # Make sure images and labels are divisible by 2^4=16 if img.shape[0] % 8 != 0 or img.shape[1] % 8 != 0: img_width = ((img.shape[1] // 8) + 1) * 8 img_height = ((img.shape[0] // 8) + 1) * 8 img = cv2.resize(img, (img_width, img_height)) # gt = cv2.resize(gt, (img_width, img_height)) else: pass # img_width = self.args.test_img_width # img_height = self.args.test_img_height # img = cv2.resize(img, (img_width, img_height)) # gt = cv2.resize(gt, (img_width, img_height)) # # For FPS # img = cv2.resize(img, (496,320)) img = np.array(img, dtype=np.float32) # if self.rgb: # img = img[:, :, ::-1] # RGB->BGR img -= self.mean_bgr img = img.transpose((2, 0, 1)) img = torch.from_numpy(img.copy()).float() if self.test_data == "CLASSIC": gt = np.zeros((img.shape[:2])) gt = torch.from_numpy(np.array([gt])).float() else: gt = np.array(gt, dtype=np.float32) if len(gt.shape) == 3: gt = gt[:, :, 0] gt /= 255. gt = torch.from_numpy(np.array([gt])).float() return img, gt # ************************************************* # ************* training ************************** # ************************************************* class BipedDataset(Dataset): train_modes = ['train', 'test', ] dataset_types = ['rgbr', ] data_types = ['aug', ] def __init__(self, data_root, img_height, img_width, train_mode='train', dataset_type='rgbr', # is_scaling=None, # Whether to crop image or otherwise resize image to match image height and width. crop_img=False, arg=None ): self.data_root = data_root self.train_mode = train_mode self.dataset_type = dataset_type self.data_type = 'aug' # be aware that this might change in the future self.img_height = img_height self.img_width = img_width self.mean_bgr = arg.mean_train if len(arg.mean_train) == 3 else arg.mean_train[:3] self.crop_img = crop_img self.arg = arg self.data_index = self._build_index() def _build_index(self): assert self.train_mode in self.train_modes, self.train_mode assert self.dataset_type in self.dataset_types, self.dataset_type assert self.data_type in self.data_types, self.data_type data_root = os.path.abspath(self.data_root) sample_indices = [] file_path = os.path.join(data_root, self.arg.train_list) if self.arg.train_data.lower() == 'bsds': with open(file_path, 'r') as f: files = f.readlines() files = [line.strip() for line in files] pairs = [line.split() for line in files] for pair in pairs: tmp_img = pair[0] tmp_gt = pair[1] sample_indices.append( (os.path.join(data_root, tmp_img), os.path.join(data_root, tmp_gt),)) else: with open(file_path) as f: files = json.load(f) for pair in files: tmp_img = pair[0] tmp_gt = pair[1] sample_indices.append( (os.path.join(data_root, tmp_img), os.path.join(data_root, tmp_gt),)) return sample_indices def __len__(self): return len(self.data_index) def __getitem__(self, idx): # get data sample image_path, label_path = self.data_index[idx] # load data image = cv2.imdecode(np.fromfile(image_path, np.uint8), cv2.IMREAD_COLOR) label = cv2.imdecode(np.fromfile(label_path), cv2.IMREAD_GRAYSCALE) image, label = self.transform(img=image, gt=label) return dict(images=image, labels=label) def transform(self, img, gt): gt = np.array(gt, dtype=np.float32) if len(gt.shape) == 3: gt = gt[:, :, 0] gt /= 255. # for LDC input and BDCN img = np.array(img, dtype=np.float32) img -= self.mean_bgr i_h, i_w, _ = img.shape # 400 for BIPEd and 352 for BSDS check with 384 crop_size = self.img_height if self.img_height == self.img_width else None # 448# MDBD=480 BIPED=480/400 BSDS=352 # # # for BSDS 352/BRIND # if i_w > crop_size and i_h > crop_size: # later 400, before crop_size # i = random.randint(0, i_h - crop_size) # j = random.randint(0, i_w - crop_size) # img = img[i:i + crop_size, j:j + crop_size] # gt = gt[i:i + crop_size, j:j + crop_size] # for BIPED/MDBD # Second augmentation if i_w> 400 and i_h>400: #before 420 h,w = gt.shape if np.random.random() > 0.4: #before i_w> 500 and i_h>500: LR_img_size = crop_size #l BIPED=256, 240 200 # MDBD= 352 BSDS= 176 i = random.randint(0, h - LR_img_size) j = random.randint(0, w - LR_img_size) # if img. img = img[i:i + LR_img_size , j:j + LR_img_size ] gt = gt[i:i + LR_img_size , j:j + LR_img_size ] else: LR_img_size = 300# 256 300 400 # l BIPED=208-352, # MDBD= 352-480- BSDS= 176-320 i = random.randint(0, h - LR_img_size) j = random.randint(0, w - LR_img_size) # if img. img = img[i:i + LR_img_size, j:j + LR_img_size] gt = gt[i:i + LR_img_size, j:j + LR_img_size] img = cv2.resize(img, dsize=(crop_size, crop_size), ) gt = cv2.resize(gt, dsize=(crop_size, crop_size)) else: # New addidings img = cv2.resize(img, dsize=(crop_size, crop_size)) gt = cv2.resize(gt, dsize=(crop_size, crop_size)) # BRIND Best for TEDD+BIPED gt[gt > 0.1] +=0.2#0.4 gt = np.clip(gt, 0., 1.) img = img.transpose((2, 0, 1)) img = torch.from_numpy(img.copy()).float() gt = torch.from_numpy(np.array([gt])).float() return img, gt