training/mono/datasets/any_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 copy
from .__base_dataset__ import BaseDataset
import mono.utils.transform as img_transform


class AnyDataset(BaseDataset):
    def __init__(self, cfg, phase, **kwargs):        
        super(AnyDataset, self).__init__(
            cfg=cfg,
            phase=phase,
            **kwargs)

        self.cfg = cfg
        self.phase = phase
        self.mldb_info = kwargs['mldb_info']
        
        # root dir for data
        self.data_root = os.path.join(self.mldb_info['mldb_root'], self.mldb_info['data_root'])
        # depth/disp data root
        disp_root = self.mldb_info['disp_root'] if 'disp_root' in self.mldb_info else None
        self.disp_root = os.path.join(self.mldb_info['mldb_root'], disp_root) if disp_root is not None else None
        depth_root = self.mldb_info['depth_root'] if 'depth_root' in self.mldb_info else None
        self.depth_root = os.path.join(self.mldb_info['mldb_root'], depth_root) if depth_root is not None \
            else self.data_root
        # meta data root
        meta_data_root = self.mldb_info['meta_data_root'] if 'meta_data_root' in self.mldb_info else None
        self.meta_data_root = os.path.join(self.mldb_info['mldb_root'], meta_data_root) if meta_data_root is not None \
            else None
        # semantic segmentation labels root
        sem_root = self.mldb_info['semantic_root'] if 'semantic_root' in self.mldb_info else None
        self.sem_root = os.path.join(self.mldb_info['mldb_root'], sem_root) if sem_root is not None \
            else None

        # data annotations path
        self.data_annos_path = '/yvan1/data/NuScenes/NuScenes/annotations/train_ring_annotations.json'  # fill this 

        # load annotations
        annotations = self.load_annotations()
        whole_data_size = len(annotations['files']) 
               
        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)
        sample_list_of_whole_data = list(range(whole_data_size))[:self.sample_size]
        self.data_size = self.sample_size
        sample_list_of_whole_data = random.sample(list(range(whole_data_size)), whole_data_size)
        self.annotations = {'files': [annotations['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.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-8

        self.tmpl_info = ['rgb_sr', 'rgb_pre', 'rgb_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']

    def __getitem__(self, idx: int) -> dict:
        return self.get_data_for_test(idx)

    def get_data_for_test(self, idx: int):
        # basic info
        anno = self.annotations['files'][idx]
        curr_rgb_path = os.path.join(self.data_root, anno['CAM_FRONT_RIGHT']['rgb']) # Lyft: CAM_FRONT_LEFT
        curr_depth_path = os.path.join(self.depth_root, anno['CAM_FRONT_RIGHT']['depth'])
        meta_data = self.load_meta_data(anno['CAM_FRONT_RIGHT'])
        ori_curr_intrinsic = meta_data['cam_in']
        
        curr_rgb, curr_depth = self.load_rgb_depth(curr_rgb_path, curr_depth_path)
        ori_h, ori_w, _ = curr_rgb.shape
        # create camera model
        curr_cam_model = self.create_cam_model(curr_rgb.shape[0], curr_rgb.shape[1], ori_curr_intrinsic)
        # load tmpl rgb info
        # tmpl_annos = self.load_tmpl_annos(anno, curr_rgb, meta_data)
        # tmpl_rgb = tmpl_annos['tmpl_rgb_list'] # list of reference rgbs

        transform_paras = dict()
        rgbs, depths, intrinsics, cam_models, other_labels, transform_paras = self.img_transforms(
                                                                   images=[curr_rgb, ], 
                                                                   labels=[curr_depth, ], 
                                                                   intrinsics=[ori_curr_intrinsic,], 
                                                                   cam_models=[curr_cam_model, ],
                                                                   transform_paras=transform_paras)
        # depth in augmented size
        # depth_out = self.clip_depth(depths[0])
        # depth in original size
        #depth_out = self.clip_depth(curr_depth)
        depth_out = curr_depth

        filename = os.path.basename(curr_rgb_path)
        curr_intrinsic_mat = self.intrinsics_list2mat(intrinsics[0])

        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)
        data = dict(input=rgbs[0],
                    target=depth_out,
                    intrinsic=curr_intrinsic_mat,
                    filename=filename,
                    dataset=self.data_name,
                    cam_model=cam_models_stacks,
                    # ref_input=rgbs[1:],
                    # tmpl_flg=tmpl_annos['w_tmpl'],
                    pad=pad,
                    scale=scale_ratio,
                    raw_rgb=raw_rgb) 
        return data


    def process_depth(self, depth):
        depth[depth>65500] = 0
        depth /= 200.0
        return depth



if __name__ == '__main__':
    from mmcv.utils import Config 
    cfg = Config.fromfile('mono/configs/Apolloscape_DDAD/convnext_base.cascade.1m.sgd.mae.py')
    dataset_i = ApolloscapeDataset(cfg['Apolloscape'], 'train', **cfg.data_basic)
    print(dataset_i)