zoedepth/data/data_mono.py

# 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: Shariq Farooq Bhat

# This file is partly inspired from BTS (https://github.com/cleinc/bts/blob/master/pytorch/bts_dataloader.py); author: Jin Han Lee

# This file may include modifications from author Zhenyu Li

import itertools
import os
import random

import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.utils.data.distributed
from zoedepth.utils.easydict import EasyDict as edict
from PIL import Image, ImageOps
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from zoedepth.utils.config import change_dataset

from .ddad import get_ddad_loader
from .diml_indoor_test import get_diml_indoor_loader
from .diml_outdoor_test import get_diml_outdoor_loader
from .diode import get_diode_loader
from .hypersim import get_hypersim_loader
from .ibims import get_ibims_loader
from .sun_rgbd_loader import get_sunrgbd_loader
from .vkitti import get_vkitti_loader
from .vkitti2 import get_vkitti2_loader
from .u4k import get_u4k_loader
from .middleburry import get_mid_loader
from .gta import get_gta_loader
from .preprocess import CropParams, get_white_border, get_black_border
import copy
from zoedepth.utils.misc import get_boundaries
from zoedepth.models.base_models.midas import Resize

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})


# def preprocessing_transforms(mode, **kwargs):
#     return transforms.Compose([
#         ToTensor(mode=mode, **kwargs)
#     ])

def preprocessing_transforms(mode, sec_stage=False, **kwargs):
    return transforms.Compose([
        ToTensor(mode=mode, sec_stage=sec_stage, **kwargs)
    ])



class DepthDataLoader(object):
    def __init__(self, config, mode, device='cpu', transform=None, **kwargs):
        """
        Data loader for depth datasets

        Args:
            config (dict): Config dictionary. Refer to utils/config.py
            mode (str): "train" or "online_eval"
            device (str, optional): Device to load the data on. Defaults to 'cpu'.
            transform (torchvision.transforms, optional): Transform to apply to the data. Defaults to None.
        """

        self.config = config

        if config.dataset == 'ibims':
            self.data = get_ibims_loader(config, batch_size=1, num_workers=1)
            return

        if config.dataset == 'sunrgbd':
            self.data = get_sunrgbd_loader(
                data_dir_root=config.sunrgbd_root, batch_size=1, num_workers=1)
            return

        if config.dataset == 'diml_indoor':
            self.data = get_diml_indoor_loader(
                data_dir_root=config.diml_indoor_root, batch_size=1, num_workers=1)
            return

        if config.dataset == 'diml_outdoor':
            self.data = get_diml_outdoor_loader(
                data_dir_root=config.diml_outdoor_root, batch_size=1, num_workers=1)
            return

        if "diode" in config.dataset:
            self.data = get_diode_loader(
                config[config.dataset+"_root"], batch_size=1, num_workers=1)
            return

        if config.dataset == 'hypersim_test':
            self.data = get_hypersim_loader(
                config.hypersim_test_root, batch_size=1, num_workers=1)
            return

        if config.dataset == 'vkitti':
            self.data = get_vkitti_loader(
                config.vkitti_root, batch_size=1, num_workers=1)
            return

        if config.dataset == 'vkitti2':
            self.data = get_vkitti2_loader(
                config.vkitti2_root, batch_size=1, num_workers=1)
            return

        if config.dataset == 'ddad':
            self.data = get_ddad_loader(config.ddad_root, resize_shape=(
                352, 1216), batch_size=1, num_workers=1)
            return
        
        # under construction
        if config.dataset == 'u4k':
            self.data = get_u4k_loader(config, mode, transform)
            return
        
        if config.dataset == 'mid':
            self.data = get_mid_loader(config, mode, transform)
            return
        
        if config.dataset == 'gta':
            self.data = get_gta_loader(config, mode, transform)
            return
        
        img_size = self.config.get("img_size", None)
        img_size = img_size if self.config.get(
            "do_input_resize", False) else None
        if transform is None:
            # transform = preprocessing_transforms(mode, size=img_size)
            transform = preprocessing_transforms(mode, size=img_size, sec_stage=config.get("sec_stage", False))

        if mode == 'train':

            Dataset = DataLoadPreprocess
            self.training_samples = Dataset(
                config, mode, transform=transform, device=device)

            if config.distributed:
                self.train_sampler = torch.utils.data.distributed.DistributedSampler(
                    self.training_samples)
            else:
                self.train_sampler = None

            self.data = DataLoader(self.training_samples,
                                   batch_size=config.batch_size,
                                   shuffle=(self.train_sampler is None),
                                   num_workers=config.workers,
                                   pin_memory=True,
                                   persistent_workers=True,
                                #    prefetch_factor=2,
                                   sampler=self.train_sampler)

        elif mode == 'online_eval':
            self.testing_samples = DataLoadPreprocess(
                config, mode, transform=transform)
            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
                self.eval_sampler = None
            else:
                self.eval_sampler = None
            self.data = DataLoader(self.testing_samples, 1,
                                   shuffle=False,
                                   num_workers=1,
                                   pin_memory=False,
                                   sampler=self.eval_sampler)

        elif mode == 'test':
            self.testing_samples = DataLoadPreprocess(
                config, mode, transform=transform)
            self.data = DataLoader(self.testing_samples,
                                   1, shuffle=False, num_workers=1)

        else:
            print(
                'mode should be one of \'train, test, online_eval\'. Got {}'.format(mode))


