zoedepth/data/middleburry.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: 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


from .u4k import U4KDataset, remove_leading_slash

import re
import numpy as np
import sys
import matplotlib.pyplot as plt

def readPFM(file):
    file = open(file, 'rb')

    color = None
    width = None
    height = None
    scale = None
    endian = None

    header = file.readline().rstrip()
    if (sys.version[0]) == '3':
        header = header.decode('utf-8')
    if header == 'PF':
        color = True
    elif header == 'Pf':
        color = False
    else:
        raise Exception('Not a PFM file.')

    if (sys.version[0]) == '3':
        dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('utf-8'))
    else:
        dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline())
    if dim_match:
        width, height = map(int, dim_match.groups())
    else:
        raise Exception('Malformed PFM header.')

    if (sys.version[0]) == '3':
        scale = float(file.readline().rstrip().decode('utf-8'))
    else:
        scale = float(file.readline().rstrip())
        
    if scale < 0: # little-endian
        endian = '<'
        scale = -scale
    else:
        endian = '>' # big-endian

    data = np.fromfile(file, endian + 'f')
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    return data, scale

class MiddleBurry(U4KDataset):
    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_l, depth_map_l = line.strip().split(" ")

                    img_info_l['depth_map_path'] = osp.join(data_root, remove_leading_slash(depth_map_l))
                    img_info_l['img_path'] = osp.join(data_root, remove_leading_slash(img_l))
                    img_info_l['depth_fields'] = []

                    filename = img_info_l['depth_map_path']
                    ext_name_l = filename.replace('disp0.pfm', 'calib.txt')
                    with open(ext_name_l, 'r') as f:
                        ext_l = f.readlines()

                    cam_info = ext_l[0].strip()
                    cam_info_f = float(cam_info.split(' ')[0].split('[')[1])
                    base = float(ext_l[3].strip().split('=')[1])
                    doffs = float(ext_l[2].strip().split('=')[1])
    
                    f = cam_info_f
                    img_info_l['focal'] = f
                    base = base
                    img_info_l['depth_factor'] = base * f
                    img_info_l['doffs'] = doffs

                    img_infos.append(img_info_l)
        else:
            raise NotImplementedError 

        # github issue:: make sure the same order
        img_infos = sorted(img_infos, key=lambda x: x['img_path'])
        if self.mode == 'train':
            img_infos = img_infos * 100
        return img_infos

    
    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
        height = 1840
        width = 2300

        image = Image.open(img_file_path).convert("RGB")
        image = np.asarray(image, dtype=np.uint8) / 255.0
        image = image.astype(np.float32)

        disp_gt, scale = readPFM(disp_path)
        disp_gt = disp_gt.astype(np.float32)

        h, w, _ = image.shape
        h_start = int(h / 2 - height / 2)
        h_end = h_start + height
        w_start = int(w / 2 - width / 2)
        w_end = w_start + width
        image = image[h_start:h_end, w_start:w_end, :]
        disp_gt = disp_gt[h_start:h_end, w_start:w_end]

        disp_gt_copy = disp_gt.copy()
        disp_gt = disp_gt[..., np.newaxis]
        invalid_mask = disp_gt == np.inf
        
        depth_gt = depth_factor / (disp_gt + self.data_infos[idx]['doffs'])
        depth_gt = depth_gt / 1000
        depth_gt[invalid_mask] = 0 # set to a invalid number
        disp_gt_copy[invalid_mask[:, :, 0]] = 0
        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])

                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


            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)
                    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['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['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 = []
                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)

                image = img_crops
                bboxs_roi = torch.stack(bboxs_roi, 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}
                if bboxs_roi is not None:
                    sample['bbox'] = bboxs_roi
                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
                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_mid_loader(config, mode, transform):
    if mode == 'train':
        log = 0
        dataset = MiddleBurry(config, mode, config.data_path, config.filenames_train)
        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 == 'online_eval':
        dataset = MiddleBurry(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 = MiddleBurry(config, mode, config.data_path, config.filenames_test)
        dataloader = DataLoader(dataset, 1, shuffle=False, num_workers=1)
        
    return dataloader