training/mono/datasets/__base_dataset__.py

import os
import json
import torch
import torchvision.transforms as transforms
import os.path
import numpy as np
import cv2
from torch.utils.data import Dataset
import random
import mono.utils.transform as img_transform
import copy 
from mono.utils.comm import get_func
import pickle
import logging
import multiprocessing as mp
import ctypes
"""
Dataset annotations are saved in a Json file. All data, including rgb, depth, pose, and so on, captured within the same frame are saved in the same dict.
All frames are organized in a list. In each frame, it may contains the some or all of following data format. 

# Annotations for the current central RGB/depth cameras.

'rgb':          rgb image in the current frame.
'depth':        depth map in the current frame.
'sem':          semantic mask in the current frame.
'cam_in':       camera intrinsic parameters of the current rgb camera. 
'cam_ex':       camera extrinsic parameters of the current rgb camera.
'cam_ex_path':  path to the extrinsic parameters.
'pose':         pose in current frame.
'timestamp_rgb':    time stamp of current rgb image.

# Annotations for the left hand RGB/depth cameras.

'rgb_l':          rgb image of the left hand camera in the current frame.
'depth_l':        depth map of the left hand camera in the current frame.
'sem_l':          semantic mask of the left hand camera in the current frame.
'cam_in_l':       camera intrinsic parameters of the left hand rgb camera in the current frame.
'cam_ex_l':       camera extrinsic parameters of the left hand rgb camera in the current frame.
'cam_ex_path':    path to the extrinsic parameters.
'pose_l':         pose of the left hand camera  in the incurrent frame.
'timestamp_rgb_l':    time stamp of the rgb img captured by the left hand camera.

# Annotations for the right RGB/depth cameras, which is on the left hand of the current central cameras.

'rgb_r':          rgb image of the right hand camera in the current frame.
'depth_r':        depth map of the right hand camera in the current frame.
'sem_r':          semantic mask of the right hand camera in the current frame.
'cam_in_r':       camera intrinsic parameters of the right hand rgb camera in the current frame.
'cam_ex_r':       camera extrinsic parameters of the right hand rgb camera in the current frame.
'cam_ex_path_r':  path to the extrinsic parameters.
'pose_r':         pose of the right hand camera  in the incurrent frame.
'timestamp_rgb_r':    time stamp of the rgb img captured by the right hand camera.

# Annotations for the central RGB/depth cameras in the last frame.

'rgb_pre':          rgb image of the central camera in the last frame.
'depth_pre':        depth map of the central camera in the last frame.
'sem_pre':          semantic mask of the central camera in the last frame.
'cam_in_pre':       camera intrinsic parameters of the central rgb camera in the last frame.
'cam_ex_pre':       camera extrinsic parameters of the central rgb camera in the last frame.
'cam_ex_path_pre':  path to the extrinsic parameters.
'pose_pre':         pose of the central camera  in the last frame.
'timestamp_rgb_pre':    time stamp of the rgb img captured by the central camera.

# Annotations for the central RGB/depth cameras in the next frame.

'rgb_next':          rgb image of the central camera in the next frame.
'depth_next':        depth map of the central camera in the next frame.
'sem_next':          semantic mask of the central camera in the next frame.
'cam_in_next':       camera intrinsic parameters of the central rgb camera in the next frame.
'cam_ex_next':       camera extrinsic parameters of the central rgb camera in the next frame.
'cam_ex_path_next':  path to the extrinsic parameters.
'pose_next':         pose of the central camera  in the next frame.
'timestamp_rgb_next':    time stamp of the rgb img captured by the central camera.
"""