def repetitive_roundrobin(*iterables):
    """
    cycles through iterables but sample wise
    first yield first sample from first iterable then first sample from second iterable and so on
    then second sample from first iterable then second sample from second iterable and so on

    If one iterable is shorter than the others, it is repeated until all iterables are exhausted
    repetitive_roundrobin('ABC', 'D', 'EF') --> A D E B D F C D E
    """
    # Repetitive roundrobin
    iterables_ = [iter(it) for it in iterables]
    exhausted = [False] * len(iterables)
    while not all(exhausted):
        for i, it in enumerate(iterables_):
            try:
                yield next(it)
            except StopIteration:
                exhausted[i] = True
                iterables_[i] = itertools.cycle(iterables[i])
                # First elements may get repeated if one iterable is shorter than the others
                yield next(iterables_[i])


class RepetitiveRoundRobinDataLoader(object):
    def __init__(self, *dataloaders):
        self.dataloaders = dataloaders

    def __iter__(self):
        return repetitive_roundrobin(*self.dataloaders)

    def __len__(self):
        # First samples get repeated, thats why the plus one
        return len(self.dataloaders) * (max(len(dl) for dl in self.dataloaders) + 1)


class MixedNYUKITTI(object):
    def __init__(self, config, mode, device='cpu', **kwargs):
        config = edict(config)
        config.workers = config.workers // 2
        self.config = config
        nyu_conf = change_dataset(edict(config), 'nyu')
        kitti_conf = change_dataset(edict(config), 'kitti')

        # make nyu default for testing
        self.config = config = nyu_conf
        img_size = self.config.get("img_size", None)
        img_size = img_size if self.config.get(
            "do_input_resize", False) else None
        if mode == 'train':
            nyu_loader = DepthDataLoader(
                nyu_conf, mode, device=device, transform=preprocessing_transforms(mode, size=img_size)).data
            kitti_loader = DepthDataLoader(
                kitti_conf, mode, device=device, transform=preprocessing_transforms(mode, size=img_size)).data
            # It has been changed to repetitive roundrobin
            self.data = RepetitiveRoundRobinDataLoader(
                nyu_loader, kitti_loader)
        else:
            self.data = DepthDataLoader(nyu_conf, mode, device=device).data


def remove_leading_slash(s):
    if s[0] == '/' or s[0] == '\\':
        return s[1:]
    return s


