# MIT License # Copyright (c) 2022 Intelligent Systems Lab Org # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # File author: Zhenyu Li # This file is partly inspired from ZoeDepth (https://github.com/isl-org/ZoeDepth/blob/main/zoedepth/data/data_mono.py); author: Shariq Farooq Bhat import os import numpy as np import torch from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision import transforms import os.path as osp import random import torch.nn as nn import cv2 import copy from zoedepth.utils.misc import get_boundaries from zoedepth.models.base_models.midas import Resize class SampleDataPairs(object): def __init__(self, num_sample_inout=50000, sampling_strategy='random', # or 'dda' dilation_factor=10, crop_size=(2160, 3840), ): self.num_sample_inout = num_sample_inout self.sampling_strategy = sampling_strategy self.dilation_factor = dilation_factor self.crop_height, self.crop_width = crop_size[0], crop_size[1] self.__init_grid() def __init_grid(self): nu = np.linspace(0, self.crop_width - 1, self.crop_width) nv = np.linspace(0, self.crop_height - 1, self.crop_height) u, v = np.meshgrid(nu, nv) self.u = u.flatten() self.v = v.flatten() def get_coords(self, gt): # get subpixel coordinates u = self.u + np.random.random_sample(self.u.size) v = self.v + np.random.random_sample(self.v.size) # use nearest neighbor to get gt for each samples d = gt[np.clip(np.rint(v).astype(np.uint16), 0, self.crop_height-1), np.clip(np.rint(u).astype(np.uint16), 0, self.crop_width-1)] # remove invalid depth values u = u[np.nonzero(d)] v = v[np.nonzero(d)] d = d[np.nonzero(d)] return np.stack((u, v, d), axis=-1) def get_boundaries(self, disp, th=1., dilation=10): edges_y = np.logical_or(np.pad(np.abs(disp[1:, :] - disp[:-1, :]) > th, ((1, 0), (0, 0))), np.pad(np.abs(disp[:-1, :] - disp[1:, :]) > th, ((0, 1), (0, 0)))) edges_x = np.logical_or(np.pad(np.abs(disp[:, 1:] - disp[:, :-1]) > th, ((0, 0), (1, 0))), np.pad(np.abs(disp[:, :-1] - disp[:,1:]) > th, ((0, 0), (0, 1)))) edges = np.logical_or(edges_y, edges_x).astype(np.float32) if dilation > 0: kernel = np.ones((dilation, dilation), np.uint8) edges = cv2.dilate(edges, kernel, iterations=1) return edges def __call__(self, results, disp_gt_copy): """Call function to load multiple types annotations. Args: results (dict): Result dict from :obj:`depth.CustomDataset`. Returns: dict: The dict contains loaded depth estimation annotations. """ gt = results['depth'] gt_squeeze = gt[:, :, 0] if self.sampling_strategy == "random": random_points = self.get_coords(gt_squeeze) idx = np.random.choice(random_points.shape[0], self.num_sample_inout) points = random_points[idx, :] elif self.sampling_strategy == "dda": disp_gt_squeeze = disp_gt_copy edges = self.get_boundaries(disp_gt_squeeze, dilation=self.dilation_factor) random_points = self.get_coords(gt_squeeze * (1. - edges)) edge_points = self.get_coords(gt_squeeze * edges) # if edge points exist if edge_points.shape[0]>0 and random_points.shape[0]>0: # Check tot num of edge points cond = edges.sum()//2 - self.num_sample_inout//2 < 0 tot= (self.num_sample_inout - int(edges.sum())//2, int(edges.sum())//2) if cond else \ (self.num_sample_inout//2, self.num_sample_inout//2) idx = np.random.choice(random_points.shape[0], tot[0]) idx_edges = np.random.choice(edge_points.shape[0], tot[1]) points = np.concatenate([edge_points[idx_edges, :], random_points[idx, :]], 0) # use uniform sample otherwise else: random_points = self.get_coords(gt_squeeze) idx = np.random.choice(random_points.shape[0], self.num_sample_inout) points = random_points[idx, :] results['sample_points'] = points return results def _is_pil_image(img): return isinstance(img, Image.Image) def _is_numpy_image(img): return isinstance(img, np.ndarray) and (img.ndim in {2, 3}) class ToTensor(object): def __init__(self, mode, do_normalize=False, size=None, sec_stage=False): self.mode = mode # don't do normalization as default self.normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if do_normalize else nn.Identity() self.size = size if size is not None: # self.resize = transforms.Resize(size=size) net_h, net_w = size self.resize = Resize(net_w, net_h, keep_aspect_ratio=False, ensure_multiple_of=32, resize_method="minimal") else: self.resize = nn.Identity() self.sec_stage = sec_stage def __call__(self, sample): image, focal = sample['image'], sample['focal'] crop_areas = sample.get('crop_area', None) if isinstance(image, list): # there must be crop_areas # only infer on eval sec_stage imgs_process = [] crp_process = [] for img, crp in zip(image, crop_areas): img = self.to_tensor(img) img = self.normalize(img) img = img.unsqueeze(dim=0) img = self.resize(img) img = img.squeeze(dim=0) imgs_process.append(img) crp = self.to_tensor(crp) crp = crp.unsqueeze(dim=0) crp = self.resize(crp) crp = crp.squeeze(dim=0) crp_process.append(crp) image = torch.cat(imgs_process, dim=0) crop_areas = torch.cat(crp_process, dim=0) img_temp = sample['img_temp'] img_temp = self.to_tensor(img_temp) img_temp = self.normalize(img_temp) img_temp = img_temp.unsqueeze(dim=0) img_temp = self.resize(img_temp) #NOTE: hack img_temp = img_temp.squeeze(dim=0) image_raw = copy.deepcopy(img_temp) else: image = self.to_tensor(image) image = self.normalize(image) if crop_areas is not None: crop_areas = self.to_tensor(crop_areas) crop_areas = crop_areas.unsqueeze(dim=0) crop_areas = self.resize(crop_areas) crop_areas = crop_areas.squeeze(dim=0) if self.sec_stage: img_temp = sample['img_temp'] img_temp = self.to_tensor(img_temp) img_temp = self.normalize(img_temp) img_temp = img_temp.unsqueeze(dim=0) img_temp = self.resize(img_temp) image_raw = img_temp.squeeze(dim=0) image = image.unsqueeze(dim=0) image = self.resize(image) image = image.squeeze(dim=0) else: # in the first stage, this hr info is reserved image_raw = copy.deepcopy(image) image = image.unsqueeze(dim=0) image = self.resize(image) image = image.squeeze(dim=0) if self.mode == 'test': return_dict = {'image': image, 'focal': focal} if crop_areas is not None: return_dict['crop_area'] = crop_areas return return_dict depth = sample['depth'] depth = self.to_tensor(depth) depth_gt_temp = sample['depth_gt_temp'] depth_gt_raw = self.to_tensor(depth_gt_temp) if self.mode == 'train': return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw, 'depth_raw': depth_gt_raw} if crop_areas is not None: return_dict['crop_area'] = crop_areas return return_dict else: has_valid_depth = sample['has_valid_depth'] # image = self.resize(image) return_dict = {**sample, 'image': image, 'depth': depth, 'focal': focal, 'image_raw': image_raw, 'has_valid_depth': has_valid_depth, 'image_path': sample['image_path'], 'depth_path': sample['depth_path'], 'depth_raw': depth_gt_raw} if crop_areas is not None: return_dict['crop_area'] = crop_areas return return_dict def to_tensor(self, pic): if isinstance(pic, np.ndarray): img = torch.from_numpy(pic.transpose((2, 0, 1))) # img here return img def preprocessing_transforms(mode, sec_stage=False, **kwargs): return transforms.Compose([ ToTensor(mode=mode, sec_stage=sec_stage, **kwargs) ]) def remove_leading_slash(s): if s[0] == '/' or s[0] == '\\': return s[1:] return s class U4KDataset(Dataset): def __init__(self, config, mode, data_root, split): self.mode = mode self.config = config self.data_root = data_root self.split = split img_size = self.config.get("img_size", None) img_size = img_size if self.config.get( "do_input_resize", False) else None self.sec_stage = self.config.get("sec_stage", False) self.transform = preprocessing_transforms(mode, size=img_size, sec_stage=self.sec_stage) self.data_infos = self.load_data_list() self.sampled_training = self.config.get("sampled_training", False) if self.sampled_training: self.data_sampler = SampleDataPairs( num_sample_inout=config.num_sample_inout, sampling_strategy=config.sampling_strategy, # or 'dda' dilation_factor=config.dilation_factor, crop_size=config.transform_sample_gt_size) self.random_crop = self.config.get("random_crop", False) self.crop_size = [540, 960] # 1/4 self.overlap = self.config.get("overlap", False) self.consistency_training = self.config.get("consistency_training", False) self.overlap_length_h = self.config.get("overlap_length_h", int(256)) self.overlap_length_w = self.config.get("overlap_length_w", int(384)) print("current overlap_length_h and overlap_length_w are {} and {}".format(self.overlap_length_h, self.overlap_length_w)) def load_data_list(self): """Load annotation from directory. Args: data_root (str): Data root for img_dir/ann_dir. split (str|None): Split txt file. If split is specified, only file with suffix in the splits will be loaded. Otherwise, all images in img_dir/ann_dir will be loaded. Default: None Returns: list[dict]: All image info of dataset. """ data_root = self.data_root split = self.split self.invalid_depth_num = 0 img_infos = [] if split is not None: with open(split) as f: for line in f: img_info_l = dict() # img_info_r = dict() img_l, img_r, depth_map_l, depth_map_r = line.strip().split(" ") # HACK: a hack to replace the png with raw to accelerate training img_l = img_l[:-3] + 'raw' # img_r = img_r[:-3] + 'raw' img_info_l['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_l)) # img_info_r['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_r)) img_info_l['img_path'] = osp.join(data_root, remove_leading_slash(img_l)) # img_info_r['filename'] = osp.join(data_root, remove_leading_slash(img_r)) img_info_l['depth_fields'] = [] filename = img_info_l['depth_map_path'] ext_name_l = filename.replace('Disp0', 'Extrinsics0') ext_name_l = ext_name_l.replace('npy', 'txt') ext_name_r = filename.replace('Disp0', 'Extrinsics1') ext_name_r = ext_name_r.replace('npy', 'txt') with open(ext_name_l, 'r') as f: ext_l = f.readlines() with open(ext_name_r, 'r') as f: ext_r = f.readlines() f = float(ext_l[0].split(' ')[0]) img_info_l['focal'] = f base = abs(float(ext_l[1].split(' ')[3]) - float(ext_r[1].split(' ')[3])) img_info_l['depth_factor'] = base * f img_infos.append(img_info_l) # img_infos.append(img_info_r) else: raise NotImplementedError # github issue:: make sure the same order img_infos = sorted(img_infos, key=lambda x: x['img_path']) return img_infos def augment_image(self, image): # gamma augmentation gamma = random.uniform(0.9, 1.1) image_aug = image ** gamma # brightness augmentation if self.config.dataset == 'nyu': brightness = random.uniform(0.75, 1.25) else: brightness = random.uniform(0.9, 1.1) image_aug = image_aug * brightness # color augmentation colors = np.random.uniform(0.9, 1.1, size=3) white = np.ones((image.shape[0], image.shape[1])) color_image = np.stack([white * colors[i] for i in range(3)], axis=2) image_aug *= color_image image_aug = np.clip(image_aug, 0, 1) return image_aug def train_preprocess(self, image, depth_gt): if self.config.aug: # Random flipping do_flip = random.random() if do_flip > 0.5: image = (image[:, ::-1, :]).copy() depth_gt = (depth_gt[:, ::-1, :]).copy() # Random gamma, brightness, color augmentation do_augment = random.random() if do_augment > 0.5: image = self.augment_image(image) return image, depth_gt def get_crop_bbox(self, img): """Randomly get a crop bounding box.""" margin_h = max(img.shape[0] - self.crop_size[0], 0) margin_w = max(img.shape[1] - self.crop_size[1], 0) offset_h = np.random.randint(0, margin_h + 1) offset_w = np.random.randint(0, margin_w + 1) crop_y1, crop_y2 = offset_h, offset_h + self.crop_size[0] crop_x1, crop_x2 = offset_w, offset_w + self.crop_size[1] return crop_y1, crop_y2, crop_x1, crop_x2 def crop(self, img, crop_bbox, tmp=False): """Crop from ``img``""" crop_y1, crop_y2, crop_x1, crop_x2 = crop_bbox if tmp: templete = np.zeros((img.shape[0], img.shape[1], 1), dtype=np.float32) templete[crop_y1:crop_y2, crop_x1:crop_x2, :] = 1.0 img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...] return img, templete else: img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...] return img def __getitem__(self, idx): img_file_path = self.data_infos[idx]['img_path'] disp_path = self.data_infos[idx]['depth_map_path'] depth_factor = self.data_infos[idx]['depth_factor'] height=2160 width=3840 image = np.fromfile(open(img_file_path, 'rb'), dtype=np.uint8).reshape(height, width, 3) / 255.0 if self.config.get("use_rgb", False): image = image.astype(np.float32)[:, :, ::-1].copy() elif self.config.get("use_brg", False): image = image.astype(np.float32)[:, :, [0, 2, 1]].copy() elif self.config.get("use_gbr", False): image = image.astype(np.float32)[:, :, [1, 0, 2]].copy() elif self.config.get("use_rbg", False): image = image.astype(np.float32)[:, :, [2, 0, 1]].copy() elif self.config.get("use_grb", False): image = image.astype(np.float32)[:, :, [1, 2, 0]].copy() else: image = image.astype(np.float32) disp_gt = np.expand_dims(np.load(disp_path, mmap_mode='c'), -1) disp_gt = disp_gt.astype(np.float32) disp_gt_copy = disp_gt[:, :, 0].copy() depth_gt = depth_factor / disp_gt depth_gt[depth_gt > self.config.max_depth] = self.config.max_depth # for vis focal = self.data_infos[idx]['focal'] bbox = None bboxs_res = None crop_areas = None bboxs_roi = None # hack for infer if self.mode == 'train': image, depth_gt = self.train_preprocess(image, depth_gt) img_temp = copy.deepcopy(image) depth_gt_temp = copy.deepcopy(depth_gt) if self.random_crop: # use in sec_stage if self.consistency_training: crop_y1, crop_y2, crop_x1, crop_x2 = self.get_crop_bbox(image) # ensure the prob of crop is the same while True: # shift_x = random.randint(self.overlap_length//3, self.overlap_length) # shift_y = random.randint(self.overlap_length//3, self.overlap_length) shift_x = self.overlap_length_w shift_y = self.overlap_length_h if random.random() > 0.5: shift_x = shift_x * -1 if random.random() > 0.5: shift_y = shift_y * -1 crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift = crop_y1 + shift_y, crop_y2 + shift_y, crop_x1 + shift_x, crop_x2 + shift_x if crop_y1_shift > 0 and crop_x1_shift > 0 and crop_y2_shift < image.shape[0] and crop_x2_shift < image.shape[1]: break bbox_ori = (crop_y1, crop_y2, crop_x1, crop_x2) bbox_shift = (crop_y1_shift, crop_y2_shift, crop_x1_shift, crop_x2_shift) image_ori, crop_area_ori = self.crop(image, bbox_ori, tmp=True) image_shift, crop_area_shift = self.crop(image, bbox_shift, tmp=True) depth_gt_ori = self.crop(depth_gt, bbox_ori) depth_gt_shift = self.crop(depth_gt, bbox_shift) disp_gt_copy_ori = self.crop(disp_gt_copy, bbox_ori) disp_gt_copy_shift = self.crop(disp_gt_copy, bbox_shift) bboxs_ori = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) bboxs_shift = torch.tensor([crop_x1_shift / width * 512, crop_y1_shift / height * 384, crop_x2_shift / width * 512, crop_y2_shift / height * 384]) bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2]) bboxs_raw_shift = torch.tensor([crop_x1_shift, crop_y1_shift, crop_x2_shift, crop_y2_shift]) else: bbox = self.get_crop_bbox(image) image, crop_area = self.crop(image, bbox, tmp=True) depth_gt = self.crop(depth_gt, bbox) disp_gt_copy = self.crop(disp_gt_copy, bbox) crop_y1, crop_y2, crop_x1, crop_x2 = bbox bboxs_res = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) # coord in 384, 512 bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2]) mask = np.logical_and(depth_gt > self.config.min_depth, depth_gt < self.config.max_depth).squeeze()[None, ...] mask_raw = np.logical_and(depth_gt_temp > self.config.min_depth, depth_gt_temp < self.config.max_depth).squeeze()[None, ...] sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'mask': mask, 'image_raw': image.copy(), 'mask_raw': mask_raw, 'image_path': img_file_path} if self.random_crop: if self.consistency_training: image = np.concatenate([image_ori, image_shift], axis=-1) depth_gt = np.concatenate([depth_gt_ori, depth_gt_shift], axis=-1) crop_area = np.concatenate([crop_area_ori, crop_area_shift], axis=-1) bboxs_res = torch.cat([bboxs_ori, bboxs_shift], dim=-1) bboxes_raw_res = torch.cat([bboxs_raw, bboxs_raw_shift], dim=-1) mask = np.logical_and(depth_gt > self.config.min_depth, depth_gt < self.config.max_depth) # hack the sample dict sample['image'] = image sample['depth'] = depth_gt sample['crop_area'] = crop_area sample['bbox'] = bboxs_res sample['bbox_raw'] = bboxes_raw_res sample['shift'] = torch.tensor([shift_y, shift_x]) # h direction, then w direction sample['mask'] = mask else: if bboxs_res is not None: sample['bbox'] = bboxs_res sample['bbox_raw'] = bboxs_raw sample['crop_area'] = crop_area if self.sampled_training: self.data_sampler(sample, disp_gt_copy) # update mask sample_points = sample['sample_points'] sample_mask = np.logical_and(sample_points[:, -1] > self.config.min_depth, sample_points[:, -1] < self.config.max_depth).squeeze()[None, ...] sample['sample_mask'] = sample_mask else: # nothing needs to be changed for consistency training. img_temp = copy.deepcopy(image) depth_gt_temp = copy.deepcopy(depth_gt) if self.sec_stage: # x_start, y_start = [0, 540, 1080, 1620], [0, 960, 1920, 2880] x_start, y_start = [0 + 3 * self.overlap / 2, 540 + self.overlap / 2, 1080 - self.overlap / 2, 1620 - 3 * self.overlap / 2], \ [0 + 3 * self.overlap / 2, 960 + self.overlap / 2, 1920 - self.overlap / 2, 2880 - 3 * self.overlap / 2] img_crops = [] bboxs_roi = [] crop_areas = [] bboxs_raw_list = [] for x in x_start: for y in y_start: bbox = (int(x), int(x+540), int(y), int(y+960)) img_crop, crop_area = self.crop(image, bbox, tmp=True) img_crops.append(img_crop) crop_areas.append(crop_area) crop_y1, crop_y2, crop_x1, crop_x2 = bbox bbox_roi = torch.tensor([crop_x1 / width * 512, crop_y1 / height * 384, crop_x2 / width * 512, crop_y2 / height * 384]) bboxs_roi.append(bbox_roi) bboxs_raw = torch.tensor([crop_x1, crop_y1, crop_x2, crop_y2]) bboxs_raw_list.append(bboxs_raw) image = img_crops bboxs_roi = torch.stack(bboxs_roi, dim=0) bboxs_raw = torch.stack(bboxs_raw_list, dim=0) # bbox = (820, 1360 ,1440, 2400) # a hack version for quick evaluation # image = self.crop(image, bbox) # depth_gt = self.crop(depth_gt, bbox) # disp_gt_copy = self.crop(disp_gt_copy, bbox) mask = np.logical_and(depth_gt > self.config.min_depth, depth_gt < self.config.max_depth).squeeze()[None, ...] disp_gt_edges = get_boundaries(disp_gt_copy, th=1, dilation=0) if self.mode == 'online_eval': sample = {'image': image, 'depth': depth_gt, 'focal': focal, 'has_valid_depth': True, 'image_path': img_file_path, 'depth_path': disp_path, 'depth_factor_path': depth_factor, 'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path} if bboxs_roi is not None: sample['bbox'] = bboxs_roi sample['bbox_raw'] = bboxs_raw if crop_areas is not None: sample['crop_area'] = crop_areas else: sample = {'image': image, 'focal': focal, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges, 'image_path': img_file_path} if bboxs_roi is not None: sample['bbox'] = bboxs_roi sample['bbox_raw'] = bboxs_raw if crop_areas is not None: sample['crop_area'] = crop_areas if self.transform: sample['img_temp'] = img_temp sample['depth_gt_temp'] = depth_gt_temp sample = self.transform(sample) sample['dataset'] = self.config.dataset return sample def __len__(self): return len(self.data_infos) def get_u4k_loader(config, mode, transform): if mode == 'train': dataset = U4KDataset(config, mode, config.data_path, config.filenames_train) dataset[0] if config.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) else: train_sampler = None dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=(train_sampler is None), num_workers=config.workers, pin_memory=True, persistent_workers=True, sampler=train_sampler) elif mode == 'train_save': dataset = U4KDataset(config, 'online_eval', config.data_path, config.filenames_train) if config.distributed: train_sampler = None else: train_sampler = None dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1, pin_memory=False, sampler=train_sampler) elif mode == 'online_eval': dataset = U4KDataset(config, mode, config.data_path, config.filenames_val) # dataset = U4KDataset(config, mode, config.data_path, config.filenames_train) if config.distributed: # redundant. here only for readability and to be more explicit # Give whole test set to all processes (and report evaluation only on one) regardless eval_sampler = None else: eval_sampler = None dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1, pin_memory=False, sampler=eval_sampler) else: dataset = U4KDataset(config, mode, config.data_path, config.filenames_test) dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1) return dataloader