class BaseDataset(Dataset):
    def __init__(self, cfg, phase, **kwargs):
        super(BaseDataset, self).__init__()
        self.cfg = cfg
        self.phase = phase
        self.db_info = kwargs['db_info']
        
        # root dir for data
        self.data_root = os.path.join(self.db_info['db_root'], self.db_info['data_root']) 
        # depth/disp data root
        disp_root = self.db_info['disp_root'] if 'disp_root' in self.db_info else None
        self.disp_root = os.path.join(self.db_info['db_root'], disp_root) if disp_root is not None else None
        depth_root = self.db_info['depth_root'] if 'depth_root' in self.db_info else None
        self.depth_root = os.path.join(self.db_info['db_root'], depth_root) if depth_root is not None \
            else self.data_root
        # meta data root
        meta_data_root = self.db_info['meta_data_root'] if 'meta_data_root' in self.db_info else None
        self.meta_data_root = os.path.join(self.db_info['db_root'], meta_data_root) if meta_data_root is not None \
            else None
        # semantic segmentation labels root
        sem_root = self.db_info['semantic_root'] if 'semantic_root' in self.db_info else None
        self.sem_root = os.path.join(self.db_info['db_root'], sem_root) if sem_root is not None \
            else None
        # depth valid mask labels root
        depth_mask_root = self.db_info['depth_mask_root'] if 'depth_mask_root' in self.db_info else None
        self.depth_mask_root = os.path.join(self.db_info['db_root'], depth_mask_root) if depth_mask_root is not None \
            else None
        # surface normal labels root
        norm_root = self.db_info['normal_root'] if 'normal_root' in self.db_info else None
        self.norm_root = os.path.join(self.db_info['db_root'], norm_root) if norm_root is not None \
            else None
        # data annotations path
        self.data_annos_path = os.path.join(self.db_info['db_root'], self.db_info['%s_annotations_path' % phase])

        # load annotations
        self.data_info = self.load_annotations()
        whole_data_size = len(self.data_info['files']) 
        
        # sample a subset for training/validation/testing
        # such method is deprecated, each training may get different sample list
        
        cfg_sample_ratio = cfg.data[phase].sample_ratio 
        cfg_sample_size = int(cfg.data[phase].sample_size)
        self.sample_size = int(whole_data_size * cfg_sample_ratio) if cfg_sample_size == -1 \
                           else (cfg_sample_size if cfg_sample_size < whole_data_size else whole_data_size)
        random.seed(100) # set the random seed
        sample_list_of_whole_data = random.sample(list(range(whole_data_size)), self.sample_size)

        self.data_size = self.sample_size
        self.annotations = {'files': [self.data_info['files'][i] for i in sample_list_of_whole_data]}
        self.sample_list = list(range(self.data_size))
        
        # config transforms for the input and label
        self.transforms_cfg = cfg.data[phase]['pipeline']
        self.transforms_lib = 'mono.utils.transform.'

        self.img_file_type = ['.png', '.jpg', '.jpeg', '.bmp', '.tif']
        self.np_file_type = ['.npz', '.npy']

        # update canonical sparce information
        self.data_basic = copy.deepcopy(kwargs)
        canonical = self.data_basic.pop('canonical_space')
        self.data_basic.update(canonical)
        self.disp_scale = 10.0
        self.depth_range = kwargs['depth_range'] # predefined depth range for the network
        self.clip_depth_range = kwargs['clip_depth_range'] # predefined depth range for data processing
        self.depth_normalize = kwargs['depth_normalize']
        
        self.img_transforms = img_transform.Compose(self.build_data_transforms())
        self.EPS = 1e-6

        # self.tmpl_info = ['rgb_sr', 'rgb_pre', 'rgb_next']
        # self.tgt2ref_pose_lookup = {'rgb_sr': 'cam_ex', 'rgb_pre': 'pose_pre', 'rgb_next': 'pose_next'}

        # dataset info
        self.data_name = cfg.data_name
        self.data_type = cfg.data_type # there are mainly four types, i.e. ['rel', 'sfm', 'stereo', 'lidar']
        self.logger = logging.getLogger()
        self.logger.info(f'{self.data_name} in {self.phase} whole data size: {whole_data_size}')
        
        # random crop size for training
        crop_size = kwargs['crop_size']
        shared_array_base = mp.Array(ctypes.c_int32, 2)
        shared_array = np.ctypeslib.as_array(shared_array_base.get_obj())
        shared_array[0] = crop_size[0]
        shared_array[1] = crop_size[1]
        # self.random_crop_size = torch.from_numpy(np.array([0,0])) #torch.from_numpy(shared_array)
        self.random_crop_size = torch.from_numpy(shared_array)

        
    def __name__(self):
        return self.data_name

    def __len__(self):
        return self.data_size

    def load_annotations(self):
        if not os.path.exists(self.data_annos_path):
            raise RuntimeError(f'Cannot find {self.data_annos_path} annotations.')

        with open(self.data_annos_path, 'r')  as f:
            annos = json.load(f)
        return annos
    
    def build_data_transforms(self):
        transforms_list = []
        for transform in self.transforms_cfg:
            args = copy.deepcopy(transform)
            # insert the canonical space configs
            args.update(self.data_basic)
      
            obj_name = args.pop('type')
            obj_path = self.transforms_lib + obj_name
            obj_cls = get_func(obj_path)
            
            obj = obj_cls(**args)
            transforms_list.append(obj)
        return transforms_list
    
        
    def load_data(self, path: str, is_rgb_img: bool=False):
        if not os.path.exists(path):
            self.logger.info(f'>>>>{path} does not exist.')
            # raise RuntimeError(f'{path} does not exist.')

        data_type = os.path.splitext(path)[-1]
        if data_type in self.img_file_type:
            if is_rgb_img:
                data = cv2.imread(path)
            else:
                data = cv2.imread(path, -1)
        elif data_type in self.np_file_type:
            data = np.load(path)
        else:
            raise RuntimeError(f'{data_type} is not supported in current version.')
        
        try:
            return data.squeeze()
        except:
            temp = 1
            raise RuntimeError(f'{path} is not successfully loaded.')
    
    def __getitem__(self, idx: int) -> dict:
        if self.phase == 'test':
            return self.get_data_for_test(idx)
        else:
            return self.get_data_for_trainval(idx)

    def get_data_for_trainval(self, idx: int):
        anno = self.annotations['files'][idx]
        meta_data = self.load_meta_data(anno)
        
        data_path = self.load_data_path(meta_data)
        data_batch = self.load_batch(meta_data, data_path)
        # if data_path['sem_path'] is not None:
        #     print(self.data_name)

        curr_rgb, curr_depth, curr_normal, curr_sem, curr_cam_model = data_batch['curr_rgb'], data_batch['curr_depth'], data_batch['curr_normal'], data_batch['curr_sem'], data_batch['curr_cam_model']
        #curr_stereo_depth = data_batch['curr_stereo_depth']
        
        # A patch for stereo depth dataloader (no need to modify specific datasets)
        if 'curr_stereo_depth' in data_batch.keys():
            curr_stereo_depth = data_batch['curr_stereo_depth']
        else:
            curr_stereo_depth = self.load_stereo_depth_label(None, H=curr_rgb.shape[0], W=curr_rgb.shape[1]) 

        curr_intrinsic = meta_data['cam_in']
        # data augmentation
        transform_paras = dict(random_crop_size = self.random_crop_size) # dict() 
        assert curr_rgb.shape[:2] == curr_depth.shape == curr_normal.shape[:2] == curr_sem.shape
        rgbs, depths, intrinsics, cam_models, normals, other_labels, transform_paras = self.img_transforms(
                                                                   images=[curr_rgb, ], 
                                                                   labels=[curr_depth, ], 
                                                                   intrinsics=[curr_intrinsic,], 
                                                                   cam_models=[curr_cam_model, ],
                                                                   normals = [curr_normal, ],
                                                                   other_labels=[curr_sem, curr_stereo_depth],
                                                                   transform_paras=transform_paras)
        # process sky masks
        sem_mask = other_labels[0].int()
        # clip depth map 
        depth_out = self.normalize_depth(depths[0])
        # set the depth of sky region to the invalid
        depth_out[sem_mask==142] = -1 # self.depth_normalize[1] - 1e-6
        # get inverse depth
        inv_depth = self.depth2invdepth(depth_out, sem_mask==142)
        filename = os.path.basename(meta_data['rgb'])[:-4] + '.jpg'
        curr_intrinsic_mat = self.intrinsics_list2mat(intrinsics[0])
        cam_models_stacks = [
            torch.nn.functional.interpolate(cam_models[0][None, :, :, :], size=(cam_models[0].shape[1]//i, cam_models[0].shape[2]//i), mode='bilinear', align_corners=False).squeeze()
            for i in [2, 4, 8, 16, 32] 
            ]

        # stereo_depth
        if 'label_scale_factor' not in transform_paras.keys():
            transform_paras['label_scale_factor'] = 1
        stereo_depth_pre_trans = other_labels[1] * (other_labels[1] > 0.3) * (other_labels[1] < 200)
        stereo_depth = stereo_depth_pre_trans * transform_paras['label_scale_factor']
        stereo_depth = self.normalize_depth(stereo_depth)

        pad = transform_paras['pad'] if 'pad' in transform_paras else [0,0,0,0]        
        data = dict(input=rgbs[0],
                    target=depth_out,
                    intrinsic=curr_intrinsic_mat,
                    filename=filename,
                    dataset=self.data_name,
                    cam_model=cam_models_stacks,
                    pad=torch.tensor(pad),
                    data_type=[self.data_type, ],
                    sem_mask=sem_mask.int(),
                    stereo_depth= stereo_depth,
                    normal=normals[0],
                    inv_depth=inv_depth,
                    scale=transform_paras['label_scale_factor'])
        return data

    def get_data_for_test(self, idx: int):
        anno = self.annotations['files'][idx]
        meta_data = self.load_meta_data(anno)
        data_path = self.load_data_path(meta_data)
        data_batch = self.load_batch(meta_data, data_path)
        # load data
        curr_rgb, curr_depth, curr_normal, curr_cam_model = data_batch['curr_rgb'], data_batch['curr_depth'], data_batch['curr_normal'], data_batch['curr_cam_model']
        ori_curr_intrinsic = meta_data['cam_in']

        # get crop size
        transform_paras = dict()
        rgbs, depths, intrinsics, cam_models, _, other_labels, transform_paras = self.img_transforms(
                                                                   images=[curr_rgb,],  #+ tmpl_rgbs, 
                                                                   labels=[curr_depth, ], 
                                                                   intrinsics=[ori_curr_intrinsic, ], # * (len(tmpl_rgbs) + 1), 
                                                                   cam_models=[curr_cam_model, ],
                                                                   transform_paras=transform_paras)
        # depth in original size and orignial metric***
        depth_out = self.clip_depth(curr_depth) * self.depth_range[1] # self.clip_depth(depths[0]) #
        inv_depth = self.depth2invdepth(depth_out, np.zeros_like(depth_out, dtype=np.bool))
        filename = os.path.basename(meta_data['rgb'])[:-4] + '.jpg'
        curr_intrinsic_mat = self.intrinsics_list2mat(intrinsics[0])
        ori_curr_intrinsic_mat = self.intrinsics_list2mat(ori_curr_intrinsic)

        pad = transform_paras['pad'] if 'pad' in transform_paras else [0,0,0,0]
        scale_ratio = transform_paras['label_scale_factor'] if 'label_scale_factor' in transform_paras else 1.0
        cam_models_stacks = [
            torch.nn.functional.interpolate(cam_models[0][None, :, :, :], size=(cam_models[0].shape[1]//i, cam_models[0].shape[2]//i), mode='bilinear', align_corners=False).squeeze()
            for i in [2, 4, 8, 16, 32] 
            ]    
        raw_rgb = torch.from_numpy(curr_rgb)
        curr_normal = torch.from_numpy(curr_normal.transpose((2,0,1)))
        

        data = dict(input=rgbs[0],
                    target=depth_out,
                    intrinsic=curr_intrinsic_mat,
                    filename=filename,
                    dataset=self.data_name,
                    cam_model=cam_models_stacks,
                    pad=pad,
                    scale=scale_ratio,
                    raw_rgb=raw_rgb,
                    sample_id=idx,
                    data_path=meta_data['rgb'],
                    inv_depth=inv_depth,
                    normal=curr_normal,
                    )
        return data
    
    def load_data_path(self, meta_data):
        curr_rgb_path = os.path.join(self.data_root, meta_data['rgb'])
        curr_depth_path = os.path.join(self.depth_root, meta_data['depth'])
        curr_sem_path = os.path.join(self.sem_root, meta_data['sem']) \
            if self.sem_root is not None and ('sem' in meta_data) and (meta_data['sem'] is not None)  \
            else None
        # matterport3d separates xyz into three images
        if ('normal' in meta_data) and (meta_data['normal'] is not None) and (self.norm_root is not None):
            if isinstance(meta_data['normal'], dict):
                curr_norm_path = {}
                for k,v in meta_data['normal'].items():
                    curr_norm_path[k] = os.path.join(self.norm_root, v)
            else:
                curr_norm_path = os.path.join(self.norm_root, meta_data['normal'])
        else:
            curr_norm_path = None
        curr_depth_mask_path = os.path.join(self.depth_mask_root, meta_data['depth_mask']) \
            if self.depth_mask_root is not None and ('depth_mask' in meta_data) and (meta_data['depth_mask'] is not None)  \
            else None

        if ('disp' in meta_data) and (meta_data['disp'] is not None) and (self.disp_root is not None):
            if isinstance(meta_data['disp'], dict):
                curr_disp_path = {}
                for k,v in meta_data['disp'].items():
                    curr_disp_path[k] = os.path.join(self.disp_root, v)
            else:
                curr_disp_path = os.path.join(self.disp_root, meta_data['disp'])
        else:
            curr_disp_path = None

        data_path=dict(
            rgb_path=curr_rgb_path,
            depth_path=curr_depth_path,
            sem_path=curr_sem_path,
            normal_path=curr_norm_path,
            disp_path=curr_disp_path,
            depth_mask_path=curr_depth_mask_path,
            )
        return data_path
    
    def load_batch(self, meta_data, data_path):
        curr_intrinsic = meta_data['cam_in']
        # load rgb/depth
        curr_rgb, curr_depth = self.load_rgb_depth(data_path['rgb_path'], data_path['depth_path'])
        # get semantic labels
        curr_sem = self.load_sem_label(data_path['sem_path'], curr_depth)
        # create camera model
        curr_cam_model = self.create_cam_model(curr_rgb.shape[0], curr_rgb.shape[1], curr_intrinsic)       
        # get normal labels
        curr_normal = self.load_norm_label(data_path['normal_path'], H=curr_rgb.shape[0], W=curr_rgb.shape[1]) 
        # get depth mask
        depth_mask = self.load_depth_valid_mask(data_path['depth_mask_path'])
        curr_depth[~depth_mask] = -1
        # get stereo depth
        curr_stereo_depth = self.load_stereo_depth_label(data_path['disp_path'], H=curr_rgb.shape[0], W=curr_rgb.shape[1]) 

        data_batch = dict(
            curr_rgb = curr_rgb,
            curr_depth = curr_depth,
            curr_sem = curr_sem,
            curr_normal = curr_normal,
            curr_cam_model=curr_cam_model,
            curr_stereo_depth=curr_stereo_depth,
        )
        return data_batch


    def clip_depth(self, depth: np.array) -> np.array:
        depth[(depth>self.clip_depth_range[1]) | (depth<self.clip_depth_range[0])] = -1
        depth /= self.depth_range[1]
        depth[depth<self.EPS] = -1
        return depth
    
    def normalize_depth(self, depth: np.array) -> np.array:
        depth /= self.depth_range[1]
        depth[depth<self.EPS] = -1
        return depth
    
    def process_depth(self, depth: np.array, rgb:np.array=None):
        return depth
    
    def create_cam_model(self, H : int, W : int, intrinsics : list) -> np.array:
        """
        Encode the camera model (focal length and principle point) to a 4-channel map. 
        """
        fx, fy, u0, v0 = intrinsics
        f = (fx + fy) / 2.0
        # principle point location
        x_row = np.arange(0, W).astype(np.float32)
        x_row_center_norm = (x_row - u0) / W
        x_center = np.tile(x_row_center_norm, (H, 1)) # [H, W]

        y_col = np.arange(0, H).astype(np.float32) 
        y_col_center_norm = (y_col - v0) / H
        y_center = np.tile(y_col_center_norm, (W, 1)).T

        # FoV
        fov_x = np.arctan(x_center / (f / W))
        fov_y =  np.arctan(y_center/ (f / H))

        cam_model = np.stack([x_center, y_center, fov_x, fov_y], axis=2)
        return cam_model
    
    def check_data(self, data_dict : dict):
        for k, v in data_dict.items():
            if v is None:
                # print(f'{self.data_name}, {k} cannot be read!')
                self.logger.info(f'{self.data_name}, {k} cannot be read!')
    
    def intrinsics_list2mat(self, intrinsics: torch.tensor) -> torch.tensor:
        """
        Create camera intrinsic matrix.
        Args:
            intrinsics (torch.tensor, [4,]): list of camera intrinsic parameters.
        returns:
            intrinsics_mat (torch.tensor, [3x3]): camera intrinsic parameters matrix.
        """
        intrinsics_mat = torch.zeros((3,3)).float()
        intrinsics_mat[0, 0] = intrinsics[0]
        intrinsics_mat[1, 1] = intrinsics[1]
        intrinsics_mat[0, 2] = intrinsics[2]
        intrinsics_mat[1, 2] = intrinsics[3]
        intrinsics_mat[2, 2] = 1.0
        return intrinsics_mat
        
    # def load_tmpl_image(self, curr_rgb: np.array, meta_data: dict) -> dict:
    #     """
    #     Load  consecutive RGB frames.
    #     Args:
    #         anno: the annotation for this group.
    #         curr_rgb: rgb image of the current frame.
    #         meta_data: meta data information.
    #     Returns:
    #         tmpl_annos: temporal rgbs.
    #     """
    #     w_tmpl = False
        
    #     tmpl_list = []
    #     # organize temporal annotations 
    #     for i in self.tmpl_info:
    #         if (i in meta_data) and (meta_data[i] is not None) and os.path.exists(os.path.join(self.data_root, meta_data[i])):
    #             tmpl_list.append(os.path.join(self.data_root, meta_data[i]))

    #     if len(tmpl_list) == 0:
    #         rgb_tmpl = curr_rgb.copy()
    #     else:
    #         id = np.random.randint(len(tmpl_list))
    #         rgb_tmpl = self.load_data(tmpl_list[id], is_rgb_img=True)
    #         w_tmpl = True

    #     tmpl_annos = dict(
    #         tmpl_rgb_list = [rgb_tmpl,],
    #         w_tmpl = w_tmpl
    #     )
    #     return tmpl_annos

    def load_meta_data(self, anno: dict) -> dict:
        """
        Load meta data information.
        """
        if self.meta_data_root is not None and ('meta_data' in anno or 'meta' in anno):
            meta_data_path = os.path.join(self.meta_data_root, anno['meta_data']) if 'meta_data' in anno else os.path.join(self.meta_data_root, anno['meta'])
            with open(meta_data_path, 'rb') as f:
                meta_data = pickle.load(f)
            meta_data.update(anno)
        else:
            meta_data = anno
        return meta_data
    
    def load_rgb_depth(self, rgb_path: str, depth_path: str):
        """
        Load the rgb and depth map with the paths.
        """
        rgb = self.load_data(rgb_path, is_rgb_img=True)
        if rgb is None:
            self.logger.info(f'>>>>{rgb_path} has errors.')
       
        depth = self.load_data(depth_path)
        if depth is None:
            self.logger.info(f'{depth_path} has errors.')
        
        # self.check_data(dict(
        #     rgb_path=rgb,
        #     depth_path=depth,
        # ))
        depth = depth.astype(np.float)
        # if depth.shape != rgb.shape[:2]:
        #     print(f'no-equal in {self.data_name}')
        #     depth = cv2.resize(depth, rgb.shape[::-1][1:])
        
        depth  = self.process_depth(depth, rgb)
        return rgb, depth
    
    def load_sem_label(self, sem_path, depth=None, sky_id=142) -> np.array:
        H, W = depth.shape
        # if sem_path is not None:
        #     print(self.data_name)
        sem_label = cv2.imread(sem_path, 0) if sem_path is not None \
            else np.ones((H, W), dtype=np.int) * -1
        if sem_label is None:
            sem_label = np.ones((H, W), dtype=np.int) * -1
        # set dtype to int before 
        sem_label = sem_label.astype(np.int) 
        sem_label[sem_label==255] = -1
        
        # mask invalid sky region
        mask_depth_valid = depth > 1e-8
        invalid_sky_region = (sem_label==142) & (mask_depth_valid)
        if self.data_type in ['lidar', 'sfm', 'denselidar', 'denselidar_nometric']:
            sem_label[invalid_sky_region] = -1
        return sem_label
    
    def load_depth_valid_mask(self, depth_mask_path, depth=None) -> np.array:
        if depth_mask_path is None:
            return np.ones_like(depth, dtype=np.bool)
        data_type = os.path.splitext(depth_mask_path)[-1]
        if data_type in self.img_file_type:
            data = cv2.imread(depth_mask_path, -1)
        elif data_type in self.np_file_type:
            data = np.load(depth_mask_path)
        else:
            raise RuntimeError(f'{data_type} is not supported in current version.')
        data = data.astype(np.bool)
        return data
        
    def load_norm_label(self, norm_path, H, W):
        norm_gt = np.zeros((H, W, 3)).astype(np.float32)
        return norm_gt

    def load_stereo_depth_label(self, disp_path, H, W):
        stereo_depth_gt = np.zeros((H, W, 1)).astype(np.float32)
        return stereo_depth_gt

    def depth2invdepth(self, depth, sky_mask):
        inv_depth = 1.0 / depth * self.disp_scale
        inv_depth[depth<1e-6] = -1.0
        inv_depth[inv_depth < 0] = -1.0
        inv_depth[sky_mask] = 0
        return inv_depth


    def set_random_crop_size(self, random_crop_size):
        self.random_crop_size[0] = random_crop_size[0]
        self.random_crop_size[1] = random_crop_size[1]