class CachedReader:
    def __init__(self, shared_dict=None):
        if shared_dict:
            self._cache = shared_dict
        else:
            self._cache = {}

    def open(self, fpath):
        im = self._cache.get(fpath, None)
        if im is None:
            im = self._cache[fpath] = Image.open(fpath)
        return im


class ImReader:
    def __init__(self):
        pass

    # @cache
    def open(self, fpath):
        return Image.open(fpath)


class DataLoadPreprocess(Dataset):
    def __init__(self, config, mode, transform=None, is_for_online_eval=False, **kwargs):
        self.config = config
        if mode == 'online_eval':
            with open(config.filenames_file_eval, 'r') as f:
                self.filenames = f.readlines()
        else:
            with open(config.filenames_file, 'r') as f:
                self.filenames = f.readlines()

        self.sec_stage = self.config.get("sec_stage", False)
        # self.crop_size = [120, 160] # 1/4
        self.crop_size = [120*2, 160*2] # 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(60))
        self.overlap_length_w = self.config.get("overlap_length_w", int(80))
        print("current overlap_length_h and overlap_length_w are {} and {}".format(self.overlap_length_h, self.overlap_length_w))
        
        self.mode = mode
        self.transform = transform
        self.to_tensor = ToTensor(mode)
        self.is_for_online_eval = is_for_online_eval
        if config.use_shared_dict:
            self.reader = CachedReader(config.shared_dict)
        else:
            self.reader = ImReader()

    def postprocess(self, sample):
        return sample

    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):
        sample_path = self.filenames[idx]
        focal = float(sample_path.split()[2])
        sample = {}

        height=480 
        width=640
        
        if self.mode == 'train':
            if self.config.dataset == 'kitti' and self.config.use_right and random.random() > 0.5:
                image_path = os.path.join(
                    self.config.data_path, remove_leading_slash(sample_path.split()[3]))
                depth_path = os.path.join(
                    self.config.gt_path, remove_leading_slash(sample_path.split()[4]))
            else:
                image_path = os.path.join(
                    self.config.data_path, remove_leading_slash(sample_path.split()[0]))
                depth_path = os.path.join(
                    self.config.gt_path, remove_leading_slash(sample_path.split()[1]))

            image = self.reader.open(image_path)
            depth_gt = self.reader.open(depth_path)
            w, h = image.size

            if self.config.do_kb_crop:
                height = image.height
                width = image.width
                top_margin = int(height - 352)
                left_margin = int((width - 1216) / 2)
                depth_gt = depth_gt.crop(
                    (left_margin, top_margin, left_margin + 1216, top_margin + 352))
                image = image.crop(
                    (left_margin, top_margin, left_margin + 1216, top_margin + 352))

            # Avoid blank boundaries due to pixel registration?
            # Train images have white border. Test images have black border.
            if self.config.dataset == 'nyu' and self.config.avoid_boundary:
                # print("Avoiding Blank Boundaries!")
                # We just crop and pad again with reflect padding to original size
                # original_size = image.size
                crop_params = get_white_border(np.array(image, dtype=np.uint8))
                image = image.crop((crop_params.left, crop_params.top, crop_params.right, crop_params.bottom))
                depth_gt = depth_gt.crop((crop_params.left, crop_params.top, crop_params.right, crop_params.bottom))

                # Use reflect padding to fill the blank
                image = np.array(image)
                image = np.pad(image, ((crop_params.top, h - crop_params.bottom), (crop_params.left, w - crop_params.right), (0, 0)), mode='reflect')
                image = Image.fromarray(image)

                depth_gt = np.array(depth_gt)
                depth_gt = np.pad(depth_gt, ((crop_params.top, h - crop_params.bottom), (crop_params.left, w - crop_params.right)), 'constant', constant_values=0)
                depth_gt = Image.fromarray(depth_gt)


            if self.config.do_random_rotate and (self.config.aug):
                # NOTE: YES!
                random_angle = (random.random() - 0.5) * 2 * self.config.degree
                image = self.rotate_image(image, random_angle)
                depth_gt = self.rotate_image(
                    depth_gt, random_angle, flag=Image.NEAREST)

            image = np.asarray(image, dtype=np.float32) / 255.0
            depth_gt = np.asarray(depth_gt, dtype=np.float32)
            depth_gt = np.expand_dims(depth_gt, axis=2)
            disp_gt_copy = depth_gt[:, :, 0].copy()

            if self.config.dataset == 'nyu':
                depth_gt = depth_gt / 1000.0
            else:
                depth_gt = depth_gt / 256.0

            # if self.config.aug and (self.config.random_crop):
            #     image, depth_gt = self.random_crop(
            #         image, depth_gt, self.config.input_height, self.config.input_width)
            
            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 * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2])
                    bboxs_shift = torch.tensor([crop_x1_shift / width * 160 * 2, crop_y1_shift / height * 120 * 2, crop_x2_shift / width * 160 * 2, crop_y2_shift / height * 120 * 2])
                    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 * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2]) # 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}
            
            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.config.aug and self.config.random_translate:
                image, depth_gt = self.random_translate(image, depth_gt, self.config.max_translation)

        else:
            if self.mode == 'online_eval':
                data_path = self.config.data_path_eval
            else:
                data_path = self.config.data_path

            image_path = os.path.join(
                data_path, remove_leading_slash(sample_path.split()[0]))
            image = np.asarray(self.reader.open(image_path),
                               dtype=np.float32) / 255.0

            if self.mode == 'online_eval':
                gt_path = self.config.gt_path_eval
                depth_path = os.path.join(
                    gt_path, remove_leading_slash(sample_path.split()[1]))
                has_valid_depth = False
                try:
                    depth_gt = self.reader.open(depth_path)
                    has_valid_depth = True
                except IOError:
                    depth_gt = False
                    # print('Missing gt for {}'.format(image_path))

                if has_valid_depth:
                    depth_gt = np.asarray(depth_gt, dtype=np.float32)
                    depth_gt = np.expand_dims(depth_gt, axis=2)
                    disp_gt_copy = depth_gt[:, :, 0].copy()
                    if self.config.dataset == 'nyu':
                        depth_gt = depth_gt / 1000.0
                    else:
                        depth_gt = depth_gt / 256.0

                    mask = np.logical_and(
                        depth_gt >= self.config.min_depth, depth_gt <= self.config.max_depth).squeeze()[None, ...]
                else:
                    mask = False

            if self.config.do_kb_crop:
                height = image.shape[0]
                width = image.shape[1]
                top_margin = int(height - 352)
                left_margin = int((width - 1216) / 2)
                image = image[top_margin:top_margin + 352,
                              left_margin:left_margin + 1216, :]
                if self.mode == 'online_eval' and has_valid_depth:
                    depth_gt = depth_gt[top_margin:top_margin +
                                        352, left_margin:left_margin + 1216, :]

            # NOTE: start insert something new for sec_stage training
            if self.sec_stage:
                img_temp = copy.deepcopy(image)
                depth_gt_temp = copy.deepcopy(depth_gt)
            
                x_start, y_start = [0, 240], [0, 320]
                # x_start, y_start = [0 + 3 * self.overlap / 2, 120 + self.overlap / 2, 240 - self.overlap / 2, 360 - 3 * self.overlap / 2], \
                #     [0 + 3 * self.overlap / 2, 160 + self.overlap / 2, 320 - self.overlap / 2, 480 - 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+240), int(y), int(y+320))
                        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 * 160 * 2, crop_y1 / height * 120 * 2, crop_x2 / width * 160 * 2, crop_y2 / height * 120 * 2])
                        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)
                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': has_valid_depth,
                          'image_path': sample_path.split()[0], 'depth_path': sample_path.split()[1],
                          'mask': mask, 'image_raw': image.copy(), 'disp_gt_edges': disp_gt_edges}
                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}

        if (self.mode == 'train') or ('has_valid_depth' in sample and sample['has_valid_depth']):
            mask = np.logical_and(depth_gt > self.config.min_depth,
                                  depth_gt < self.config.max_depth).squeeze()[None, ...]
            sample['mask'] = mask

        if self.transform:
            # sample = self.transform(sample)
            sample['img_temp'] = img_temp
            sample['depth_gt_temp'] = depth_gt_temp
            sample = self.transform(sample)

        sample = self.postprocess(sample)
        sample['dataset'] = self.config.dataset
        sample = {**sample, 'image_path': sample_path.split()[0], 'depth_path': sample_path.split()[1]}

        return sample

    def rotate_image(self, image, angle, flag=Image.BILINEAR):
        result = image.rotate(angle, resample=flag)
        return result

    def random_crop(self, img, depth, height, width):
        assert img.shape[0] >= height
        assert img.shape[1] >= width
        assert img.shape[0] == depth.shape[0]
        assert img.shape[1] == depth.shape[1]
        x = random.randint(0, img.shape[1] - width)
        y = random.randint(0, img.shape[0] - height)
        img = img[y:y + height, x:x + width, :]
        depth = depth[y:y + height, x:x + width, :]

        return img, depth
    
    def random_translate(self, img, depth, max_t=20):
        assert img.shape[0] == depth.shape[0]
        assert img.shape[1] == depth.shape[1]
        p = self.config.translate_prob
        do_translate = random.random()
        if do_translate > p:
            return img, depth
        x = random.randint(-max_t, max_t)
        y = random.randint(-max_t, max_t)
        M = np.float32([[1, 0, x], [0, 1, y]])
        # print(img.shape, depth.shape)
        img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
        depth = cv2.warpAffine(depth, M, (depth.shape[1], depth.shape[0]))
        depth = depth.squeeze()[..., None]  # add channel dim back. Affine warp removes it
        # print("after", img.shape, depth.shape)
        return img, depth

    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 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 __len__(self):
        return len(self.filenames)


# class ToTensor(object):
#     def __init__(self, mode, do_normalize=False, size=None):
#         self.mode = mode
#         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)
#         else:
#             self.resize = nn.Identity()

#     def __call__(self, sample):
#         image, focal = sample['image'], sample['focal']
#         image = self.to_tensor(image)
#         image = self.normalize(image)
#         image = self.resize(image)

#         if self.mode == 'test':
#             return {'image': image, 'focal': focal}

#         depth = sample['depth']
#         if self.mode == 'train':
#             depth = self.to_tensor(depth)
#             return {**sample, 'image': image, 'depth': depth, 'focal': focal}
#         else:
#             has_valid_depth = sample['has_valid_depth']
#             image = self.resize(image)
#             return {**sample, 'image': image, 'depth': depth, 'focal': focal, 'has_valid_depth': has_valid_depth,
#                     'image_path': sample['image_path'], 'depth_path': sample['depth_path']}

#     def to_tensor(self, pic):
#         if not (_is_pil_image(pic) or _is_numpy_image(pic)):
#             raise TypeError(
#                 'pic should be PIL Image or ndarray. Got {}'.format(type(pic)))

#         if isinstance(pic, np.ndarray):
#             img = torch.from_numpy(pic.transpose((2, 0, 1)))
#             return img

#         # handle PIL Image
#         if pic.mode == 'I':
#             img = torch.from_numpy(np.array(pic, np.int32, copy=False))
#         elif pic.mode == 'I;16':
#             img = torch.from_numpy(np.array(pic, np.int16, copy=False))
#         else:
#             img = torch.ByteTensor(
#                 torch.ByteStorage.from_buffer(pic.tobytes()))
#         # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
#         if pic.mode == 'YCbCr':
#             nchannel = 3
#         elif pic.mode == 'I;16':
#             nchannel = 1
#         else:
#             nchannel = len(pic.mode)
#         img = img.view(pic.size[1], pic.size[0], nchannel)

#         img = img.transpose(0, 1).transpose(0, 2).contiguous()
#         if isinstance(img, torch.ByteTensor):
#             return img.float()
#         else:
#             return img


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