from .kitti_dataset import KITTIRAWDataset, KITTIOdomDataset, KITTIDepthDataset
from .make3d_dataset import Make3DDataset
from .nyuv2_dataset import NYUDataset
from .cityscapes_dataset import CityscapesDataset
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import Sampler
import torch

class CustomSampler(Sampler):
    def __init__(self, dataset, seed=0):
        self.len = len(dataset)
        self.start_iter = 0
        self.epoch = 0
        self.seed = seed

    def __iter__(self):
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        indices = torch.randperm(self.len, generator=g).tolist()
        indices = indices[self.start_iter:]
        return iter(indices)
    
    def __len__(self):
        return self.len

    def set_start_iter(self, start_iter):
        self.start_iter = start_iter

    def set_epoch(self, epoch):
        self.epoch = epoch


class CustomDistributedSampler(DistributedSampler):
    """
    More fine-grained state DataSampler that uses training iteration and epoch
    both for shuffling data. PyTorch DistributedSampler only uses epoch
    for the shuffling and starts sampling data from the start. In case of training
    on very large data, we train for one epoch only and when we resume training,
    we want to resume the data sampler from the training iteration.
    """

    def __init__(self, dataset, seed: int = 0):
        """
        Initializes the instance of StatefulDistributedSampler. Random seed is set
        for the epoch set and data is shuffled. For starting the sampling, use
        the start_iter (set to 0 or set by checkpointing resuming) to
        sample data from the remaining images.
        Args:
            dataset (Dataset): Pytorch dataset that sampler will shuffle
            batch_size (int): batch size we want the sampler to sample
            seed (int): Seed for the torch generator.
        """
        super().__init__(dataset, shuffle=False, seed=seed)

        self.start_iter = 0
        self.total_size = len(dataset) - (len(dataset) % self.num_replicas)
        self.num_samples = self.total_size // self.num_replicas
        # self.num_samples = self.total_size 
        # logging.info(f"rank: {self.rank}: Sampler created...")
        # logging.info(f"sample_num: {len(dataset) * self.num_replicas}")
        # print(f"rank: {self.rank}: Sampler created...")

    def __iter__(self):
        # partition data into num_replicas and optionally shuffle within a rank
        g = torch.Generator()
        g.manual_seed(self.epoch + self.seed)
        indices = torch.randperm(len(self.dataset), generator=g).tolist()
        indices = indices[:self.total_size]
        assert len(indices) == self.total_size

        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples
                
        # resume the sampler
        indices = indices[self.start_iter:]
        return iter(indices)

    def set_start_iter(self, start_iter):
        """
        Set the iteration number from which the sampling should start. This is
        used to find the marker in the data permutation order from where the
        sampler should start sampling.
        """
        self.start_iter = start_iter
datasets/__init__.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.
import os
import json
import numpy as np
import PIL.Image as pil

from .mono_dataset import MonoDataset


class CityscapesDataset(MonoDataset):
    """Cityscapes dataset - this expects triplets of images concatenated into a single wide image,
    which have had the ego car removed (bottom 25% of the image cropped)
    """

    def __init__(self, *args, **kwargs):
        super(CityscapesDataset, self).__init__(*args, **kwargs)

        if self.is_train:
            self.RAW_WIDTH = 1024
            self.RAW_HEIGHT = 384
        else:
            self.RAW_WIDTH = 2048
            self.RAW_HEIGHT = 1024       

    def index_to_folder_and_frame_idx(self, index):
        """Convert index in the dataset to a folder name, frame_idx and any other bits

        txt file is of format:
            ulm ulm_000064_000012
        """
        city, frame_name = self.filenames[index].split()
        side = None
        return city, frame_name, side

    def check_depth(self):
        return False

    def load_intrinsics(self, city, frame_name):
        # adapted from sfmlearner
        if self.is_train:
            camera_file = os.path.join(self.data_path, city, "{}_cam.txt".format(frame_name))
            camera = np.loadtxt(camera_file, delimiter=",")
            fx = camera[0]
            fy = camera[4]
            u0 = camera[2]
            v0 = camera[5]
            intrinsics = np.array([[fx, 0, u0, 0],
                                [0, fy, v0, 0],
                                [0,  0,  1, 0],
                                [0,  0,  0, 1]]).astype(np.float32)

            intrinsics[0, :] /= self.RAW_WIDTH
            intrinsics[1, :] /= self.RAW_HEIGHT
        else:
            split = "test"  # if self.is_train else "val"
            camera_file = os.path.join(self.data_path, 'camera',
                                    split, city, frame_name + '_camera.json')
            with open(camera_file, 'r') as f:
                camera = json.load(f)
            fx = camera['intrinsic']['fx']
            fy = camera['intrinsic']['fy']
            u0 = camera['intrinsic']['u0']
            v0 = camera['intrinsic']['v0']
            intrinsics = np.array([[fx, 0, u0, 0],
                                [0, fy, v0, 0],
                                [0, 0, 1, 0],
                                [0, 0, 0, 1]]).astype(np.float32)
            intrinsics[0, :] /= self.RAW_WIDTH
            intrinsics[1, :] /= self.RAW_HEIGHT * 0.75            
        return intrinsics

    def get_color(self, city, frame_name, side, do_flip):
        if side is not None:
            raise ValueError("Cityscapes dataset doesn't know how to deal with sides")

        inputs = {}
        if self.is_train:
            color = self.loader(self.get_image_path(city, frame_name))
            color = np.array(color)

            h = color.shape[0] // 3
            inputs[("color", -1, -1)] = pil.fromarray(color[:h, :])
            inputs[("color", 0, -1)] = pil.fromarray(color[h:2*h, :])
            inputs[("color", 1, -1)] = pil.fromarray(color[2*h:, :])

            if do_flip:
                for key in inputs:
                    inputs[key] = inputs[key].transpose(pil.FLIP_LEFT_RIGHT)
        else:
            color = self.loader(self.get_image_path(city, frame_name))
            w, h = color.size
            crop_h = h * 3 // 4
            color = color.crop((0, 0, w, crop_h))
            inputs[("color", 0, -1)] = color
        return inputs

    def get_image_path(self, city, frame_name):
        if self.is_train:
            return os.path.join(self.data_path, city, "{}.png".format(frame_name))
        else:
            folder = "leftImg8bit_sequence"
            split = "test"
            image_path = os.path.join(
                self.data_path, folder, split, city, frame_name + '_leftImg8bit.png')
            return image_path            

datasets/cityscapes_dataset.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import skimage.transform
import numpy as np
import PIL.Image as pil

from kitti_utils import generate_depth_map
from .mono_dataset import MonoDataset


class KITTIDataset(MonoDataset):
    """Superclass for different types of KITTI dataset loaders
    """
    def __init__(self, *args, **kwargs):
        super(KITTIDataset, self).__init__(*args, **kwargs)

        # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size.
        # To normalize you need to scale the first row by 1 / image_width and the second row
        # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered.
        # If your principal point is far from the center you might need to disable the horizontal
        # flip augmentation.
        self.K = np.array([[0.58, 0, 0.5, 0],
                           [0, 1.92, 0.5, 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        self.full_res_shape = (1242, 375)
        self.side_map = {"2": 2, "3": 3, "l": 2, "r": 3}

    def check_depth(self):
        line = self.filenames[0].split()
        scene_name = line[0]
        frame_index = int(line[1])

        velo_filename = os.path.join(
            self.data_path,
            scene_name,
            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))

        return os.path.isfile(velo_filename)

    def index_to_folder_and_frame_idx(self, index):
        """Convert index in the dataset to a folder name, frame_idx and any other bits
        """
        line = self.filenames[index].split()
        folder = line[0]

        if len(line) == 3:
            frame_index = int(line[1])
        else:
            frame_index = 0

        if len(line) == 3:
            side = line[2]
        else:
            side = None

        return folder, frame_index, side
    
    def get_color(self, folder, frame_index, side, do_flip):
        color = self.loader(self.get_image_path(folder, frame_index, side))

        if do_flip:
            color = color.transpose(pil.FLIP_LEFT_RIGHT)

        return color


class KITTIRAWDataset(KITTIDataset):
    """KITTI dataset which loads the original velodyne depth maps for ground truth
    """
    def __init__(self, *args, **kwargs):
        super(KITTIRAWDataset, self).__init__(*args, **kwargs)

    def get_image_path(self, folder, frame_index, side):
        f_str = "{:010d}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(
            self.data_path, folder, "image_0{}/data".format(self.side_map[side]), f_str)
        return image_path

    def get_depth(self, folder, frame_index, side, do_flip):
        calib_path = os.path.join(self.data_path, folder.split("/")[0])

        velo_filename = os.path.join(
            self.data_path,
            folder,
            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))

        depth_gt = generate_depth_map(calib_path, velo_filename, self.side_map[side])
        depth_gt = skimage.transform.resize(
            depth_gt, self.full_res_shape[::-1], order=0, preserve_range=True, mode='constant')

        if do_flip:
            depth_gt = np.fliplr(depth_gt)

        return depth_gt


class KITTIOdomDataset(KITTIDataset):
    """KITTI dataset for odometry training and testing
    """
    def __init__(self, *args, **kwargs):
        super(KITTIOdomDataset, self).__init__(*args, **kwargs)

    def get_image_path(self, folder, frame_index, side):
        f_str = "{:06d}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(
            self.data_path,
            "sequences/{:02d}".format(int(folder)),
            "image_{}".format(self.side_map[side]),
            f_str)
        return image_path


class KITTIDepthDataset(KITTIDataset):
    """KITTI dataset which uses the updated ground truth depth maps
    """
    def __init__(self, *args, **kwargs):
        super(KITTIDepthDataset, self).__init__(*args, **kwargs)

    def get_image_path(self, folder, frame_index, side):
        f_str = "{:010d}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(
            self.data_path,
            folder,
            "image_0{}/data".format(self.side_map[side]),
            f_str)
        return image_path

    def get_depth(self, folder, frame_index, side, do_flip):
        f_str = "{:010d}.png".format(frame_index)
        depth_path = os.path.join(
            self.data_path,
            folder,
            "proj_depth/groundtruth/image_0{}".format(self.side_map[side]),
            f_str)

        depth_gt = pil.open(depth_path)
        depth_gt = depth_gt.resize(self.full_res_shape, pil.NEAREST)
        depth_gt = np.array(depth_gt).astype(np.float32) / 256

        if do_flip:
            depth_gt = np.fliplr(depth_gt)

        return depth_gt

datasets/kitti_dataset.py
Sauf4896/BDEdepth
from __future__ import absolute_import, division, print_function

import os
import numpy as np
from PIL import Image, ImageFile
from scipy.io import loadmat

import torch
import torch.utils.data as data
from torchvision import transforms

ImageFile.LOAD_TRUNCATED_IMAGES = True

# refer to https://github.com/nianticlabs/monodepth2/issues/392

class Make3DDataset(data.Dataset):
    def __init__(self, data_path, filenames, input_resolution=None):
        super().__init__()

        self.data_path = data_path
        self.filenames = filenames
        self.input_resolution = input_resolution #  H*W
        self.path_format_dict = {
                        'color': ('Test134', 'img-', 'jpg'),
                        'depth': ('Gridlaserdata', 'depth_sph_corr-', 'mat')
                    }
        
        if self.input_resolution is not None:
            ## for the interpolation mode
            ## Image.LANCZOS can replicate the results in monodepth2 paper
            ## Image.NEAREST can be better than LANCZOS
            self.resize = transforms.Resize(self.input_resolution,
                                               interpolation=Image.LANCZOS)
        self.to_tensor = transforms.ToTensor()

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):

        line = self.filenames[index]
        base_path = os.path.join(self.data_path, '{}',
                                 '{}' + line + '.{}')
        
        color_path = base_path.format(*self.path_format_dict['color'])
        raw_img = Image.open(color_path).convert('RGB')

        depth_path = base_path.format(*self.path_format_dict['depth'])
        raw_depth = loadmat(depth_path)['Position3DGrid'][:, :, 3]
  
        inputs = {}

        img = raw_img.crop((0, 710, 1704, 1562))
        if self.input_resolution is not None:
            img = self.resize(img)
        img = self.to_tensor(img)
        inputs['color'] = img

        depth = raw_depth[17:38, :]
        depth = torch.from_numpy(depth)
        inputs['depth']  = depth

        return inputs


datasets/make3d_dataset.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import random
import numpy as np
import copy
from PIL import Image  # using pillow-simd for increased speed

import torch
import torch.utils.data as data
from torchvision import transforms


def pil_loader(path):
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')

class MonoDataset(data.Dataset):
    """Superclass for monocular dataloaders

    Args:
        data_path
        filenames
        height
        width
        frame_idxs
        num_scales
        is_train
        img_ext
    """
    def __init__(self,
                 data_path,
                 filenames,
                 height,
                 width,
                 frame_idxs,
                 num_scales,
                 resize_ratio=[1.2, 2.0],
                 local_crop=False,
                 split_ratio=[0.1, 0.9],
                 patch_reshuffle=False,
                 is_train=False,
                 img_ext='.png'):
        super(MonoDataset, self).__init__()

        self.data_path = data_path
        self.filenames = filenames
        self.height = height
        self.width = width
        self.num_scales = num_scales
        self.interp = Image.ANTIALIAS

        self.frame_idxs = frame_idxs

        self.is_train = is_train
        self.img_ext = img_ext

        self.resize_ratio_lower = resize_ratio[0]
        self.resize_ratio_upper = resize_ratio[1]
        self.local_crop = local_crop

        self.split_ratio_lower = split_ratio[0]
        self.split_ratio_upper = split_ratio[1]
        self.patch_reshuffle = patch_reshuffle

        self.loader = pil_loader
        self.to_tensor = transforms.ToTensor()

        # We need to specify augmentations differently in newer versions of torchvision.
        # We first try the newer tuple version; if this fails we fall back to scalars
        try:
            self.brightness = (0.8, 1.2)
            self.contrast = (0.8, 1.2)
            self.saturation = (0.8, 1.2)
            self.hue = (-0.1, 0.1)
            transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)
        except TypeError:
            self.brightness = 0.2
            self.contrast = 0.2
            self.saturation = 0.2
            self.hue = 0.1

        self.resize = {}
        for i in range(self.num_scales):
            s = 2 ** i
            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
                                               interpolation=self.interp)

        # self.load_depth = self.check_depth()
        self.load_depth = False

    def preprocess(self, inputs, color_aug):
        """Resize colour images to the required scales and augment if required

        We create the color_aug object in advance and apply the same augmentation to all
        images in this item. This ensures that all images input to the pose network receive the
        same augmentation.
        """

        if self.local_crop:
            resize_ratio = (self.resize_ratio_upper - self.resize_ratio_lower) * random.random() + self.resize_ratio_lower
            height_re = int(self.height * resize_ratio)
            width_re = int(self.width * resize_ratio)        
            w0 = int((width_re - self.width) * random.random())
            h0 = int((height_re - self.height) * random.random())
            self.resize_local = transforms.Resize((height_re, width_re), interpolation=self.interp)
            box = (w0, h0, w0+self.width, h0+self.height)
            gridx, gridy = np.meshgrid(np.linspace(-1, 1, width_re), np.linspace(-1, 1, height_re))
            gridx = torch.from_numpy(gridx)
            gridy = torch.from_numpy(gridy)
            grid = torch.stack([gridx, gridy], dim=0)
            inputs[("grid_local")] = grid[:, h0 : h0+self.height, w0 : w0+self.width].clone()  
            inputs[("ratio_local")] = torch.tensor([resize_ratio])


        for k in list(inputs):
            frame = inputs[k]
            if "color" in k:
                n, im, i = k
                for i in range(self.num_scales):
                    inputs[(n, im, i)] = self.resize[i](inputs[(n, im, i - 1)])
                    if self.local_crop:
                        if i == 0:
                            inputs[(n + "_local", im, i)] = self.resize_local(inputs[(n, im, -1)]).crop(box)
                        else:
                            inputs[(n + "_local", im, i)] = self.resize[i](inputs[(n + "_local", im, i - 1)])
                if self.patch_reshuffle:
                    if im == 0:
                        ### Split-Permute as depicted in paper (vertical + horizontal)
                        img = inputs[(n, im, 0)]
                        newimg = img.copy()
                        ratio_x = random.random() * (self.split_ratio_upper - self.split_ratio_lower) + self.split_ratio_lower
                        ratio_y = random.random() * (self.split_ratio_upper - self.split_ratio_lower) + self.split_ratio_lower
                        w_i = int(self.width * ratio_x)
                        patch1 = img.crop((0, 0, w_i, self.height)).copy()
                        patch2 = img.crop((w_i, 0, self.width, self.height)).copy()
                        newimg.paste(patch2, (0, 0))
                        newimg.paste(patch1, (self.width-w_i, 0))     
                        h_i = int(self.height * ratio_y)
                        patch1 = newimg.crop((0, 0, self.width, h_i)).copy()
                        patch2 = newimg.crop((0, h_i, self.width, self.height)).copy()
                        newimg.paste(patch2, (0, 0))
                        newimg.paste(patch1, (0, self.height-h_i))      
                        inputs[(n + "_reshuffle", im, 0)] = newimg
                        inputs[("split_xy")] = torch.tensor([self.width-w_i, self.height-h_i])         

                        ### Split-Permute (vertical or horizontal, randomly choose one)
                        # img = inputs[(n, im, 0)]
                        # newimg = img.copy()
                        # ratio = random.random() * (self.split_ratio_upper - self.split_ratio_lower) + self.split_ratio_lower
                        # if random.random() > 0.5:
                        #     w_i = int(self.width * ratio)
                        #     patch1 = img.crop((0, 0, w_i, self.height)).copy()
                        #     patch2 = img.crop((w_i, 0, self.width, self.height)).copy()
                        #     newimg.paste(patch2, (0, 0))
                        #     newimg.paste(patch1, (self.width-w_i, 0))
                        #     inputs[(n + "_reshuffle", im, 0)] = newimg
                        #     inputs[("split_xy")] = torch.tensor([self.width-w_i, 0])
                        # else:
                        #     h_i = int(self.height * ratio)
                        #     patch1 = img.crop((0, 0, self.width, h_i)).copy()
                        #     patch2 = img.crop((0, h_i, self.width, self.height)).copy()
                        #     newimg.paste(patch2, (0, 0))
                        #     newimg.paste(patch1, (0, self.height-h_i))                           
                        #     inputs[(n + "_reshuffle", im, 0)] = newimg
                        #     inputs[("split_xy")] = torch.tensor([0, self.height-h_i])          

        for k in list(inputs):
            f = inputs[k]
            if "color" in k[0]:
                n, im, i = k
                inputs[(n, im, i)] = self.to_tensor(f)
                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f))

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        """Returns a single training item from the dataset as a dictionary.

        Values correspond to torch tensors.
        Keys in the dictionary are either strings or tuples:

            ("color", <frame_id>, <scale>)          for raw colour images,
            ("color_aug", <frame_id>, <scale>)      for augmented colour images,
            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
            "stereo_T"                              for camera extrinsics, and
            "depth_gt"                              for ground truth depth maps.

        <frame_id> is either:
            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
        or
            "s" for the opposite image in the stereo pair.

        <scale> is an integer representing the scale of the image relative to the fullsize image:
            -1      images at native resolution as loaded from disk
            0       images resized to (self.width,      self.height     )
            1       images resized to (self.width // 2, self.height // 2)
            2       images resized to (self.width // 4, self.height // 4)
            3       images resized to (self.width // 8, self.height // 8)
        """
        inputs = {}

        do_color_aug = self.is_train and random.random() > 0.5
        do_flip = self.is_train and random.random() > 0.5

        folder, frame_index, side = self.index_to_folder_and_frame_idx(index)

        if type(self).__name__ in ["CityscapesDataset"]:
            inputs.update(self.get_color(folder, frame_index, side, do_flip))       
            self.K = self.load_intrinsics(folder, frame_index)
        else:
            for i in self.frame_idxs:
                if i == "s":
                    other_side = {"r": "l", "l": "r"}[side]
                    inputs[("color", i, -1)] = self.get_color(folder, frame_index, other_side, do_flip)
                else:
                    inputs[("color", i, -1)] = self.get_color(folder, frame_index + i, side, do_flip)

        # adjusting intrinsics to match each scale in the pyramid
        for scale in range(self.num_scales):
            K = self.K.copy()

            K[0, :] *= self.width // (2 ** scale)
            K[1, :] *= self.height // (2 ** scale)

            inv_K = np.linalg.pinv(K)

            inputs[("K", scale)] = torch.from_numpy(K)
            inputs[("inv_K", scale)] = torch.from_numpy(inv_K)

        if do_color_aug:
            color_aug = transforms.ColorJitter(
                self.brightness, self.contrast, self.saturation, self.hue)
        else:
            color_aug = (lambda x: x)

        self.preprocess(inputs, color_aug)

        for i in self.frame_idxs:
            del inputs[("color", i, -1)]
            del inputs[("color_aug", i, -1)]

        if self.load_depth:
            depth_gt = self.get_depth(folder, frame_index, side, do_flip)
            inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
            inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))

        if "s" in self.frame_idxs:
            stereo_T = np.eye(4, dtype=np.float32)
            baseline_sign = -1 if do_flip else 1
            side_sign = -1 if side == "l" else 1
            stereo_T[0, 3] = side_sign * baseline_sign * 0.1

            inputs["stereo_T"] = torch.from_numpy(stereo_T)

        return inputs

    def get_color(self, folder, frame_index, side, do_flip):
        raise NotImplementedError

    def check_depth(self):
        raise NotImplementedError

    def get_depth(self, folder, frame_index, side, do_flip):
        raise NotImplementedError

datasets/mono_dataset.py
Sauf4896/BDEdepth
from __future__ import absolute_import, division, print_function

import os
import cv2
import random
import numpy as np
import h5py
from PIL import Image
import torch
from torchvision import transforms
from .mono_dataset import MonoDataset

CROP = 16

def h5_loader(path):
    h5f = h5py.File(path, "r")
    rgb = np.array(h5f['rgb'])
    rgb = np.transpose(rgb, (1, 2, 0))
    depth = np.array(h5f['depth'])
    # norm = np.array(h5f['norm'])
    # norm = np.transpose(norm, (1,2,0))
    # valid_mask = np.array(h5f['mask'])

    return rgb, depth

class NYUDataset(MonoDataset):

    def __init__(self, *args, **kwargs):
        super(NYUDataset, self).__init__(*args, **kwargs)

        if self.is_train:
            self.full_res_shape = (480-CROP*2, 640-CROP*2) 
        else:
            self.full_res_shape = (427, 561) 
            self.loader = h5_loader

        self.K = self._get_intrinsics()

    def __getitem__(self, index):
        # for testing
        if not self.is_train:
            line = self.filenames[index]
            line = os.path.join(self.data_path, line)
            rgb, depth = self.loader(line) 
            h, w, c = rgb.shape

            rgb = rgb[44: 471, 40: 601, :]
            depth = depth[44: 471, 40: 601]
            rgb = Image.fromarray(rgb)
            rgb = self.to_tensor(self.resize[0](rgb))
            depth = torch.tensor(depth)

            # K = self.K.copy()
            # K[0, :] *= self.width
            # K[1, :] *= self.height
            return rgb, depth

        inputs = {}

        do_color_aug = random.random() > 0.5
        do_flip = random.random() > 0.5

        line = self.filenames[index].split()
        frame_name = line[0]

        # if do_flip:
        #     if frame_name in ['nyu_depth_v2/bedroom_0086/r-1315251252.453265-3131793654.ppm', 'nyu_depth_v2/living_room_0011/r-1295837629.903855-1295712361.ppm']:
        #         do_flip = False
        #         print("-----------")
        #         print("meet not flip vps frame {}\n do_flip is False".format(frame_name))
        #         print("-----------")

        line = [os.path.join(self.data_path, l) for l in line]
        
        for ind, i in enumerate([0, -4, -3, -2, -1, 1, 2, 3, 4]):
            if not i in set(self.frame_idxs):
                continue
            inputs[("color", i, -1)] = self.get_color(line[ind], do_flip)


        # adjusting intrinsics to match each scale in the pyramid
        for scale in range(self.num_scales):
            K = self.K.copy()

            K[0, :] *= self.width // (2 ** scale)
            K[1, :] *= self.height // (2 ** scale)
            
            inv_K = np.linalg.pinv(K)

            inputs[("K", scale)] = torch.from_numpy(K)
            inputs[("inv_K", scale)] = torch.from_numpy(inv_K)

        if do_color_aug:
            color_aug = transforms.ColorJitter(
                self.brightness, self.contrast, self.saturation, self.hue)
        else:
            color_aug = (lambda x: x)

        self.preprocess(inputs, color_aug)

        for i in [0, -4, -3, -2, -1, 1, 2, 3, 4]:
            if not i in set(self.frame_idxs):
                continue

            del inputs[("color", i, -1)]
            del inputs[("color_aug", i, -1)]
    
        return inputs

    def get_color(self, fp, do_flip):
        color = self.loader(fp)
        color = self._undistort(np.array(color))

        if do_flip:
            color = cv2.flip(color, 1)

        h, w, c = color.shape
        color = color[CROP: h-CROP, CROP: w-CROP, :]

        return Image.fromarray(color)


    def _get_intrinsics(self):
        h, w = self.full_res_shape
        fx = 5.1885790117450188e+02 / w
        fy = 5.1946961112127485e+02 / h
        if self.is_train:
            cx = (3.2558244941119034e+02 - CROP) / w
            cy = (2.5373616633400465e+02 - CROP) / h
        else:
            cx = (3.2558244941119034e+02 - 40) / w
            cy = (2.5373616633400465e+02 - 44) / h
        intrinsics =np.array([[fx, 0., cx, 0.], 
                               [0., fy, cy, 0.],
                               [0., 0., 1., 0.],
                               [0., 0., 0., 1.]], dtype="float32")
        return intrinsics

    def _undistort(self, image):
        k1 =  2.0796615318809061e-01
        k2 = -5.8613825163911781e-01
        p1 = 7.2231363135888329e-04
        p2 = 1.0479627195765181e-03
        k3 = 4.9856986684705107e-01

        fx = 5.1885790117450188e+02
        fy = 5.1946961112127485e+02
        cx = 3.2558244941119034e+02
        cy = 2.5373616633400465e+02

        kmat = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
        dist = np.array([[k1, k2, p1, p2, k3]])
        image = cv2.undistort(image, kmat, dist)
        return image
datasets/nyuv2_dataset.py
Sauf4896/BDEdepth
import argparse
import torch
from torch.utils.data import DataLoader
import torch.nn.functional as F
import datasets
from networks import *
from utils import *
from layers import disp_to_depth, compute_depth_errors
import warnings

warnings.filterwarnings("ignore")

STEREO_SCALE_FACTOR = 5.4

def eval_args():
    parser = argparse.ArgumentParser(description='Evaluation Parser')

    parser.add_argument('--pretrained_path',
                        type=str,
                        help="path of model checkpoint to load")
    parser.add_argument("--backbone",
                        type=str,
                        help="backbone of depth encoder",
                        default="hrnet",
                        choices=["resnet", "hrnet"])
    parser.add_argument("--num_layers",
                        type=int,
                        help="number of resnet layers",
                        default=18,
                        choices=[18, 34, 50, 101, 152])
    parser.add_argument("--batch_size",
                        type=int,
                        help="batch size",
                        default=4)
    parser.add_argument("--height",
                        type=int,
                        help="input image height",
                        default=192)
    parser.add_argument("--width",
                        type=int,
                        help="input image width",
                        default=640)
    parser.add_argument("--num_workers",
                        type=int,
                        help="number of dataloader workers",
                        default=12)
    parser.add_argument("--min_depth",
                        type=float,
                        help="minimum depth",
                        default=0.1)
    parser.add_argument("--max_depth",
                        type=float,
                        help="maximum depth",
                        default=100.0)
    parser.add_argument("--post_process",
                        help="if set will perform the flipping post processing "
                            "from the original monodepth paper",
                        action="store_true")
    parser.add_argument("--use_stereo",
                        help="if set, uses stereo pair for training",
                        action="store_true")
    ## paths of test datasets
    parser.add_argument('--kitti_path',
                        type=str,
                        help="data path of KITTI, do not set if you do not want to evaluate on this dataset")
    parser.add_argument('--make3d_path',
                        type=str,
                        help="data path of Make3D, do not set if you do not want to evaluate on this dataset")
    parser.add_argument('--nyuv2_path',
                        type=str,
                        help="data path of NYU v2, do not set if you do not want to evaluate on this dataset")
    parser.add_argument('--cityscapes_path',
                        type=str,
                        help="data path of Cityscapes, do not set if you do not want to evaluate on this dataset")
    args = parser.parse_args()
    return args

def compute_errors(gt, pred):
    rmse = (gt - pred) ** 2
    rmse = torch.sqrt(rmse.mean())

    rmse_log = (torch.log10(gt) - torch.log10(pred)) ** 2
    rmse_log = torch.sqrt(rmse_log.mean())

    abs_rel = torch.mean(torch.abs(gt - pred) / gt)

    sq_rel = torch.mean(((gt - pred)**2) / gt)

    return abs_rel, sq_rel, rmse, rmse_log


def batch_post_process_disparity(l_disp, r_disp):
    """Apply the disparity post-processing method as introduced in Monodepthv1
    """
    _, h, w = l_disp.shape
    m_disp = 0.5 * (l_disp + r_disp)
    l, _ = np.meshgrid(np.linspace(0, 1, w), np.linspace(0, 1, h))
    l_mask = (1.0 - np.clip(20 * (l - 0.05), 0, 1))[None, ...]
    r_mask = l_mask[:, :, ::-1]
    l_mask = torch.tensor(l_mask).cuda()
    r_mask = torch.tensor(r_mask.copy()).cuda()
    return r_mask * l_disp + l_mask * r_disp + (1.0 - l_mask - r_mask) * m_disp


def load_model(args):
    print("-> Loading weights from {}".format(args.pretrained_path))

    model = torch.load(args.pretrained_path, map_location='cpu')
    if args.backbone == "hrnet":
        depth_encoder = depth_hrnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_hrnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))
    if args.backbone == "resnet":
        depth_encoder = depth_resnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_resnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))

    depth_encoder.load_state_dict({k: v for k, v in model["encoder"].items() if k in depth_encoder.state_dict()})
    depth_decoder.load_state_dict({k: v for k, v in model["depth"].items() if k in depth_decoder.state_dict()})
    
    depth_encoder.cuda().eval()
    depth_decoder.cuda().eval()
    
    input_color = torch.ones(1, 3, args.height, args.width).cuda()
    flops, params, flops_e, params_e, flops_d, params_d = profile_once(depth_encoder, depth_decoder, input_color)
    print("\n  " + ("flops: {0}, params: {1}, flops_e: {2}, params_e:{3}, flops_d:{4}, params_d:{5}").format(flops, params, flops_e, params_e, flops_d, params_d) + "\n")
  
    return depth_encoder, depth_decoder

def test_kitti(args, dataloader, depth_encoder, depth_decoder, eval_split='eigen'):
    MIN_DEPTH = 1e-3
    MAX_DEPTH = 80

    gt_path = os.path.join(os.path.dirname(__file__), "splits", "kitti", eval_split, "gt_depths.npz")
    gt_depths = np.load(gt_path, fix_imports=True, encoding='latin1', allow_pickle=True)["data"]

    pred_disps = []
    for data in dataloader:
        input_color = data[("color", 0, 0)].cuda()
        if args.post_process:
            input_color = torch.cat((input_color, torch.flip(input_color, [3])), 0)
        output = depth_decoder(depth_encoder(input_color))
        pred_disp, _ = disp_to_depth(output[("disp", 0)], args.min_depth, args.max_depth)
        pred_disp = pred_disp[:, 0]
        if args.post_process:
            N = pred_disp.shape[0] // 2
            pred_disp = batch_post_process_disparity(pred_disp[:N], torch.flip(pred_disp[N:], [2]))
        pred_disps.append(pred_disp)
    pred_disps = torch.cat(pred_disps, dim=0)

    errors = []
    ratios = []
    for i in range(pred_disps.shape[0]):
        gt_depth = torch.from_numpy(gt_depths[i]).cuda()
        gt_height, gt_width = gt_depth.shape[:2]
        pred_disp = pred_disps[i:i+1].unsqueeze(0)
        pred_disp = F.interpolate(pred_disp, (gt_height, gt_width), mode="bilinear", align_corners=True)
        pred_depth = 1 / pred_disp[0, 0, :]
        if eval_split == "eigen":
            mask = (gt_depth > MIN_DEPTH) & (gt_depth < MAX_DEPTH)
            crop_mask = torch.zeros_like(mask)
            crop_mask[
                    int(0.40810811 * gt_height):int(0.99189189 * gt_height),
                    int(0.03594771 * gt_width):int(0.96405229 * gt_width)] = 1
            mask = mask * crop_mask
        else:
            mask = gt_depth > 0

        pred_depth = pred_depth[mask]
        gt_depth = gt_depth[mask]
        if args.use_stereo:
            pred_depth *= STEREO_SCALE_FACTOR
        else:
            ratio = torch.median(gt_depth) / torch.median(pred_depth)
            ratios.append(ratio)
            pred_depth *= ratio  
        pred_depth = torch.clamp(pred_depth, MIN_DEPTH, MAX_DEPTH)
        errors.append(compute_depth_errors(gt_depth, pred_depth))

    if not args.use_stereo:
        ratios = torch.tensor(ratios)
        med = torch.median(ratios)
        std = torch.std(ratios / med)
        print(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

    mean_errors = torch.tensor(errors).mean(0)

    print(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
    print(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))

def test_make3d(args, dataloader, depth_encoder, depth_decoder):
    pred_depths = []
    gt_depths = []
    for data in dataloader:
        input_color = data["color"].cuda()
        if args.post_process:
            input_color = torch.cat((input_color, torch.flip(input_color, [3])), 0)

        output = depth_decoder(depth_encoder(input_color))
        pred_disp, _ = disp_to_depth(output[("disp", 0)], args.min_depth, args.max_depth)
        pred_disp = pred_disp[:, 0]

        if args.post_process:
            N = pred_disp.shape[0] // 2
            pred_disp = batch_post_process_disparity(pred_disp[:N], torch.flip(pred_disp[N:], [2]))
        
        gt_depth = data["depth"]
        _, h, w = gt_depth.shape
        pred_depth = 1 / pred_disp
        pred_depth = F.interpolate(pred_depth.unsqueeze(0), (h, w), mode="nearest")[0]
        pred_depths.append(pred_depth)
        gt_depths.append(gt_depth)
    pred_depths = torch.cat(pred_depths, dim=0)
    gt_depths = torch.cat(gt_depths, dim=0).cuda()

    errors = []
    ratios = []
    for i in range(pred_depths.shape[0]):    
        pred_depth = pred_depths[i]
        gt_depth = gt_depths[i]
        mask = (gt_depth > 0) & (gt_depth < 70)
        pred_depth = pred_depth[mask]
        gt_depth = gt_depth[mask]
        if args.use_stereo:
            pred_depth *= STEREO_SCALE_FACTOR
        else:
            ratio = torch.median(gt_depth) / torch.median(pred_depth)
            ratios.append(ratio)
            pred_depth *= ratio         
        pred_depth[pred_depth > 70] = 70
        errors.append(compute_errors(gt_depth, pred_depth))

    if not args.use_stereo:
        ratios = torch.tensor(ratios)
        med = torch.median(ratios)
        std = torch.std(ratios / med)
        print(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

    mean_errors = torch.tensor(errors).mean(0)

    print(("{:>8} | " * 4).format( "abs_rel", "sq_rel", "rmse", "rmse_log"))
    print(("{: 8.3f} | " * 4 + "\n").format(*mean_errors.tolist()))

def test_nyuv2(args, dataloader, depth_encoder, depth_decoder):
    pred_depths = []
    gt_depths = []
    for (color, depth) in dataloader:
        input_color = color.cuda()
        if args.post_process:
            input_color = torch.cat((input_color, torch.flip(input_color, [3])), 0)

        output = depth_decoder(depth_encoder(input_color))
        pred_disp, _ = disp_to_depth(output[("disp", 0)], args.min_depth, args.max_depth)
        pred_disp = pred_disp[:, 0]

        if args.post_process:
            N = pred_disp.shape[0] // 2
            pred_disp = batch_post_process_disparity(pred_disp[:N], torch.flip(pred_disp[N:], [2]))
        
        gt_depth = depth
        _, h, w = gt_depth.shape
        pred_depth = 1 / pred_disp
        pred_depth = F.interpolate(pred_depth.unsqueeze(0), (h, w), mode="nearest")[0]
        pred_depths.append(pred_depth)
        gt_depths.append(gt_depth)
    pred_depths = torch.cat(pred_depths, dim=0)
    gt_depths = torch.cat(gt_depths, dim=0).cuda()

    errors = []
    ratios = []
    for i in range(pred_depths.shape[0]):    
        pred_depth = pred_depths[i]
        gt_depth = gt_depths[i]
        mask = (gt_depth > 0) & (gt_depth < 10)
        pred_depth = pred_depth[mask]
        gt_depth = gt_depth[mask]
        ratio = torch.median(gt_depth) / torch.median(pred_depth)
        ratios.append(ratio)
        pred_depth *= ratio         
        pred_depth[pred_depth > 10] = 10
        errors.append(compute_depth_errors(gt_depth, pred_depth))

    ratios = torch.tensor(ratios)
    med = torch.median(ratios)
    std = torch.std(ratios / med)
    print(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

    mean_errors = torch.tensor(errors).mean(0)

    print(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
    print(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))


def test_cityscapes(args, dataloader, depth_encoder, depth_decoder):
    MIN_DEPTH = 1e-3
    MAX_DEPTH = 80

    gt_path = os.path.join(os.path.dirname(__file__), "splits", "cityscapes", "gt_depths")

    pred_disps = []
    for data in dataloader:
        input_color = data[("color", 0, 0)].cuda()
        if args.post_process:
            input_color = torch.cat((input_color, torch.flip(input_color, [3])), 0)
        output = depth_decoder(depth_encoder(input_color))
        pred_disp, _ = disp_to_depth(output[("disp", 0)], args.min_depth, args.max_depth)
        pred_disp = pred_disp[:, 0]
        if args.post_process:
            N = pred_disp.shape[0] // 2
            pred_disp = batch_post_process_disparity(pred_disp[:N], torch.flip(pred_disp[N:], [2]))
        pred_disps.append(pred_disp)
    pred_disps = torch.cat(pred_disps, dim=0)

    errors = []
    ratios = []
    for i in range(pred_disps.shape[0]):
        gt_depth = np.load(os.path.join(gt_path, str(i).zfill(3) + '_depth.npy'))
        gt_height, gt_width = gt_depth.shape[:2]
        # crop ground truth to remove ego car -> this has happened in the dataloader for inputs
        gt_height = int(round(gt_height * 0.75))
        gt_depth = torch.from_numpy(gt_depth[:gt_height]).cuda()
        pred_disp = pred_disps[i:i+1].unsqueeze(0)
        pred_disp = F.interpolate(pred_disp, (gt_height, gt_width), mode="bilinear", align_corners=True)
        pred_depth = 1 / pred_disp[0, 0, :]

        # when evaluating cityscapes, we centre crop to the middle 50% of the image.
        # Bottom 25% has already been removed - so crop the sides and the top here
        gt_depth = gt_depth[256:, 192:1856]
        pred_depth = pred_depth[256:, 192:1856]

        mask = (gt_depth > MIN_DEPTH) & (gt_depth < MAX_DEPTH)
        pred_depth = pred_depth[mask]
        gt_depth = gt_depth[mask]
       
        if args.use_stereo:
            pred_depth *= STEREO_SCALE_FACTOR
        else:
            ratio = torch.median(gt_depth) / torch.median(pred_depth)
            ratios.append(ratio)
            pred_depth *= ratio  
        pred_depth = torch.clamp(pred_depth, MIN_DEPTH, MAX_DEPTH)
        errors.append(compute_depth_errors(gt_depth, pred_depth))

    if not args.use_stereo:
        ratios = torch.tensor(ratios)
        med = torch.median(ratios)
        std = torch.std(ratios / med)
        print(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

    mean_errors = torch.tensor(errors).mean(0)

    print(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
    print(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))  

def main(args):

    depth_encoder, depth_decoder = load_model(args)
    input_resolution = (args.height, args.width)
    
    print(" Evaluated at resolution {} * {}".format(input_resolution[0], input_resolution[1]))
    if args.post_process:
        print(" Post-process is used")
    else:
        print(" No post-process")
    if args.use_stereo:
        print(" Stereo evaluation - disabling median scaling")
        print(" Scaling by {} \n".format(STEREO_SCALE_FACTOR))
    else:
        print(" Mono evaluation - using median scaling \n")

    splits_dir = os.path.join(os.path.dirname(__file__), "splits")

    if args.kitti_path:
        ## evaluate on eigen split
        print(" Evaluate on KITTI with eigen split:")
        filenames = readlines(os.path.join(splits_dir, "kitti", "eigen", "test_files.txt")) 
        dataset = datasets.KITTIRAWDataset(args.kitti_path, filenames, input_resolution[0], input_resolution[1], [0], 1, is_train=False, img_ext='.png')
        dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, drop_last=False)
        with torch.no_grad():
            test_kitti(args, dataloader, depth_encoder, depth_decoder, "eigen")

        ## evaluate on eigen_benchmark split
        print(" Evaluate on KITTI with eigen_benchmark split (improved groundtruth):")
        filenames = readlines(os.path.join(splits_dir, "kitti", "eigen_benchmark", "test_files.txt")) 
        dataset = datasets.KITTIRAWDataset(args.kitti_path, filenames, input_resolution[0], input_resolution[1], [0], 1, is_train=False, img_ext='.png')
        dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, drop_last=False)
        with torch.no_grad():
            test_kitti(args, dataloader, depth_encoder, depth_decoder, "eigen_benchmark")

    if args.make3d_path:
        print(" Evaluate on Make3D:")
        filenames = readlines(os.path.join(splits_dir, "make3d", "test_files.txt"))
        dataset = datasets.Make3DDataset(args.make3d_path, filenames, input_resolution)
        dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, drop_last=False)
        with torch.no_grad():
            test_make3d(args, dataloader, depth_encoder, depth_decoder)

    if args.nyuv2_path:
        print(" Evaluate on NYU Depth v2:")
        filenames = readlines(os.path.join(splits_dir, "nyuv2", "test_files.txt"))
        dataset = datasets.NYUDataset(args.nyuv2_path, filenames, input_resolution[0], input_resolution[1], [0], 1, is_train=False)
        dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, drop_last=False)
        with torch.no_grad():
            test_nyuv2(args, dataloader, depth_encoder, depth_decoder)

    if args.cityscapes_path:
        print(" Evaluate on Cisyscapes:")
        filenames = readlines(os.path.join(splits_dir, "cityscapes", "test_files.txt"))
        dataset = datasets.CityscapesDataset(args.cityscapes_path, filenames, input_resolution[0], input_resolution[1], [0], 1, is_train=False)
        dataloader = DataLoader(dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=False, drop_last=False)
        with torch.no_grad():
            test_cityscapes(args, dataloader, depth_encoder, depth_decoder)  

if __name__ == '__main__':
    args = eval_args()
    main(args)

evaluate_depth.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os

import argparse
import numpy as np
import PIL.Image as pil

from utils import readlines
from kitti_utils import generate_depth_map


def export_gt_depths_kitti():

    parser = argparse.ArgumentParser(description='export_gt_depth')

    parser.add_argument('--data_path',
                        type=str,
                        help='path to the root of the KITTI data',
                        required=True)
    parser.add_argument('--split',
                        type=str,
                        help='which split to export gt from',
                        required=True,
                        choices=["eigen", "eigen_benchmark"])
    opt = parser.parse_args()

    split_folder = os.path.join(os.path.dirname(__file__), "splits/kitti", opt.split)
    lines = readlines(os.path.join(split_folder, "test_files.txt"))

    print("Exporting ground truth depths for {}".format(opt.split))

    gt_depths = []
    for line in lines:

        folder, frame_id, _ = line.split()
        frame_id = int(frame_id)

        if opt.split == "eigen":
            calib_dir = os.path.join(opt.data_path, folder.split("/")[0])
            velo_filename = os.path.join(opt.data_path, folder,
                                         "velodyne_points/data", "{:010d}.bin".format(frame_id))
            gt_depth = generate_depth_map(calib_dir, velo_filename, 2, True)
        elif opt.split == "eigen_benchmark":
            gt_depth_path = os.path.join(opt.data_path, folder, "proj_depth",
                                         "groundtruth", "image_02", "{:010d}.png".format(frame_id))
            gt_depth = np.array(pil.open(gt_depth_path)).astype(np.float32) / 256

        gt_depths.append(gt_depth.astype(np.float32))

    output_path = os.path.join(split_folder, "gt_depths.npz")

    print("Saving to {}".format(output_path))

    np.savez_compressed(output_path, data=np.array(gt_depths))


if __name__ == "__main__":
    export_gt_depths_kitti()
export_gt_depth.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F


def disp_to_depth(disp, min_depth, max_depth):
    """Convert network's sigmoid output into depth prediction
    The formula for this conversion is given in the 'additional considerations'
    section of the paper.
    """
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = min_disp + (max_disp - min_disp) * disp
    depth = 1 / scaled_disp
    return scaled_disp, depth


def transformation_from_parameters(axisangle, translation, invert=False):
    """Convert the network's (axisangle, translation) output into a 4x4 matrix
    """
    R = rot_from_axisangle(axisangle)
    t = translation.clone()

    if invert:
        R = R.transpose(1, 2)
        t *= -1

    T = get_translation_matrix(t)

    if invert:
        M = torch.matmul(R, T)
    else:
        M = torch.matmul(T, R)

    return M


def get_translation_matrix(translation_vector):
    """Convert a translation vector into a 4x4 transformation matrix
    """
    T = torch.zeros(translation_vector.shape[0], 4, 4).to(device=translation_vector.device)

    t = translation_vector.contiguous().view(-1, 3, 1)

    T[:, 0, 0] = 1
    T[:, 1, 1] = 1
    T[:, 2, 2] = 1
    T[:, 3, 3] = 1
    T[:, :3, 3, None] = t

    return T


def rot_from_axisangle(vec):
    """Convert an axisangle rotation into a 4x4 transformation matrix
    (adapted from https://github.com/Wallacoloo/printipi)
    Input 'vec' has to be Bx1x3
    """
    angle = torch.norm(vec, 2, 2, True)
    axis = vec / (angle + 1e-7)

    ca = torch.cos(angle)
    sa = torch.sin(angle)
    C = 1 - ca

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    xs = x * sa
    ys = y * sa
    zs = z * sa
    xC = x * C
    yC = y * C
    zC = z * C
    xyC = x * yC
    yzC = y * zC
    zxC = z * xC

    rot = torch.zeros((vec.shape[0], 4, 4)).to(device=vec.device)

    rot[:, 0, 0] = torch.squeeze(x * xC + ca)
    rot[:, 0, 1] = torch.squeeze(xyC - zs)
    rot[:, 0, 2] = torch.squeeze(zxC + ys)
    rot[:, 1, 0] = torch.squeeze(xyC + zs)
    rot[:, 1, 1] = torch.squeeze(y * yC + ca)
    rot[:, 1, 2] = torch.squeeze(yzC - xs)
    rot[:, 2, 0] = torch.squeeze(zxC - ys)
    rot[:, 2, 1] = torch.squeeze(yzC + xs)
    rot[:, 2, 2] = torch.squeeze(z * zC + ca)
    rot[:, 3, 3] = 1

    return rot


class ConvBlock(nn.Module):
    """Layer to perform a convolution followed by ELU
    """
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()

        self.conv = Conv3x3(in_channels, out_channels)
        self.nonlin = nn.ELU(inplace=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.nonlin(out)
        return out


class Conv3x3(nn.Module):
    """Layer to pad and convolve input
    """
    def __init__(self, in_channels, out_channels, use_refl=True):
        super(Conv3x3, self).__init__()

        if use_refl:
            self.pad = nn.ReflectionPad2d(1)
        else:
            self.pad = nn.ZeroPad2d(1)
        self.conv = nn.Conv2d(int(in_channels), int(out_channels), 3)
        # self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.pad(x)
        out = self.conv(out)
        # out = self.bn(out)
        return out


class BackprojectDepth(nn.Module):
    """Layer to transform a depth image into a point cloud
    """
    def __init__(self, batch_size, height, width):
        super(BackprojectDepth, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width

        meshgrid = np.meshgrid(range(self.width), range(self.height), indexing='xy')
        self.id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
        self.id_coords = nn.Parameter(torch.from_numpy(self.id_coords),
                                      requires_grad=False)

        self.ones = nn.Parameter(torch.ones(self.batch_size, 1, self.height * self.width),
                                 requires_grad=False)

        self.pix_coords = torch.unsqueeze(torch.stack(
            [self.id_coords[0].view(-1), self.id_coords[1].view(-1)], 0), 0)
        self.pix_coords = self.pix_coords.repeat(batch_size, 1, 1)
        self.pix_coords = nn.Parameter(torch.cat([self.pix_coords, self.ones], 1),
                                       requires_grad=False)

    def forward(self, depth, inv_K):
        cam_points = torch.matmul(inv_K[:, :3, :3], self.pix_coords)
        cam_points = depth.view(self.batch_size, 1, -1) * cam_points
        cam_points = torch.cat([cam_points, self.ones], 1)

        return cam_points

def ResidualBlock(in_channels, out_channels, stride=1):
	return torch.nn.Sequential(
		nn.PReLU(),
		nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True),
        	nn.PReLU(),
		nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True)
	)

# upsample block
def UpsampleBlock(in_channels, out_channels, stride=1):
	return torch.nn.Sequential(
		nn.Upsample(scale_factor=2, mode='bilinear'),
		nn.PReLU(),
		nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True),
        # 	nn.PReLU(),
		# nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True),
	)


class Project3D(nn.Module):
    """Layer which projects 3D points into a camera with intrinsics K and at position T
    """
    def __init__(self, batch_size, height, width, eps=1e-7):
        super(Project3D, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width
        self.eps = eps

    def forward(self, points, K, T):
        P = torch.matmul(K, T)[:, :3, :]

        cam_points = torch.matmul(P, points)

        pix_coords = cam_points[:, :2, :] / (cam_points[:, 2, :].unsqueeze(1) + self.eps)
        pix_coords = pix_coords.view(self.batch_size, 2, self.height, self.width)
        pix_coords = pix_coords.permute(0, 2, 3, 1)
        pix_coords[..., 0] /= self.width - 1
        pix_coords[..., 1] /= self.height - 1
        pix_coords = (pix_coords - 0.5) * 2
        return pix_coords


def upsample(x, scale_factor=2, mode="nearest"):
    """Upsample input tensor by a factor of 2
    """
    return F.interpolate(x, scale_factor=scale_factor, mode=mode)


def get_smooth_loss(disp, img):
    """Computes the smoothness loss for a disparity image
    The color image is used for edge-aware smoothness
    """
    grad_disp_x = torch.abs(disp[:, :, :, :-1] - disp[:, :, :, 1:])
    grad_disp_y = torch.abs(disp[:, :, :-1, :] - disp[:, :, 1:, :])

    grad_img_x = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]), 1, keepdim=True)
    grad_img_y = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]), 1, keepdim=True)
    grad_disp_x *= torch.exp(-grad_img_x)
    grad_disp_y *= torch.exp(-grad_img_y)
    return grad_disp_x.mean() + grad_disp_y.mean()


class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images
    """
    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool   = nn.AvgPool2d(3, 1)
        self.mu_y_pool   = nn.AvgPool2d(3, 1)
        self.sig_x_pool  = nn.AvgPool2d(3, 1)
        self.sig_y_pool  = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01 ** 2
        self.C2 = 0.03 ** 2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x  = self.sig_x_pool(x ** 2) - mu_x ** 2
        sigma_y  = self.sig_y_pool(y ** 2) - mu_y ** 2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x ** 2 + mu_y ** 2 + self.C1) * (sigma_x + sigma_y + self.C2)
        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)


def compute_depth_errors(gt, pred):
    """Computation of error metrics between predicted and ground truth depths
    """
    thresh = torch.max((gt / pred), (pred / gt))
    a1 = (thresh < 1.25     ).float().mean()
    a2 = (thresh < 1.25 ** 2).float().mean()
    a3 = (thresh < 1.25 ** 3).float().mean()

    rmse = (gt - pred) ** 2
    rmse = torch.sqrt(rmse.mean())

    rmse_log = (torch.log(gt) - torch.log(pred)) ** 2
    rmse_log = torch.sqrt(rmse_log.mean())

    abs_rel = torch.mean(torch.abs(gt - pred) / gt)

    sq_rel = torch.mean((gt - pred) ** 2 / gt)

    return abs_rel, sq_rel, rmse, rmse_log, a1, a2, a3

layers.py
Sauf4896/BDEdepth
from .depth_resnet import *
from .depth_hrnet import *
from .posenet import *
from .hrnet import hrnet18

networks/__init__.py
Sauf4896/BDEdepth


from __future__ import absolute_import, division, print_function

import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models
import torch.utils.model_zoo as model_zoo

from collections import OrderedDict
from layers import *
from timm.models.layers import trunc_normal_
from .hrnet import hrnet18

class DepthEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained):
        super(DepthEncoder, self).__init__()
        assert num_layers==18

        if num_layers==18:
            self.encoder = hrnet18(pretrained)

        self.num_ch_enc = self.encoder.num_ch_enc

        self.context_conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.context_bn = nn.BatchNorm2d(64)
        self.context_conv.weight = nn.Parameter(self.encoder.conv1.weight.clone())
        self.context_bn.weight = nn.Parameter(self.encoder.bn1.weight.clone())
        self.context_bn.bias = nn.Parameter(self.encoder.bn1.bias.clone())    


    def forward(self, x):
        x = (x - 0.45) / 0.225
        self.features = self.encoder(x)
        context = self.encoder.relu(self.context_bn(self.context_conv(x)))
        self.features = [context] + self.features
        return self.features


class DepthDecoder(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), columns=3, num_output_channels=1):
        super(DepthDecoder, self).__init__()

        self.num_output_channels = num_output_channels
        self.scales = scales
        self.num_ch_enc = list(num_ch_enc)
        self.num_ch_enc = [self.num_ch_enc[0]] + self.num_ch_enc
        self.grid_ch = [ch//2 for ch in self.num_ch_enc]
        self.grid_ch[0] = 16
        # self.grid_ch = [16,32,32, 64, 128, 256]
        self.columns = columns

        self.convs = OrderedDict()
        for i in range(len(self.grid_ch)):
            for j in range(self.columns):
                if j == 0:
                    self.convs[('lateral', i, j)] = ResidualBlock(self.num_ch_enc[i], self.grid_ch[i], stride=1)
                else:
                    self.convs[('lateral', i, j)] = ResidualBlock(self.grid_ch[i], self.grid_ch[i], stride=1)

                if i < len(self.grid_ch)-1:
                    self.convs[('upconv', i, j)] = UpsampleBlock(self.grid_ch[i+1], self.grid_ch[i], stride=1)

            if i in self.scales:
                self.convs[("dispconv", i)] = Conv3x3(self.grid_ch[i], self.num_output_channels)
                

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, input_features):
        self.outputs = {}
        lateral_tmp = [0 for _ in range(self.columns)]
      
        for i in range(len(self.grid_ch)-1, -1, -1):
            x = input_features[i]
            for j in range(self.columns):
                x = self.convs[('lateral', i, j)](x)
                if i < len(self.grid_ch)-1:
                    x = x + lateral_tmp[j]
                if i > 0:
                    lateral_tmp[j] = self.convs[('upconv', i-1, j)](x)

            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))
  
        return self.outputs


networks/depth_hrnet.py
Sauf4896/BDEdepth

from __future__ import absolute_import, division, print_function

import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models
import torch.utils.model_zoo as model_zoo

from collections import OrderedDict
from layers import *
from timm.models.layers import trunc_normal_



class DepthEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained):
        super(DepthEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        self.encoder = resnets[num_layers](pretrained)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4
        
        self.conv1 = nn.Conv2d(3, self.num_ch_enc[0], kernel_size=7, stride=1, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.num_ch_enc[0])

        self.conv1.weight = nn.Parameter(self.encoder.conv1.weight.clone())
        self.bn1.weight = nn.Parameter(self.encoder.bn1.weight.clone())
        self.bn1.bias = nn.Parameter(self.encoder.bn1.bias.clone())


    def forward(self, input_image):

        self.features = []
        x = (input_image - 0.45) / 0.225
        context = self.encoder.relu(self.bn1(self.conv1(x)))
        x = self.encoder.conv1(x)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))

        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        self.features = [context] + self.features
        return self.features


class DepthDecoder(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), columns=3,num_output_channels=1):
        super(DepthDecoder, self).__init__()

        self.num_output_channels = num_output_channels
        self.scales = scales
        self.num_ch_enc = list(num_ch_enc)
        self.num_ch_enc = [self.num_ch_enc[0]] + self.num_ch_enc
        self.grid_ch = [ch//2 for ch in self.num_ch_enc]
        self.grid_ch[0] = 16
        self.columns = columns

        self.convs = OrderedDict()
        for i in range(len(self.grid_ch)):
            for j in range(self.columns):
                if j == 0:
                    self.convs[('lateral', i, j)] = ResidualBlock(self.num_ch_enc[i], self.grid_ch[i], stride=1)
                else:
                    self.convs[('lateral', i, j)] = ResidualBlock(self.grid_ch[i], self.grid_ch[i], stride=1)

                if i < len(self.grid_ch)-1:
                    self.convs[('upconv', i, j)] = UpsampleBlock(self.grid_ch[i+1], self.grid_ch[i], stride=1)

            if i in self.scales:
                self.convs[("dispconv", i)] = Conv3x3(self.grid_ch[i], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, input_features):
        self.outputs = {}
        lateral_tmp = [0 for _ in range(self.columns)]
      
        for i in range(len(self.grid_ch)-1, -1, -1):
            x = input_features[i]
            for j in range(self.columns):
                x = self.convs[('lateral', i, j)](x)
                if i < len(self.grid_ch)-1:
                    x = x + lateral_tmp[j]
                if i > 0:
                    lateral_tmp[j] = self.convs[('upconv', i-1, j)](x)

            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))
  
        return self.outputs
networks/depth_resnet.py
Sauf4896/BDEdepth
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models.utils import load_state_dict_from_url
import matplotlib.pyplot as plt

__all__ = ['hrnet18', 'hrnet32', 'hrnet48','hrnet64']


model_urls = {
    'hrnet18_imagenet': 'https://opr0mq.dm.files.1drv.com/y4mIoWpP2n-LUohHHANpC0jrOixm1FZgO2OsUtP2DwIozH5RsoYVyv_De5wDgR6XuQmirMV3C0AljLeB-zQXevfLlnQpcNeJlT9Q8LwNYDwh3TsECkMTWXCUn3vDGJWpCxQcQWKONr5VQWO1hLEKPeJbbSZ6tgbWwJHgHF7592HY7ilmGe39o5BhHz7P9QqMYLBts6V7QGoaKrr0PL3wvvR4w',
    'hrnet32_imagenet': 'https://opr74a.dm.files.1drv.com/y4mKOuRSNGQQlp6wm_a9bF-UEQwp6a10xFCLhm4bqjDu6aSNW9yhDRM7qyx0vK0WTh42gEaniUVm3h7pg0H-W0yJff5qQtoAX7Zze4vOsqjoIthp-FW3nlfMD0-gcJi8IiVrMWqVOw2N3MbCud6uQQrTaEAvAdNjtjMpym1JghN-F060rSQKmgtq5R-wJe185IyW4-_c5_ItbhYpCyLxdqdEQ',
    'hrnet48_imagenet': 'https://optgaw.dm.files.1drv.com/y4mWNpya38VArcDInoPaL7GfPMgcop92G6YRkabO1QTSWkCbo7djk8BFZ6LK_KHHIYE8wqeSAChU58NVFOZEvqFaoz392OgcyBrq_f8XGkusQep_oQsuQ7DPQCUrdLwyze_NlsyDGWot0L9agkQ-M_SfNr10ETlCF5R7BdKDZdupmcMXZc-IE3Ysw1bVHdOH4l-XEbEKFAi6ivPUbeqlYkRMQ',
    'hrnet48_cityscapes': 'https://optgaw.dm.files.1drv.com/y4mWNpya38VArcDInoPaL7GfPMgcop92G6YRkabO1QTSWkCbo7djk8BFZ6LK_KHHIYE8wqeSAChU58NVFOZEvqFaoz392OgcyBrq_f8XGkusQep_oQsuQ7DPQCUrdLwyze_NlsyDGWot0L9agkQ-M_SfNr10ETlCF5R7BdKDZdupmcMXZc-IE3Ysw1bVHdOH4l-XEbEKFAi6ivPUbeqlYkRMQ',
}


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):

    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class HighResolutionModule(nn.Module):
    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
                 num_channels, fuse_method, multi_scale_output=True, norm_layer=None):
        super(HighResolutionModule, self).__init__()
        self._check_branches(
            num_branches, blocks, num_blocks, num_inchannels, num_channels)

        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.norm_layer = norm_layer

        self.num_inchannels = num_inchannels
        self.fuse_method = fuse_method
        self.num_branches = num_branches

        self.multi_scale_output = multi_scale_output

        self.branches = self._make_branches(
            num_branches, blocks, num_blocks, num_channels)
        self.fuse_layers = self._make_fuse_layers()
        self.relu = nn.ReLU(inplace=True)

    def _check_branches(self, num_branches, blocks, num_blocks,
                        num_inchannels, num_channels):
        if num_branches != len(num_blocks):
            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
                num_branches, len(num_blocks))
            raise ValueError(error_msg)

        if num_branches != len(num_channels):
            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
                num_branches, len(num_channels))
            raise ValueError(error_msg)

        if num_branches != len(num_inchannels):
            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
                num_branches, len(num_inchannels))
            raise ValueError(error_msg)

    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
                         stride=1):
        downsample = None
        if stride != 1 or \
                self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.num_inchannels[branch_index],
                          num_channels[branch_index] * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                self.norm_layer(num_channels[branch_index] * block.expansion),
            )

        layers = []
        layers.append(block(self.num_inchannels[branch_index],
                            num_channels[branch_index], stride, downsample, norm_layer=self.norm_layer))
        self.num_inchannels[branch_index] = \
            num_channels[branch_index] * block.expansion
        for i in range(1, num_blocks[branch_index]):
            layers.append(block(self.num_inchannels[branch_index],
                                num_channels[branch_index], norm_layer=self.norm_layer))

        return nn.Sequential(*layers)

    def _make_branches(self, num_branches, block, num_blocks, num_channels):
        branches = []

        for i in range(num_branches):
            branches.append(
                self._make_one_branch(i, block, num_blocks, num_channels))

        return nn.ModuleList(branches)

    def _make_fuse_layers(self):
        if self.num_branches == 1:
            return None

        num_branches = self.num_branches
        num_inchannels = self.num_inchannels
        fuse_layers = []
        for i in range(num_branches if self.multi_scale_output else 1):
            fuse_layer = []
            for j in range(num_branches):
                if j > i:
                    fuse_layer.append(nn.Sequential(
                        nn.Conv2d(num_inchannels[j],
                                  num_inchannels[i],
                                  1,
                                  1,
                                  0,
                                  bias=False),
                        self.norm_layer(num_inchannels[i])))
                elif j == i:
                    fuse_layer.append(None)
                else:
                    conv3x3s = []
                    for k in range(i-j):
                        if k == i - j - 1:
                            num_outchannels_conv3x3 = num_inchannels[i]
                            conv3x3s.append(nn.Sequential(
                                nn.Conv2d(num_inchannels[j],
                                          num_outchannels_conv3x3,
                                          3, 2, 1, bias=False),
                                self.norm_layer(num_outchannels_conv3x3)))
                        else:
                            num_outchannels_conv3x3 = num_inchannels[j]
                            conv3x3s.append(nn.Sequential(
                                nn.Conv2d(num_inchannels[j],
                                          num_outchannels_conv3x3,
                                          3, 2, 1, bias=False),
                                self.norm_layer(num_outchannels_conv3x3),
                                nn.ReLU(inplace=True)))
                    fuse_layer.append(nn.Sequential(*conv3x3s))
            fuse_layers.append(nn.ModuleList(fuse_layer))

        return nn.ModuleList(fuse_layers)

    def get_num_inchannels(self):
        return self.num_inchannels

    def forward(self, x):
        if self.num_branches == 1:
            return [self.branches[0](x[0])]

        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])

        x_fuse = []
        for i in range(len(self.fuse_layers)):
            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
            for j in range(1, self.num_branches):
                if i == j:
                    y = y + x[j]
                elif j > i:
                    width_output = x[i].shape[-1]
                    height_output = x[i].shape[-2]
                    y = y + F.interpolate(
                        self.fuse_layers[i][j](x[j]),
                        size=[height_output, width_output],
                        mode='bilinear',
                        align_corners=True
                        )
                else:
                    y = y + self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))

        return x_fuse


blocks_dict = {
    'BASIC': BasicBlock,
    'BOTTLENECK': Bottleneck
}


class HighResolutionNet(nn.Module):

    def __init__(self,
                 cfg,
                 norm_layer=None):
        super(HighResolutionNet, self).__init__()

        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.norm_layer = norm_layer
        # stem network
        # stem net
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn1 = self.norm_layer(64)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
                               bias=False)
        self.bn2 = self.norm_layer(64)
        self.relu = nn.ReLU(inplace=True)

        # stage 1
        self.stage1_cfg = cfg['STAGE1']
        num_channels = self.stage1_cfg['NUM_CHANNELS'][0]
        block = blocks_dict[self.stage1_cfg['BLOCK']]
        num_blocks = self.stage1_cfg['NUM_BLOCKS'][0]
        self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
        stage1_out_channel = block.expansion*num_channels

        # stage 2
        self.stage2_cfg = cfg['STAGE2']
        num_channels = self.stage2_cfg['NUM_CHANNELS']
        block = blocks_dict[self.stage2_cfg['BLOCK']]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))]
        self.transition1 = self._make_transition_layer(
            [stage1_out_channel], num_channels)
        self.stage2, pre_stage_channels = self._make_stage(
            self.stage2_cfg, num_channels)

        # stage 3
        self.stage3_cfg = cfg['STAGE3']
        num_channels = self.stage3_cfg['NUM_CHANNELS']
        block = blocks_dict[self.stage3_cfg['BLOCK']]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))]
        self.transition2 = self._make_transition_layer(
            pre_stage_channels, num_channels)
        self.stage3, pre_stage_channels = self._make_stage(
            self.stage3_cfg, num_channels)

        # stage 4
        self.stage4_cfg = cfg['STAGE4']
        num_channels = self.stage4_cfg['NUM_CHANNELS']
        block = blocks_dict[self.stage4_cfg['BLOCK']]
        num_channels = [
            num_channels[i] * block.expansion for i in range(len(num_channels))]
        self.transition3 = self._make_transition_layer(
            pre_stage_channels, num_channels)
        self.stage4, pre_stage_channels = self._make_stage(
            self.stage4_cfg, num_channels, multi_scale_output=True)
    

    def _make_transition_layer(
            self, num_channels_pre_layer, num_channels_cur_layer):
        num_branches_cur = len(num_channels_cur_layer)
        num_branches_pre = len(num_channels_pre_layer)

        transition_layers = []
        for i in range(num_branches_cur):
            if i < num_branches_pre:
                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
                    transition_layers.append(nn.Sequential(
                        nn.Conv2d(num_channels_pre_layer[i],
                                  num_channels_cur_layer[i],
                                  3,
                                  1,
                                  1,
                                  bias=False),
                        self.norm_layer(num_channels_cur_layer[i]),
                        nn.ReLU(inplace=True)))
                else:
                    transition_layers.append(None)
            else:
                conv3x3s = []
                for j in range(i+1-num_branches_pre):
                    inchannels = num_channels_pre_layer[-1]
                    outchannels = num_channels_cur_layer[i] \
                        if j == i-num_branches_pre else inchannels
                    conv3x3s.append(nn.Sequential(
                        nn.Conv2d(
                            inchannels, outchannels, 3, 2, 1, bias=False),
                        self.norm_layer(outchannels),
                        nn.ReLU(inplace=True)))
                transition_layers.append(nn.Sequential(*conv3x3s))

        return nn.ModuleList(transition_layers)

    def _make_layer(self, block, inplanes, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                self.norm_layer(planes * block.expansion),
            )

        layers = []
        layers.append(block(inplanes, planes, stride, downsample, norm_layer=self.norm_layer))
        inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(inplanes, planes, norm_layer=self.norm_layer))

        return nn.Sequential(*layers)

    def _make_stage(self, layer_config, num_inchannels,
                    multi_scale_output=True):
        num_modules = layer_config['NUM_MODULES']
        num_branches = layer_config['NUM_BRANCHES']
        num_blocks = layer_config['NUM_BLOCKS']
        num_channels = layer_config['NUM_CHANNELS']
        block = blocks_dict[layer_config['BLOCK']]
        fuse_method = layer_config['FUSE_METHOD']

        modules = []
        for i in range(num_modules):
            # multi_scale_output is only used last module
            if not multi_scale_output and i == num_modules - 1:
                reset_multi_scale_output = False
            else:
                reset_multi_scale_output = True

            modules.append(
                HighResolutionModule(num_branches,
                                     block,
                                     num_blocks,
                                     num_inchannels,
                                     num_channels,
                                     fuse_method,
                                     reset_multi_scale_output,
                                     norm_layer=self.norm_layer)
            )
            num_inchannels = modules[-1].get_num_inchannels()

        return nn.Sequential(*modules), num_inchannels


    def forward(self, x):
        outputs = []

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        outputs.append(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)

        x = self.layer1(x)
        
        x_list = []
        for i in range(self.stage2_cfg['NUM_BRANCHES']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)
        y_list = self.stage2(x_list)

        
        x_list = []
        for i in range(self.stage3_cfg['NUM_BRANCHES']):
            if self.transition2[i] is not None:
                if i < self.stage2_cfg['NUM_BRANCHES']:
                    x_list.append(self.transition2[i](y_list[i]))
                else:
                    x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage3(x_list)
        
        x_list = []
        for i in range(self.stage4_cfg['NUM_BRANCHES']):
            if self.transition3[i] is not None:
                if i < self.stage3_cfg['NUM_BRANCHES']:
                    x_list.append(self.transition3[i](y_list[i]))
                else:
                    x_list.append(self.transition3[i](y_list[-1]))
                    # here generate new scale features (downsample) 
            else:
                x_list.append(y_list[i])
        y_list = self.stage4(x_list)

        outputs += y_list #0:[12,64,96,320]; 1:[12,18,48,160]; 2:[12,36,24,80]; 3:[12,72,12,40]; 4:[12,144,6,20]

        return outputs
    
    def weight_parameters(self):
        for name, param in self.named_parameters():
            if 'alpha' not in name:
                yield param
        

def _hrnet(arch, pretrained, progress, **kwargs):
    from .hrnet_config import MODEL_CONFIGS
    model = HighResolutionNet(MODEL_CONFIGS[arch], **kwargs)
    if pretrained:
        if arch == 'hrnet64':
            arch = 'hrnet32_imagenet'
            model_url = model_urls[arch]
            loaded_state_dict = load_state_dict_from_url(model_url, progress=progress)
            #add weights demention to adopt input change
            exp_layers = ['conv1.weight', 'bn1.weight', 'bn1.bias', 'bn1.running_mean', 'bn1.running_var', 'conv2.weight', 'bn2.weight', 'bn2.bias', 'bn2.running_mean', 'bn2.running_var']
            lista = ['transition1.0.0.weight', 'transition1.1.0.0.weight', 'transition2.2.0.0.weight', 'transition3.3.0.0.weight']
            for k,v in loaded_state_dict.items() :
                if k not in exp_layers:
                    if ('layer' not in k) and 'conv' in k or k in lista and len(v.size()) > 1:
                        if k in ['transition1.0.0.weight' , 'transition1.1.0.0.weight']:
                            loaded_state_dict[k] = torch.cat([loaded_state_dict[k]] * 2,0)
                        else:
                            loaded_state_dict[k] = torch.cat([v] * 2, 1) / 2
                            loaded_state_dict[k] = torch.cat([loaded_state_dict[k]] * 2,0)

                    if 'fuse_layer' in k and 'weight' in k and len(v.size()) > 1:
                        loaded_state_dict[k] = torch.cat([v] * 2, 1) / 2
                        loaded_state_dict[k] = torch.cat([loaded_state_dict[k]] * 2,0)

                    if 'layer' not in k and len(v.size()) == 1:
                        v = v.unsqueeze(1)
                        v = torch.cat([v] * 2, 0) 
                        loaded_state_dict[k] = v.squeeze(1) 
                    if 'fuse_layer' in k and len(v.size()) == 1:
                        v = v.unsqueeze(1)
                        v = torch.cat([v] * 2, 0) 
                        loaded_state_dict[k] = v.squeeze(1) 
                    if len(loaded_state_dict[k].size()) == 2:
                        loaded_state_dict[k] = loaded_state_dict[k].squeeze(1)
                    # for multi-input 
                    #if k == 'conv1.weight':
                    #  loaded_state_dict[k] = torch.cat([v] * 2, 1) / 2
        else:
            arch = arch + '_imagenet'
            model_url = model_urls[arch]
            # loaded_state_dict = load_state_dict_from_url(model_url, progress=progress)

            # pth_path = "/hrnet_IN_pretrained/hrnetv2_w18_imagenet_pretrained.pth"
            # loaded_state_dict = torch.load(pth_path)
            pth_path = "./hrnet_IN_pretrained/HRNet_W18_C_cosinelr_cutmix_300epoch.pth.tar"
            loaded_state_dict = torch.load(pth_path, map_location='cpu')
        #if k == 'conv1.weight':
        #    loaded_state_dict[k] = torch.cat([v] * 2, 1) / 2
                   
        model.load_state_dict({k: v for k,v in loaded_state_dict.items() if k in model.state_dict()})
    return model


def hrnet18(pretrained=True, progress=True, **kwargs):
    r"""HRNet-18 model
    """
    model = _hrnet('hrnet18', pretrained, progress, **kwargs)
    model.num_ch_enc = [ 64, 18, 36, 72, 144 ]
    return model


def hrnet32(pretrained=True, progress=True, **kwargs):
    r"""HRNet-32 model
    """
    return _hrnet('hrnet32', pretrained, progress,
                   **kwargs)


def hrnet48(pretrained=True, progress=True, **kwargs):
    r"""HRNet-48 model
    """
    return _hrnet('hrnet48', pretrained, progress,
                   **kwargs)

def hrnet64(pretrained=True, progress=True, **kwargs):
    r"""HRNet-64 model
    """
    return _hrnet('hrnet64', pretrained, progress,
                   **kwargs)



networks/hrnet.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models
import torch.utils.model_zoo as model_zoo
from collections import OrderedDict


class ResNetMultiImageInput(models.ResNet):
    """Constructs a resnet model with varying number of input images.
    Adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
    """
    def __init__(self, block, layers, num_classes=1000, num_input_images=1):
        super(ResNetMultiImageInput, self).__init__(block, layers)
        self.inplanes = 64
        self.conv1 = nn.Conv2d(
            num_input_images * 3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


def resnet_multiimage_input(num_layers, pretrained=False, num_input_images=1):
    """Constructs a ResNet model.
    Args:
        num_layers (int): Number of resnet layers. Must be 18 or 50
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        num_input_images (int): Number of frames stacked as input
    """
    assert num_layers in [18, 50], "Can only run with 18 or 50 layer resnet"
    blocks = {18: [2, 2, 2, 2], 50: [3, 4, 6, 3]}[num_layers]
    block_type = {18: models.resnet.BasicBlock, 50: models.resnet.Bottleneck}[num_layers]
    model = ResNetMultiImageInput(block_type, blocks, num_input_images=num_input_images)

    if pretrained:
        loaded = model_zoo.load_url(models.resnet.model_urls['resnet{}'.format(num_layers)])
        loaded['conv1.weight'] = torch.cat(
            [loaded['conv1.weight']] * num_input_images, 1) / num_input_images
        model.load_state_dict(loaded)
    return model


class ResnetEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained, num_input_images=1):
        super(ResnetEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        if num_input_images > 1:
            self.encoder = resnet_multiimage_input(num_layers, pretrained, num_input_images)
        else:
            self.encoder = resnets[num_layers](pretrained)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4


    def forward(self, input_image):
        self.features = []
        x = (input_image - 0.45) / 0.225
        x = self.encoder.conv1(x)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))

        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        return self.features


class PoseDecoder(nn.Module):
    def __init__(self, num_ch_enc, num_input_features, num_frames_to_predict_for=None, stride=1):
        super(PoseDecoder, self).__init__()

        self.num_ch_enc = num_ch_enc
        self.num_input_features = num_input_features

        if num_frames_to_predict_for is None:
            num_frames_to_predict_for = num_input_features - 1
        self.num_frames_to_predict_for = num_frames_to_predict_for

        self.convs = OrderedDict()
        self.convs[("squeeze")] = nn.Conv2d(self.num_ch_enc[-1], 256, 1)
        self.convs[("pose", 0)] = nn.Conv2d(num_input_features * 256, 256, 3, stride, 1)
        self.convs[("pose", 1)] = nn.Conv2d(256, 256, 3, stride, 1)
        self.convs[("pose", 2)] = nn.Conv2d(256, 6 * num_frames_to_predict_for, 1)

        self.relu = nn.ReLU()

        self.net = nn.ModuleList(list(self.convs.values()))


    def forward(self, input_features):
        last_features = [f[-1] for f in input_features]

        cat_features = [self.relu(self.convs["squeeze"](f)) for f in last_features]
        cat_features = torch.cat(cat_features, 1)

        out = cat_features
        for i in range(3):
            out = self.convs[("pose", i)](out)
            if i != 2:
                out = self.relu(out)

        out = out.mean(3).mean(2)

        out = 0.01 * out.view(-1, self.num_frames_to_predict_for, 1, 6)

        axisangle = out[..., :3]
        translation = out[..., 3:]

        return axisangle, translation

networks/posenet.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import argparse

file_dir = os.path.dirname(__file__)  # the directory that options.py resides in


parser = argparse.ArgumentParser(description="Monodepthv2 options")

# Distributed Training. Do not modify!
parser.add_argument('--local_rank', 
                    default=0, 
                    type=int,
                    help='node rank for distributed training')
parser.add_argument('--global_rank', 
                    default=0, 
                    type=int,
                    help='global node rank for distributed training')
parser.add_argument('--world_size', 
                    default=1, 
                    type=int,
                    help='GPU numbers')

# PATHS
parser.add_argument("--data_path",
                    type=str,
                    help="path to the training data",
                    default=os.path.join(file_dir, "kitti_data"))
parser.add_argument("--data_path_pre",
                    type=str,
                    help="path to the preprocessed training data, specify when training on Cityscapes.")
parser.add_argument("--log_dir",
                    type=str,
                    help="log directory",
                    default=os.path.join(os.path.expanduser("~"), "tmp"))

# TRAINING options
parser.add_argument("--exp_name",
                    type=str,
                    help="the name of the folder to save the model in",
                    default="mdp")
parser.add_argument("--split",
                    type=str,
                    help="which training split to use",
                    choices=["eigen_zhou", "eigen_full", "odom", "benchmark"],
                    default="eigen_zhou")
parser.add_argument("--backbone",
                    type=str,
                    help="backbone of depth encoder",
                    default="hrnet",
                    choices=["resnet", "hrnet"])
parser.add_argument("--num_layers",
                    type=int,
                    help="number of resnet layers",
                    default=18,
                    choices=[18, 34, 50, 101, 152])
parser.add_argument("--dataset",
                    type=str,
                    help="dataset to train on",
                    default="kitti",
                    choices=["kitti", "kitti_odom", "kitti_depth", "kitti_test", "nyuv2", "cityscapes"])
parser.add_argument("--jpg",
                    help="if set, trains from raw KITTI jpg files (instead of pngs)",
                    action="store_true")
parser.add_argument("--height",
                    type=int,
                    help="input image height",
                    default=192)
parser.add_argument("--width",
                    type=int,
                    help="input image width",
                    default=640)
parser.add_argument("--disparity_smoothness",
                    type=float,
                    help="disparity smoothness weight",
                    default=1e-3)
parser.add_argument("--num_scales",
                    type=int,
                    help="scale numbers used in the loss",
                    default=1)
parser.add_argument("--min_depth",
                    type=float,
                    help="minimum depth",
                    default=0.1)
parser.add_argument("--max_depth",
                    type=float,
                    help="maximum depth",
                    default=100.0)
parser.add_argument("--use_stereo",
                    help="if set, uses stereo pair for training",
                    action="store_true")
parser.add_argument("--frame_ids",
                    nargs="+",
                    type=int,
                    help="frames to load",
                    default=[0, -1, 1])
parser.add_argument("--use_local_crop",
                    help="if set, use the local crop resize augument)",
                    action="store_true")
parser.add_argument("--use_patch_reshuffle",
                    help="if set, use the patch reshuffle augument)",
                    action="store_true")

# OPTIMIZATION options
parser.add_argument('--optimizer',
                    type=str,
                    default='adamw',
                    help='choose optimizer, (default:adamw)',
                    choices=["adamw", "adam", "sgd"])
parser.add_argument('--lr_sche_type',
                    type=str,
                    default='step',
                    help='choose lr_scheduler, (default:cos)',
                    choices=["cos", "step"])
parser.add_argument('--warm_up',
                    type=float,
                    default=0.0,
                    help='percent of steps for warmup when using cosine annealing.')
parser.add_argument('--eta_min',
                    type=float,
                    default=5e-6,
                    help='eta_min is the final learning rate when using cosine annealing.')
parser.add_argument("--batch_size",
                    type=int,
                    help="batch size",
                    default=12)
parser.add_argument('--learning_rate',
                    default=0.0001,
                    type=float,
                    help='learning rate')
parser.add_argument('--decay_rate',
                    dest='decay_rate',
                    type=float,
                    default=0.1,
                    help='# of optimizer')
parser.add_argument('--decay_step',
                    dest='decay_step',
                    type=int,
                    nargs='+',
                    default=[15],
                    help='# of optimizer')
parser.add_argument('--weight_decay',
                    type=float,
                    default=0.01,
                    help='Weight decay for adamw.')
parser.add_argument('--beta1',
                    type=float,
                    default=0.9,
                    help='beta1 for adamw.')
parser.add_argument('--beta2',
                    type=float,
                    default=0.999,
                    help='beta2 for adamw.')
parser.add_argument('--momentum',
                    default=0.9,
                    type=float,
                    help='momentum for sgd')
parser.add_argument('--clip_grad',
                    dest='clip_grad',
                    type=float,
                    default=-1,
                    help='# of the parameters')
parser.add_argument("--num_epochs",
                    type=int,
                    help="number of epochs",
                    default=20)
parser.add_argument("--seed",
                    type=int,
                    help="seed for random",
                    default=1234)
parser.add_argument("--resume",
                    help="if set, try to resume from a exsited ckpt",
                    action="store_true")

        # ABLATION options
parser.add_argument("--v1_multiscale",
                    help="if set, uses monodepth v1 multiscale",
                    action="store_true")
parser.add_argument("--avg_reprojection",
                    help="if set, uses average reprojection loss",
                    action="store_true")
parser.add_argument("--disable_automasking",
                    help="if set, doesn't do auto-masking",
                    action="store_true")
parser.add_argument("--predictive_mask",
                    help="if set, uses a predictive masking scheme as in Zhou et al",
                    action="store_true")
parser.add_argument("--no_ssim",
                    help="if set, disables ssim in the loss",
                    action="store_true")
parser.add_argument("--weights_init",
                    type=str,
                    help="pretrained or scratch",
                    default="pretrained",
                    choices=["pretrained", "scratch"])
parser.add_argument("--pose_model_input",
                    type=str,
                    help="how many images the pose network gets",
                    default="pairs",
                    choices=["pairs", "all"])


# SYSTEM options
parser.add_argument("--no_cuda",
                    help="if set disables CUDA",
                    action="store_true")
parser.add_argument("--num_workers",
                    type=int,
                    help="number of dataloader workers",
                    default=12)

# LOADING options
parser.add_argument("--pretrained_path",
                    type=str,
                    help="path of pretrained model to load")
parser.add_argument("--models_to_load",
                    nargs="+",
                    type=str,
                    help="models to load",
                    default=["encoder", "depth", "pose_encoder", "pose"])

# LOGGING options
parser.add_argument("--log_frequency",
                    type=int,
                    help="number of batches between each log",
                    default=250)
parser.add_argument("--save_frequency",
                    type=int,
                    help="number of epochs between each save",
                    default=500)

# EVALUATION options
parser.add_argument("--eval_stereo",
                    help="if set evaluates in stereo mode",
                    action="store_true")
parser.add_argument("--eval_mono",
                    help="if set evaluates in mono mode",
                    action="store_true")
parser.add_argument("--disable_median_scaling",
                    help="if set disables median scaling in evaluation",
                    action="store_true")
parser.add_argument("--pred_depth_scale_factor",
                    help="if set multiplies predictions by this number",
                    type=float,
                    default=1)
parser.add_argument("--ext_disp_to_eval",
                    type=str,
                    help="optional path to a .npy disparities file to evaluate")
parser.add_argument("--eval_split",
                    type=str,
                    default="eigen",
                    choices=["eigen", "eigen_benchmark", "benchmark", "odom_9", "odom_10"],
                    help="which split to run eval on")
parser.add_argument("--save_pred_disps",
                    help="if set saves predicted disparities",
                    action="store_true")
parser.add_argument("--no_eval",
                    help="if set disables evaluation",
                    action="store_true")
parser.add_argument("--eval_eigen_to_benchmark",
                    help="if set assume we are loading eigen results from npy but " "we want to evaluate using the new benchmark.",
                    action="store_true")
parser.add_argument("--eval_out_dir",
                    help="if set will output the disparities to this folder",
                    type=str)
parser.add_argument("--post_process",
                    help="if set will perform the flipping post processing "
                         "from the original monodepth paper",
                    action="store_true")


opts = parser.parse_args()
options.py
Sauf4896/BDEdepth
from __future__ import division
import argparse
import PIL.Image as pil
import numpy as np
from glob import glob
from joblib import Parallel, delayed
import json
import os

parser = argparse.ArgumentParser()
parser.add_argument("--dataset_dir", type=str, required=True, help="where the dataset is stored")
parser.add_argument("--dump_root", type=str, required=True, help="Where to dump the data")
parser.add_argument("--seq_length", type=int, required=True, help="Length of each training sequence")
parser.add_argument("--img_height", type=int, default=128, help="image height")
parser.add_argument("--img_width", type=int, default=416, help="image width")
parser.add_argument("--num_threads", type=int, default=4, help="number of threads to use")
args = parser.parse_args()


class cityscapes_loader(object):
    def __init__(self, 
                 dataset_dir,
                 split='train',
                 crop_bottom=True, # Get rid of the car logo
                 sample_gap=2,  # Sample every two frames to match KITTI frame rate
                 img_height=171, 
                 img_width=416,
                 seq_length=5):
        self.dataset_dir = dataset_dir
        self.split = split
        # Crop out the bottom 25% of the image to remove the car logo
        self.crop_bottom = crop_bottom
        self.sample_gap = sample_gap
        self.img_height = img_height
        self.img_width = img_width
        self.seq_length = seq_length
        assert seq_length % 2 != 0, 'seq_length must be odd!'
        self.frames = self.collect_frames(split)
        self.num_frames = len(self.frames)
        if split == 'train':
            self.num_train = self.num_frames
        else:
            self.num_test = self.num_frames
        print('Total frames collected: %d' % self.num_frames)
        
    def collect_frames(self, split):
        img_dir = self.dataset_dir + '/leftImg8bit_sequence/' + split + '/'
        city_list = os.listdir(img_dir)
        frames = []
        for city in city_list:
            img_files = glob(img_dir + city + '/*.png')
            for f in img_files:
                frame_id = os.path.basename(f).split('leftImg8bit')[0]
                frames.append(frame_id)
        return frames

    def get_train_example_with_idx(self, tgt_idx):
        tgt_frame_id = self.frames[tgt_idx]
        if not self.is_valid_example(tgt_frame_id):
            return False
        example = self.load_example(self.frames[tgt_idx])
        return example

    def load_intrinsics(self, frame_id, split):
        city, seq, _, _ = frame_id.split('_')
        camera_file = os.path.join(self.dataset_dir, 'camera',
                                   split, city, city + '_' + seq + '_*_camera.json')
        camera_file = glob(camera_file)[0]
        with open(camera_file, 'r') as f: 
            camera = json.load(f)
        fx = camera['intrinsic']['fx']
        fy = camera['intrinsic']['fy']
        u0 = camera['intrinsic']['u0']
        v0 = camera['intrinsic']['v0']
        intrinsics = np.array([[fx, 0, u0],
                               [0, fy, v0],
                               [0,  0,  1]])
        return intrinsics

    def is_valid_example(self, tgt_frame_id):
        city, snippet_id, tgt_local_frame_id, _ = tgt_frame_id.split('_')
        half_offset = int((self.seq_length - 1)/2 * self.sample_gap)
        for o in range(-half_offset, half_offset + 1, self.sample_gap):
            curr_local_frame_id = '%.6d' % (int(tgt_local_frame_id) + o)
            curr_frame_id = '%s_%s_%s_' % (city, snippet_id, curr_local_frame_id)
            curr_image_file = os.path.join(self.dataset_dir, 'leftImg8bit_sequence', 
                                self.split, city, curr_frame_id + 'leftImg8bit.png')
            if not os.path.exists(curr_image_file):
                return False
        return True

    def load_image_sequence(self, tgt_frame_id, seq_length, crop_bottom):
        city, snippet_id, tgt_local_frame_id, _ = tgt_frame_id.split('_')
        half_offset = int((self.seq_length - 1)/2 * self.sample_gap)
        image_seq = []
        for o in range(-half_offset, half_offset + 1, self.sample_gap):
            curr_local_frame_id = '%.6d' % (int(tgt_local_frame_id) + o)
            curr_frame_id = '%s_%s_%s_' % (city, snippet_id, curr_local_frame_id)
            curr_image_file = os.path.join(self.dataset_dir, 'leftImg8bit_sequence', 
                                self.split, city, curr_frame_id + 'leftImg8bit.png')
            curr_img = pil.open(curr_image_file).convert('RGB')
            raw_shape = np.copy(curr_img.size)
            if o == 0:
                zoom_y = self.img_height/raw_shape[1]
                zoom_x = self.img_width/raw_shape[0]
            curr_img = curr_img.resize((self.img_width, self.img_height), pil.LANCZOS)
            if crop_bottom:
                ymax = int(curr_img.size[1] * 0.75)
                curr_img = curr_img.crop((0, 0, curr_img.size[0], ymax))
            image_seq.append(curr_img)
        return image_seq, zoom_x, zoom_y
    
    def load_example(self, tgt_frame_id, load_gt_pose=False):
        image_seq, zoom_x, zoom_y = self.load_image_sequence(tgt_frame_id, self.seq_length, self.crop_bottom)
        intrinsics = self.load_intrinsics(tgt_frame_id, self.split)
        intrinsics = self.scale_intrinsics(intrinsics, zoom_x, zoom_y)
        example = {}
        example['intrinsics'] = intrinsics
        example['image_seq'] = image_seq
        example['folder_name'] = tgt_frame_id.split('_')[0]
        example['file_name'] = tgt_frame_id[:-1]
        return example

    def scale_intrinsics(self, mat, sx, sy):
        out = np.copy(mat)
        out[0,0] *= sx
        out[0,2] *= sx
        out[1,1] *= sy
        out[1,2] *= sy
        return out
    

def concat_image_seq(seq):
    for i, im in enumerate(seq):
        if i == 0:
            res = np.array(im)
        else:
            im = np.array(im)
            res = np.concatenate((res, im))
    res = pil.fromarray(res)
    return res

def dump_example(n, args):
    if n % 2000 == 0:
        print('Progress %d/%d....' % (n, data_loader.num_train))
    example = data_loader.get_train_example_with_idx(n)
    if example == False:
        return
    image_seq = concat_image_seq(example['image_seq'])
    intrinsics = example['intrinsics']
    fx = intrinsics[0, 0]
    fy = intrinsics[1, 1]
    cx = intrinsics[0, 2]
    cy = intrinsics[1, 2]
    dump_dir = os.path.join(args.dump_root, example['folder_name'])
    # if not os.path.isdir(dump_dir):
    #     os.makedirs(dump_dir, exist_ok=True)
    try: 
        os.makedirs(dump_dir)
    except OSError:
        if not os.path.isdir(dump_dir):
            raise
    dump_img_file = dump_dir + '/%s.png' % example['file_name']
    image_seq.save(dump_img_file)
    dump_cam_file = dump_dir + '/%s_cam.txt' % example['file_name']
    with open(dump_cam_file, 'w') as f:
        f.write('%f,0.,%f,0.,%f,%f,0.,0.,1.' % (fx, cx, fy, cy))

def main():
    if not os.path.exists(args.dump_root):
        os.makedirs(args.dump_root)

    global data_loader

    data_loader = cityscapes_loader(args.dataset_dir,
                                    img_height=args.img_height,
                                    img_width=args.img_width,
                                    seq_length=args.seq_length)

    Parallel(n_jobs=args.num_threads)(delayed(dump_example)(n, args) for n in range(data_loader.num_train))

    # # Split into train/val
    # np.random.seed(8964)
    # subfolders = os.listdir(args.dump_root)
    # with open(args.dump_root + 'train.txt', 'w') as tf:
    #     with open(args.dump_root + 'val.txt', 'w') as vf:
    #         for s in subfolders:
    #             if not os.path.isdir(args.dump_root + '/%s' % s):c
    #                 continue
    #             imfiles = glob(os.path.join(args.dump_root, s, '*.jpg'))
    #             frame_ids = [os.path.basename(fi).split('.')[0] for fi in imfiles]
    #             for frame in frame_ids:
    #                 if np.random.random() < 0.1:
    #                     vf.write('%s %s\n' % (s, frame))
    #                 else:
    #                     tf.write('%s %s\n' % (s, frame))

main()
prepare_cityscapes.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import sys
import glob
import argparse
import numpy as np
import PIL.Image as pil
import matplotlib as mpl
import matplotlib.cm as cm

import torch
from torchvision import transforms, datasets

from networks import *
from layers import disp_to_depth
from evaluate_depth import STEREO_SCALE_FACTOR


def parse_args():
    parser = argparse.ArgumentParser(
        description='Simple testing funtion for Monodepthv2 models.')

    parser.add_argument('--image_path', type=str,
                        help='path to a test image or folder of images', required=True)
    parser.add_argument('--pretrained_path',
                        type=str,
                        help="path of model checkpoint to load")
    parser.add_argument("--backbone",
                        type=str,
                        help="backbone of depth encoder",
                        default="hrnet",
                        choices=["resnet", "hrnet"])
    parser.add_argument("--num_layers",
                        type=int,
                        help="number of resnet layers",
                        default=18,
                        choices=[18, 34, 50, 101, 152])
    parser.add_argument("--height",
                        type=int,
                        help="input image height",
                        default=192)
    parser.add_argument("--width",
                        type=int,
                        help="input image width",
                        default=640)
    parser.add_argument("--min_depth",
                        type=float,
                        help="minimum depth",
                        default=0.1)
    parser.add_argument("--max_depth",
                        type=float,
                        help="maximum depth",
                        default=100.0)
    parser.add_argument('--ext', type=str,
                        help='image extension to search for in folder', default="png")
    parser.add_argument("--no_cuda",
                        help='if set, disables CUDA',
                        action='store_true')
    parser.add_argument("--save_npy",
                        help='if set, saves numpy files of predicted disps/depths',
                        action='store_true')

    return parser.parse_args()


def test_simple(args):
    """Function to predict for a single image or folder of images
    """
    assert args.pretrained_path is not None

    if torch.cuda.is_available() and not args.no_cuda:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    print("-> Loading model from ", args.pretrained_path)
    model = torch.load(args.pretrained_path, map_location='cpu')

    if args.backbone == "hrnet":
        depth_encoder = depth_hrnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_hrnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))
    if args.backbone == "resnet":
        depth_encoder = depth_resnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_resnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))

    depth_encoder.load_state_dict({k: v for k, v in model["encoder"].items() if k in depth_encoder.state_dict()})
    depth_decoder.load_state_dict({k: v for k, v in model["depth"].items() if k in depth_decoder.state_dict()})
    depth_encoder.to(device).eval()
    depth_decoder.to(device).eval()
    feed_height = args.height
    feed_width = args.width


    # FINDING INPUT IMAGES
    if os.path.isfile(args.image_path):
        # Only testing on a single image
        paths = [args.image_path]
        output_directory = os.path.dirname(args.image_path)
    elif os.path.isdir(args.image_path):
        # Searching folder for images
        paths = glob.glob(os.path.join(args.image_path, '*.{}'.format(args.ext)))
        output_directory = args.image_path
    else:
        raise Exception("Can not find args.image_path: {}".format(args.image_path))

    print("-> Predicting on {:d} test images".format(len(paths)))

    # PREDICTING ON EACH IMAGE IN TURN
    with torch.no_grad():
        for idx, image_path in enumerate(paths):

            if image_path.endswith("_disp.jpg"):
                # don't try to predict disparity for a disparity image!
                continue

            # Load image and preprocess
            input_image = pil.open(image_path).convert('RGB')
            original_width, original_height = input_image.size
            input_image = input_image.resize((feed_width, feed_height), pil.LANCZOS)
            input_image = transforms.ToTensor()(input_image).unsqueeze(0)

            # PREDICTION
            input_image = input_image.to(device)
            features = depth_encoder(input_image)
            outputs = depth_decoder(features)

            disp = outputs[("disp", 0)]
            disp_resized = torch.nn.functional.interpolate(
                disp, (original_height, original_width), mode="bilinear", align_corners=False)

            # Saving numpy file
            output_name = os.path.splitext(os.path.basename(image_path))[0]
            scaled_disp, depth = disp_to_depth(disp, args.min_depth, args.max_depth)
            if args.save_npy:
                name_dest_npy = os.path.join(output_directory, "{}_disp.npy".format(output_name))
                np.save(name_dest_npy, scaled_disp.cpu().numpy())

            # Saving colormapped depth image
            disp_resized_np = disp_resized.squeeze().cpu().numpy()
            vmax = np.percentile(disp_resized_np, 95)
            normalizer = mpl.colors.Normalize(vmin=disp_resized_np.min(), vmax=vmax)
            mapper = cm.ScalarMappable(norm=normalizer, cmap='magma')
            colormapped_im = (mapper.to_rgba(disp_resized_np)[:, :, :3] * 255).astype(np.uint8)
            im = pil.fromarray(colormapped_im)

            name_dest_im = os.path.join(output_directory, "{}_disp.jpeg".format(output_name))
            im.save(name_dest_im)

            print("   Processed {:d} of {:d} images - saved predictions to:".format(
                idx + 1, len(paths)))
            print("   - {}".format(name_dest_im))
            if args.save_npy:
                print("   - {}".format(name_dest_npy))

    print('-> Done!')


if __name__ == '__main__':
    args = parse_args()
    test_simple(args)
test_simple.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import sys
import glob
import argparse
import numpy as np
import PIL.Image as pil
import matplotlib as mpl
import matplotlib.cm as cm

import torch
from torchvision import transforms, datasets

from networks import *
from layers import disp_to_depth
from evaluate_depth import STEREO_SCALE_FACTOR


def parse_args():
    parser = argparse.ArgumentParser(
        description='Simple testing funtion for Monodepthv2 models.')

    parser.add_argument('--image_path', type=str,
                        help='path to a test image or folder of images', required=True)
    parser.add_argument('--pretrained_path',
                        type=str,
                        help="path of model checkpoint to load")
    parser.add_argument("--backbone",
                        type=str,
                        help="backbone of depth encoder",
                        default="hrnet",
                        choices=["resnet", "hrnet"])
    parser.add_argument("--num_layers",
                        type=int,
                        help="number of resnet layers",
                        default=18,
                        choices=[18, 34, 50, 101, 152])
    parser.add_argument("--height",
                        type=int,
                        help="input image height",
                        default=192)
    parser.add_argument("--width",
                        type=int,
                        help="input image width",
                        default=640)
    parser.add_argument("--min_depth",
                        type=float,
                        help="minimum depth",
                        default=0.1)
    parser.add_argument("--max_depth",
                        type=float,
                        help="maximum depth",
                        default=100.0)
    parser.add_argument('--ext', type=str,
                        help='image extension to search for in folder', default="png")
    parser.add_argument("--no_cuda",
                        help='if set, disables CUDA',
                        action='store_true')

    return parser.parse_args()


def test_simple(args):
    """Function to predict for a single image or folder of images
    """
    assert args.pretrained_path is not None

    if torch.cuda.is_available() and not args.no_cuda:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    print("-> Loading model from ", args.pretrained_path)
    model = torch.load(args.pretrained_path, map_location='cpu')
    if args.backbone == "hrnet":
        depth_encoder = depth_hrnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_hrnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))
    if args.backbone == "resnet":
        depth_encoder = depth_resnet.DepthEncoder(args.num_layers, False)
        depth_decoder = depth_resnet.DepthDecoder(depth_encoder.num_ch_enc, scales=range(1))
    depth_encoder.load_state_dict({k: v for k, v in model["encoder"].items() if k in depth_encoder.state_dict()})
    depth_decoder.load_state_dict({k: v for k, v in model["depth"].items() if k in depth_decoder.state_dict()})
    depth_encoder.to(device).eval()
    depth_decoder.to(device).eval()
    feed_height = args.height
    feed_width = args.width


    assert os.path.isdir(args.image_path)

    paths = glob.glob(os.path.join(args.image_path, '*.{}'.format(args.ext)))
    paths = sorted(paths)
 
    output_directory = args.image_path
   

    print("-> Predicting on {:d} video frames".format(len(paths)))

    # PREDICTING ON EACH IMAGE IN TURN
    with torch.no_grad():
        ims = []
        for idx, image_path in enumerate(paths):

            if image_path.endswith("_disp.jpg"):
                # don't try to predict disparity for a disparity image!
                continue

            # Load image and preprocess
            input_image = pil.open(image_path).convert('RGB')
            original_width, original_height = input_image.size
            input_image = input_image.resize((feed_width, feed_height), pil.LANCZOS)
            img_ori = np.array(input_image).copy()
            input_image = transforms.ToTensor()(input_image).unsqueeze(0)

            # PREDICTION
            input_image = input_image.to(device)
            features = depth_encoder(input_image)
            outputs = depth_decoder(features)

            disp = outputs[("disp", 0)]
            
            # Saving colormapped depth image
            disp_np = disp.squeeze().cpu().numpy()
            vmax = np.percentile(disp_np, 95)
            normalizer = mpl.colors.Normalize(vmin=disp_np.min(), vmax=vmax)
            mapper = cm.ScalarMappable(norm=normalizer, cmap='magma')
            colormapped_im = (mapper.to_rgba(disp_np)[:, :, :3] * 255).astype(np.uint8)

            out = np.concatenate([img_ori, colormapped_im], axis=0)

            out = pil.fromarray(out)
            # out = out.resize((feed_width//2, feed_height), pil.LANCZOS)
            ims.append(out)
        
        # name_dest = os.path.join(output_directory, "demo.gif")
        name_dest = "demo.gif"
        ims[0].save(name_dest, save_all=True, append_images=ims[1:], duration=150, loop=0)

        print(" Save output gif to:   {}".format(name_dest))

    print('-> Done!')


if __name__ == '__main__':
    args = parse_args()
    test_simple(args)

test_video.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np
import time
import logging

import torch
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.utils.tensorboard import SummaryWriter
from options import opts
import torch.distributed as dist
import os
import json
from utils import *
from kitti_utils import *
from layers import *


import datasets
from networks import *
# from IPython import embed


class Trainer:
    def __init__(self, options):
        self.opt = options

        self.log_path = os.path.join(self.opt.log_dir, self.opt.exp_name)
        if self.opt.global_rank == 0:
            if not os.path.exists(self.log_path):
                os.makedirs(self.log_path)
            self.save_opts()
            if not os.path.exists(os.path.join(self.log_path, 'ckpt.pth')):
                setup_logging(os.path.join(self.log_path, 'logger.log'), rank=self.opt.global_rank)
                logging.info("Experiment is named: %s", self.opt.exp_name)
                logging.info("Saving to: %s", os.path.abspath(self.log_path))
                logging.info("GPU numbers: %d", self.opt.world_size)
            else:
                setup_logging(os.path.join(self.log_path, 'logger.log'), filemode='a', rank=self.opt.global_rank)

            self.writers = {}
            for mode in ["train"]:
                self.writers[mode] = SummaryWriter(os.path.join(self.log_path, "tensorboard", mode))

        if self.opt.world_size > 1:
            dist.barrier()

        self.device = torch.device('cuda', self.opt.local_rank)

        if self.opt.seed > 0:
            self.set_seed(self.opt.seed)
        else:
            cudnn.benchmark = True

        # checking height and width are multiples of 32
        assert self.opt.height % 32 == 0, "'height' must be a multiple of 32"
        assert self.opt.width % 32 == 0, "'width' must be a multiple of 32"

        self.models = {}
        self.parameters_to_train = []

        self.num_input_frames = len(self.opt.frame_ids)
        self.num_pose_frames = 2 if self.opt.pose_model_input == "pairs" else self.num_input_frames

        self.ep_start = 0
        self.batch_start = 0
        self.step = 0

        assert self.opt.frame_ids[0] == 0, "frame_ids must start with 0"

        self.use_pose_net = not (self.opt.use_stereo and self.opt.frame_ids == [0])

        if self.opt.use_stereo:
            self.opt.frame_ids.append("s")


        # data
        datasets_dict = {"kitti": datasets.KITTIRAWDataset,
                         "kitti_odom": datasets.KITTIOdomDataset,
                         "cityscapes": datasets.CityscapesDataset,
                         "nyuv2": datasets.NYUDataset}
        self.dataset = datasets_dict[self.opt.dataset]

        if self.opt.dataset == "kitti":
            fpath = os.path.join(os.path.dirname(__file__), "splits/kitti", self.opt.split, "{}_files.txt")
            fpath_test = os.path.join(os.path.dirname(__file__), "splits/kitti", self.opt.eval_split, "{}_files.txt")
        elif self.opt.dataset == "kitti_odom":
            fpath = os.path.join(os.path.dirname(__file__), "splits/kitti", "odom", "{}_files.txt")
            fpath_test = os.path.join(os.path.dirname(__file__), "splits/kitti", "odom", "{}_files_09.txt")
        elif self.opt.dataset == "nyuv2":
            fpath = os.path.join(os.path.dirname(__file__), "splits/nyuv2", "{}_files.txt")
            fpath_test = os.path.join(os.path.dirname(__file__), "splits/nyuv2", "{}_files.txt")
        elif self.opt.dataset == "cityscapes":
            fpath = os.path.join(os.path.dirname(__file__), "splits/cityscapes", "{}_files.txt")
            fpath_test = os.path.join(os.path.dirname(__file__), "splits/cityscapes", "{}_files.txt")            
        else:
            pass


        train_filenames = readlines(fpath.format("train"))
        test_filenames = readlines(fpath_test.format("test"))
        img_ext = '.jpg' if self.opt.jpg else '.png'

        num_train_samples = len(train_filenames)
        self.num_steps_per_epoch = num_train_samples // self.opt.world_size // self.opt.batch_size
        self.num_total_steps = self.num_steps_per_epoch * self.opt.num_epochs

        if self.opt.dataset == "cityscapes":
            train_dataset = self.dataset(
                self.opt.data_path_pre, train_filenames, self.opt.height, self.opt.width, self.opt.frame_ids, self.opt.num_scales, local_crop=self.opt.use_local_crop, patch_reshuffle=self.opt.use_patch_reshuffle, is_train=True, img_ext=img_ext)
        else:
            train_dataset = self.dataset(
                self.opt.data_path, train_filenames, self.opt.height, self.opt.width, self.opt.frame_ids, self.opt.num_scales, local_crop=self.opt.use_local_crop, patch_reshuffle=self.opt.use_patch_reshuffle, is_train=True, img_ext=img_ext)
        if self.opt.world_size > 1:
            self.sampler = datasets.CustomDistributedSampler(train_dataset, self.opt.seed)
        else:
            self.sampler = datasets.CustomSampler(train_dataset, self.opt.seed)
        self.train_loader = DataLoader(
            train_dataset, self.opt.batch_size, shuffle=False, sampler=self.sampler, num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)

        # for testing the model at the end of each epoch
        test_dataset = self.dataset(
            self.opt.data_path, test_filenames, self.opt.height, self.opt.width,
            [0], self.opt.num_scales, is_train=False, img_ext=img_ext)
        self.test_loader = DataLoader(
            test_dataset, self.opt.batch_size, shuffle=False,
            num_workers=self.opt.num_workers, pin_memory=True, drop_last=False)
        if self.opt.dataset == "kitti":
            gt_path = os.path.join(os.path.dirname(__file__), "splits/kitti", self.opt.eval_split, "gt_depths.npz")
            self.gt_depths = np.load(gt_path, fix_imports=True, encoding='latin1', allow_pickle=True)["data"]
        elif self.opt.dataset == "cityscapes":
            gt_path = os.path.join(os.path.dirname(__file__), "splits", "cityscapes", "gt_depths")
            self.gt_depths = []
            for i in range(len(test_dataset)):
                gt_depth = np.load(os.path.join(gt_path, str(i).zfill(3) + '_depth.npy'))
                self.gt_depths.append(gt_depth)
        else:
            pass

        # create models
        if self.opt.backbone == "resnet":
            self.models["encoder"] = depth_resnet.DepthEncoder(
                self.opt.num_layers, self.opt.weights_init == "pretrained")
            self.models["depth"] = depth_resnet.DepthDecoder(
                self.models["encoder"].num_ch_enc, range(self.opt.num_scales))
        if self.opt.backbone == "hrnet":
            self.models["encoder"] = depth_hrnet.DepthEncoder(
                self.opt.num_layers, self.opt.weights_init == "pretrained")
            self.models["depth"] = depth_hrnet.DepthDecoder(
                self.models["encoder"].num_ch_enc, range(self.opt.num_scales))
                
        self.models["encoder"].to(self.device)
        self.models["depth"].to(self.device)


        if self.use_pose_net:
            self.models["pose_encoder"] = posenet.ResnetEncoder(
                self.opt.num_layers,
                self.opt.weights_init == "pretrained",
                num_input_images=self.num_pose_frames)

            self.models["pose_encoder"].to(self.device)

            self.models["pose"] = posenet.PoseDecoder(
                self.models["pose_encoder"].num_ch_enc,
                num_input_features=1,
                num_frames_to_predict_for=2)

            self.models["pose"].to(self.device)
 

        if self.opt.predictive_mask:
            assert self.opt.disable_automasking, \
                "When using predictive_mask, please disable automasking with --disable_automasking"

            # Our implementation of the predictive masking baseline has the the same architecture
            # as our depth decoder. We predict a separate mask for each source frame.
            self.models["predictive_mask"] = networks.DepthDecoder(
                self.models["encoder"].num_ch_enc, range(self.opt.num_scales),
                num_output_channels=(len(self.opt.frame_ids) - 1))
            self.models["predictive_mask"].to(self.device)

        for k in self.models.keys():
            self.parameters_to_train += list(self.models[k].parameters())

        if self.opt.pretrained_path:
            if not self.opt.resume:
                self.load_pretrained_model()
            elif not os.path.exists(os.path.join(self.log_path, 'ckpt.pth')):
                self.load_pretrained_model()

        if self.opt.resume:
            checkpoint = self.load_ckpt()

        if self.opt.world_size > 1:
            for k in self.models.keys():
                self.models[k] = nn.SyncBatchNorm.convert_sync_batchnorm(
                self.models[k])
                self.models[k] = nn.parallel.DistributedDataParallel(self.models[k], device_ids=[self.opt.local_rank], output_device=self.opt.local_rank, find_unused_parameters=True)

        # optimizer settings
        if self.opt.optimizer == 'adamw':
            self.model_optimizer = torch.optim.AdamW(self.parameters_to_train,lr=self.opt.learning_rate, betas=(self.opt.beta1, self.opt.beta2),weight_decay=self.opt.weight_decay)
        elif self.opt.optimizer == 'adam':
            self.model_optimizer = torch.optim.Adam(self.parameters_to_train,lr=self.opt.learning_rate, betas=(self.opt.beta1, self.opt.beta2)) 
        elif self.opt.optimizer == 'sgd':
            self.model_optimizer = torch.optim.SGD(self.parameters_to_train,lr=self.opt.learning_rate, momentum=self.opt.momentum)
        else:
            logging.error("Optimizer '%s' not defined. Use (adamw|adam|sgd) instead", self.opt.optimizer)


        if self.opt.lr_sche_type == 'cos':
            self.model_lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.model_optimizer, T_max=self.num_total_steps, eta_min=self.opt.eta_min)
        elif self.opt.lr_sche_type == 'step':
            self.model_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(self.model_optimizer, self.opt.decay_step, self.opt.decay_rate)

        if checkpoint:
            self.model_optimizer.load_state_dict(checkpoint["optimizer"])
            self.model_lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
            del checkpoint

        if not self.opt.no_ssim:
            self.ssim = SSIM()
            self.ssim.to(self.device)

        self.backproject_depth = {}
        self.project_3d = {}
        for scale in range(self.opt.num_scales):
            h = self.opt.height // (2 ** scale)
            w = self.opt.width // (2 ** scale)

            self.backproject_depth[scale] = BackprojectDepth(self.opt.batch_size, h, w)
            self.backproject_depth[scale].to(self.device)

            self.project_3d[scale] = Project3D(self.opt.batch_size, h, w)
            self.project_3d[scale].to(self.device)

        self.depth_metric_names = [
            "de/abs_rel", "de/sq_rel", "de/rms", "de/log_rms", "da/a1", "da/a2", "da/a3"]

        logging.info("Using split: %s", self.opt.split)
        logging.info("There are {:d} training items and {:d} test items\n".format(len(train_dataset), len(test_dataset)))

        if self.opt.world_size > 1:
            dist.barrier()


    def set_seed(self, seed=1234):
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    def set_train(self):
        """Convert all models to training mode
        """
        for m in self.models.values():
            m.train()

    def set_eval(self):
        """Convert all models to testing/evaluation mode
        """
        for m in self.models.values():
            m.eval()


    def test_nyuv2(self):
        logging.info(" ")
        logging.info("Test the model at epoch {} \n".format(self.epoch))

        self.set_eval()
        pred_depths = []
        gt_depths = []

        for idx, (color, depth) in enumerate(self.test_loader):
            if self.opt.global_rank == 0:
                print("{}/{}".format(idx+1, len(self.test_loader)), end='\r')
            input_color = color.to(self.device)

            output = self.models["depth"](self.models["encoder"](input_color))
            pred_disp, _ = disp_to_depth(output[("disp", 0)], self.opt.min_depth, self.opt.max_depth)
            pred_disp = pred_disp[:, 0]
            
            gt_depth = depth
            _, h, w = gt_depth.shape
            pred_depth = 1 / pred_disp
            pred_depth = F.interpolate(pred_depth.unsqueeze(0), (h, w), mode="nearest")[0]
            pred_depths.append(pred_depth)
            gt_depths.append(gt_depth)
        pred_depths = torch.cat(pred_depths, dim=0)
        gt_depths = torch.cat(gt_depths, dim=0).to(self.device)

        errors = []
        ratios = []
        for i in range(pred_depths.shape[0]):    
            pred_depth = pred_depths[i]
            gt_depth = gt_depths[i]
            mask = (gt_depth > 0) & (gt_depth < 10)
            pred_depth = pred_depth[mask]
            gt_depth = gt_depth[mask]

            ratio = torch.median(gt_depth) / torch.median(pred_depth)
            ratios.append(ratio)
            pred_depth *= ratio         
            pred_depth[pred_depth > 10] = 10
            errors.append(compute_depth_errors(gt_depth, pred_depth))

        ratios = torch.tensor(ratios)
        med = torch.median(ratios)
        std = torch.std(ratios / med)

        logging.info(" Mono evaluation - using median scaling")
        logging.info(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

        mean_errors = torch.tensor(errors).mean(0)

        logging.info(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
        logging.info(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))
        self.set_train()


    def test_cityscapes(self):
        logging.info(" ")
        logging.info("Test the model at epoch {} \n".format(self.epoch))
        MIN_DEPTH = 1e-3
        MAX_DEPTH = 80
        STEREO_SCALE_FACTOR = 5.4
        self.set_eval()       

        pred_disps = []
        for idx, data in enumerate(self.test_loader):
            if self.opt.global_rank == 0:
                print("{}/{}".format(idx+1, len(self.test_loader)), end='\r')
            input_color = data[("color", 0, 0)].to(self.device)
            output = self.models["depth"](self.models["encoder"](input_color))
            pred_disp, _ = disp_to_depth(output[("disp", 0)], self.opt.min_depth, self.opt.max_depth)
            pred_disps.append(pred_disp[:, 0])
        pred_disps = torch.cat(pred_disps, dim=0)

        errors = []
        ratios = []
        for i in range(pred_disps.shape[0]):
            gt_depth = torch.from_numpy(self.gt_depths[i]).cuda()
            gt_height, gt_width = gt_depth.shape[:2]

            # crop ground truth to remove ego car -> this has happened in the dataloader for inputs
            gt_height = int(round(gt_height * 0.75))
            gt_depth = gt_depth[:gt_height]
            pred_disp = pred_disps[i:i+1].unsqueeze(0)
            pred_disp = F.interpolate(pred_disp, (gt_height, gt_width), mode="bilinear", align_corners=True)
            pred_depth = 1 / pred_disp[0, 0, :]

            # when evaluating cityscapes, we centre crop to the middle 50% of the image.
            # Bottom 25% has already been removed - so crop the sides and the top here
            gt_depth = gt_depth[256:, 192:1856]
            pred_depth = pred_depth[256:, 192:1856]

            mask = (gt_depth > MIN_DEPTH) & (gt_depth < MAX_DEPTH)
            pred_depth = pred_depth[mask]
            gt_depth = gt_depth[mask]

            if self.opt.use_stereo:
                pred_depth *= STEREO_SCALE_FACTOR
            else:
                ratio = torch.median(gt_depth) / torch.median(pred_depth)
                ratios.append(ratio)
                pred_depth *= ratio  
            pred_depth = torch.clamp(pred_depth, MIN_DEPTH, MAX_DEPTH)
            errors.append(compute_depth_errors(gt_depth, pred_depth))

        if self.opt.use_stereo:
            logging.info(" Stereo evaluation - disabling median scaling")
            logging.info(" Scaling by {}".format(STEREO_SCALE_FACTOR))
        else:
            ratios = torch.tensor(ratios)
            med = torch.median(ratios)
            std = torch.std(ratios / med)
            logging.info(" Mono evaluation - using median scaling")
            logging.info(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

        mean_errors = torch.tensor(errors).mean(0)

        logging.info(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
        logging.info(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))
        self.set_train()

    def test_kitti(self):
        """Test the model on a single minibatch
        """
        logging.info(" ")
        logging.info("Test the model at epoch {} \n".format(self.epoch))

        MIN_DEPTH = 1e-3
        MAX_DEPTH = 80
        # Models which were trained with stereo supervision were trained with a nominal baseline of 0.1 units. The KITTI rig has a baseline of 54cm. Therefore, to convert our stereo predictions to real-world scale we multiply our depths by 5.4.
        STEREO_SCALE_FACTOR = 5.4
        self.set_eval()

        pred_disps = []
        for idx, data in enumerate(self.test_loader):
            if self.opt.global_rank == 0:
                print("{}/{}".format(idx+1, len(self.test_loader)), end='\r')
            input_color = data[("color", 0, 0)].to(self.device)
            output = self.models["depth"](self.models["encoder"](input_color))
            pred_disp, _ = disp_to_depth(output[("disp", 0)], self.opt.min_depth, self.opt.max_depth)
            pred_disps.append(pred_disp[:, 0])
        pred_disps = torch.cat(pred_disps, dim=0)

        errors = []
        ratios = []
        for i in range(pred_disps.shape[0]):
            gt_depth = torch.from_numpy(self.gt_depths[i]).cuda()
            gt_height, gt_width = gt_depth.shape[:2]
            pred_disp = pred_disps[i:i+1].unsqueeze(0)
            pred_disp = F.interpolate(pred_disp, (gt_height, gt_width), mode="bilinear", align_corners=False)
            pred_depth = 1 / pred_disp[0, 0, :]
            if self.opt.eval_split == "eigen":
                mask = (gt_depth > MIN_DEPTH) & (gt_depth < MAX_DEPTH)
                crop_mask = torch.zeros_like(mask)
                crop_mask[
                        int(0.40810811 * gt_height):int(0.99189189 * gt_height),
                        int(0.03594771 * gt_width):int(0.96405229 * gt_width)] = 1
                mask = mask * crop_mask
            else:
                mask = gt_depth > 0

            pred_depth = pred_depth[mask]
            gt_depth = gt_depth[mask]
            if self.opt.use_stereo:
                pred_depth *= STEREO_SCALE_FACTOR
            else:
                ratio = torch.median(gt_depth) / torch.median(pred_depth)
                ratios.append(ratio)
                pred_depth *= ratio  
            pred_depth = torch.clamp(pred_depth, MIN_DEPTH, MAX_DEPTH)
            errors.append(compute_depth_errors(gt_depth, pred_depth))

        if self.opt.use_stereo:
            logging.info(" Stereo evaluation - disabling median scaling")
            logging.info(" Scaling by {}".format(STEREO_SCALE_FACTOR))
        else:
            ratios = torch.tensor(ratios)
            med = torch.median(ratios)
            std = torch.std(ratios / med)
            logging.info(" Mono evaluation - using median scaling")
            logging.info(" Scaling ratios | med: {:0.3f} | std: {:0.3f}".format(med, std))

        mean_errors = torch.tensor(errors).mean(0)

        logging.info(("{:>8} | " * 7).format("abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"))
        logging.info(("{: 8.3f} | " * 7 + "\n").format(*mean_errors.tolist()))
        self.set_train()


    def train(self):
        """Run the entire training pipeline
        """
        for self.epoch in range(self.ep_start, self.opt.num_epochs):
            self.run_epoch()
            if self.opt.lr_sche_type == "step":
                self.model_lr_scheduler.step()
            with torch.no_grad():
                if self.opt.dataset == "kitti":
                    self.test_kitti()
                elif self.opt.dataset == "nyuv2":
                    self.test_nyuv2()
                elif self.opt.dataset == "cityscapes":
                    self.test_cityscapes()
                else:
                    pass
            if self.opt.global_rank == 0:
                self.save_model(ep_end=True)


    def run_epoch(self):
        """Run a single epoch of training and validation
        """
        logging.info("Training epoch {}\n".format(self.epoch))

        self.sampler.set_epoch(self.epoch)
        self.sampler.set_start_iter(self.batch_start*self.opt.batch_size)
        self.set_train()

        if self.opt.world_size > 1:
            dist.barrier()

        start_data_time = time.time()
        for batch_idx, inputs in enumerate(self.train_loader):
            self.step += 1
            start_fp_time = time.time()
            outputs, losses = self.process_batch(inputs)

            start_bp_time = time.time()
            self.model_optimizer.zero_grad()

            losses["loss"].backward()
            if self.opt.clip_grad != -1:
                for params in self.model_optimizer.param_groups:
                    params = params['params']
                    nn.utils.clip_grad_norm_(params, max_norm=self.opt.clip_grad)

            self.model_optimizer.step()

            if self.opt.lr_sche_type == "cos":
                self.model_lr_scheduler.step()

            # compute the process time
            data_time = start_fp_time - start_data_time
            fp_time = start_bp_time - start_fp_time
            bp_time = time.time() - start_bp_time

            # logging
            if ((batch_idx+self.batch_start) % self.opt.log_frequency == 0):
                if "depth_gt" in inputs:
                    self.compute_depth_losses(inputs, outputs, losses)
                if self.opt.world_size > 1:
                    dist.barrier()
                    for k in losses.keys():
                        dist.all_reduce(losses[k], op=dist.ReduceOp.SUM)
                        losses[k] = losses[k] / self.opt.world_size
                    dist.barrier()
                if self.opt.global_rank == 0:
                    self.log_time(batch_idx+self.batch_start, data_time, fp_time, bp_time, losses)
                    self.log_tensorboard("train", inputs, outputs, losses)

            # save ckpt
            if ((batch_idx+self.batch_start)>0 and (batch_idx+self.batch_start) % self.opt.save_frequency == 0):
                if self.opt.global_rank == 0:
                    self.save_model(batch_idx=batch_idx+self.batch_start+1)
            if self.opt.world_size > 1:
                dist.barrier()
            start_data_time = time.time()

        self.batch_start = 0

    def process_batch(self, inputs):
        """Pass a minibatch through the network and generate images and losses
        """
        for key, ipt in inputs.items():
            try:
                inputs[key] = ipt.to(self.device)
            except:
                pass

        features = self.models["encoder"](inputs[("color_aug", 0, 0)])
        outputs = self.models["depth"](features)

        if self.opt.use_local_crop:
            x = self.models["encoder"](inputs[("color_local_aug", 0, 0)])
            o = self.models["depth"](x)
            for i in range(self.opt.num_scales):
                outputs[("disp_local", i)] = o[("disp", i)]
        if self.opt.use_patch_reshuffle:
            x = self.models["encoder"](inputs[("color_reshuffle_aug", 0, 0)])
            o = self.models["depth"](x)
            for i in range(self.opt.num_scales):
                outputs[("disp_reshuffle", i)] = o[("disp", i)]

                # restore the reshuffled disparity maps
                all_disp = []
                for b in range(self.opt.batch_size):
                    ### Split-Permute as depicted in paper (vertical + horizontal)
                    split_x = inputs[("split_xy")][b][0].item()
                    split_y = inputs[("split_xy")][b][1].item()
                    split_x = round(split_x / (2 ** i))
                    split_y = round(split_y / (2 ** i))
                    disp_reshuffle = outputs[("disp_reshuffle", i)][b]   #1*H*W
                    patch1 = disp_reshuffle[:, 0:split_y, :]
                    patch2 = disp_reshuffle[:, split_y:, :]
                    disp_restore = torch.cat([patch2, patch1], dim=1)
                    patch1 = disp_restore[:, :, 0:split_x]
                    patch2 = disp_restore[:, :, split_x:]
                    disp_restore = torch.cat([patch2, patch1], dim=2)
                    all_disp.append(disp_restore)

                    ### Split-Permute (vertical or horizontal, randomly choose one)
                    # split_x = inputs[("split_xy", i)][b][0].item()
                    # split_y = inputs[("split_xy", i)][b][1].item()
                    # split_x = round(split_x / (2 ** i))
                    # split_y = round(split_y / (2 ** i))
                    # disp_reshuffle = outputs[("disp_reshuffle", i)][b]   #1*H*W
                    # if split_x == 0:
                    #     patch1 = disp_reshuffle[:, 0:split_y, :]
                    #     patch2 = disp_reshuffle[:, split_y:, :]
                    #     disp_restore = torch.cat([patch2, patch1], dim=1)
                    # else:
                    #     patch1 = disp_reshuffle[:, :, 0:split_x]
                    #     patch2 = disp_reshuffle[:, :, split_x:]
                    #     disp_restore = torch.cat([patch2, patch1], dim=2)
                    # all_disp.append(disp_restore)
                disp_restore = torch.stack(all_disp, dim=0)
                outputs[("disp_reshuffle", i)] = disp_restore

        if self.opt.predictive_mask:
            outputs["predictive_mask"] = self.models["predictive_mask"](features)

        if self.use_pose_net:
            outputs.update(self.predict_poses(inputs, features))

        self.generate_images_pred(inputs, outputs)
        losses = self.compute_losses(inputs, outputs)

        return outputs, losses

    def predict_poses(self, inputs, features):
        """Predict poses between input frames for monocular sequences.
        """
        outputs = {}
        if self.num_pose_frames == 2:
            # In this setting, we compute the pose to each source frame via a
            # separate forward pass through the pose network.

            pose_feats = {f_i: inputs["color_aug", f_i, 0] for f_i in self.opt.frame_ids}

            for f_i in self.opt.frame_ids[1:]:
                if f_i != "s":
                    # To maintain ordering we always pass frames in temporal order
                    if f_i < 0:
                        pose_inputs = [pose_feats[f_i], pose_feats[0]]
                    else:
                        pose_inputs = [pose_feats[0], pose_feats[f_i]]

                    pose_inputs = [self.models["pose_encoder"](torch.cat(pose_inputs, 1))]

                    axisangle, translation = self.models["pose"](pose_inputs)
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation

                    # Invert the matrix if the frame id is negative
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i < 0))

        else:
            # Here we input all frames to the pose net (and predict all poses) together
            pose_inputs = torch.cat(
                [inputs[("color_aug", i, 0)] for i in self.opt.frame_ids if i != "s"], 1)

            pose_inputs = [self.models["pose_encoder"](pose_inputs)]

            axisangle, translation = self.models["pose"](pose_inputs)

            for i, f_i in enumerate(self.opt.frame_ids[1:]):
                if f_i != "s":
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, i], translation[:, i])

        return outputs

      
    def generate_images_pred(self, inputs, outputs):
        """Generate the warped (reprojected) color images for a minibatch.
        Generated images are saved into the `outputs` dictionary.
        """
        for scale in range(self.opt.num_scales):
            # original
            disp = outputs[("disp", scale)]
            disp = F.interpolate(disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            source_scale = 0

            _, depth = disp_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, scale)] = depth

            for i, frame_id in enumerate(self.opt.frame_ids[1:]):
                if frame_id == "s":
                    T = inputs["stereo_T"]
                else:
                    T = outputs[("cam_T_cam", 0, frame_id)]
                cam_points = self.backproject_depth[source_scale](
                    depth, inputs[("inv_K", source_scale)])
                pix_coords = self.project_3d[source_scale](
                    cam_points, inputs[("K", source_scale)], T)
                outputs[("sample", frame_id, scale)] = pix_coords
                outputs[("color", frame_id, scale)] = F.grid_sample(
                    inputs[("color", frame_id, source_scale)],
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border", align_corners=True)

            # patch_reshuffled
            if self.opt.use_patch_reshuffle:
                disp = outputs[("disp_reshuffle", scale)]
                disp = F.interpolate(disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
                _, depth = disp_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
                for _, frame_id in enumerate(self.opt.frame_ids[1:]):
                    if frame_id == "s":
                        T = inputs["stereo_T"]
                    else:
                        T = outputs[("cam_T_cam", 0, frame_id)]
                    cam_points = self.backproject_depth[source_scale](
                        depth, inputs[("inv_K", source_scale)])
                    pix_coords = self.project_3d[source_scale](
                        cam_points, inputs[("K", source_scale)], T)
                    outputs[("color_reshuffle", frame_id, scale)] = F.grid_sample(
                        inputs[("color", frame_id, source_scale)].clone(),
                        pix_coords,
                        padding_mode="border", align_corners=True)  

            # local crop
            if self.opt.use_local_crop:
                disp = outputs[("disp_local", scale)]
                disp = F.interpolate(disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
                _, depth = disp_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
                for _, frame_id in enumerate(self.opt.frame_ids[1:]):
                    if frame_id == "s":
                        T = inputs["stereo_T"]
                    else: 
                        T = outputs[("cam_T_cam", 0, frame_id)]
                    Rt_Rc = torch.zeros_like(T).to(self.device)
                    gx0 = (inputs[("grid_local")][:, 0, 0, -1] + inputs[("grid_local")][:, 0, 0, 0]) / 2.
                    gy0 = (inputs[("grid_local")][:, 1, -1, 0] + inputs[("grid_local")][:, 1, 0, 0]) / 2.
                    f = (inputs[("grid_local")][:, 0, 0, -1] - inputs[("grid_local")][:, 0, 0, 0]) / 2.
                    fx = inputs[("K", 0)][0, 0, 0] / self.opt.width
                    fy = inputs[("K", 0)][0, 1, 1] / self.opt.height
                    Rc_v = torch.stack([-gx0/(2*fx), -gy0/(2*fy), f], dim=1)
                    Rc = torch.eye(3).to(self.device)
                    Rc = Rc[None, :, :].repeat(Rc_v.shape[0], 1, 1)
                    Rc[:, :, 2] = Rc_v
                    # outputs[("Rc", f_i)] = Rc
                    Rt_Rc[:, :3, :3] = torch.matmul(Rc, torch.matmul(T[:, :3, :3], torch.inverse(Rc)))
                    Rt_Rc[:, :3, 3:4] = torch.matmul(Rc, T[:, :3, 3:4])                    
                    T = Rt_Rc
                    cam_points = self.backproject_depth[source_scale](
                        depth, inputs[("inv_K", source_scale)])
                    pix_coords = self.project_3d[source_scale](
                        cam_points, inputs[("K", source_scale)], T)
                    outputs[("color_local", frame_id, scale)] = F.grid_sample(
                        inputs[("color_local", frame_id, source_scale)],
                        pix_coords, padding_mode="border",align_corners=True)                        


    def compute_reprojection_loss(self, pred, target):
        """Computes reprojection loss between a batch of predicted and target images
        """
        abs_diff = torch.abs(target - pred)
        l1_loss = abs_diff.mean(1, True)

        if self.opt.no_ssim:
            reprojection_loss = l1_loss
        else:
            ssim_loss = self.ssim(pred, target).mean(1, True)
            reprojection_loss = 0.85 * ssim_loss + 0.15 * l1_loss

        return reprojection_loss


    def compute_SI_log_depth_loss(self, pred, target, mask=None, lamda=0.5):
        # B*1*H*W  ->  B*H*W
        if mask is None:
            mask = torch.ones_like(pred).to(self.device)

        mask = mask[:, 0] 
        log_pred = torch.log(pred[:, 0]+1e-8) * mask
        log_tgt = torch.log(target[:, 0]+1e-8) * mask
     
        log_diff = log_pred - log_tgt
        valid_num = mask.sum(1).sum(1) + 1e-8
        log_diff_squre_sum = (log_diff ** 2).sum(1).sum(1)
        log_diff_sum_squre = (log_diff.sum(1).sum(1)) ** 2
        loss = log_diff_squre_sum/valid_num - lamda*log_diff_sum_squre/(valid_num**2)

        return loss.mean()


    def compute_losses_base(self, inputs, outputs, ty=""):
        loss_all = 0
        for scale in range(self.opt.num_scales):
            loss = 0
            reprojection_losses = []
            source_scale = 0

            if ty == "_local":
                disp = outputs[("disp"+ty, scale)]
                color = inputs[("color"+ty, 0, scale)]
                target = inputs[("color"+ty, 0, source_scale)]
            else:
                disp = outputs[("disp"+ty, scale)]
                color = inputs[("color", 0, scale)]
                target = inputs[("color", 0, source_scale)]                

            for frame_id in self.opt.frame_ids[1:]:
                pred = outputs[("color"+ty, frame_id, scale)]
                reprojection_losses.append(self.compute_reprojection_loss(pred, target))

            reprojection_losses = torch.cat(reprojection_losses, 1)
            
            if not self.opt.disable_automasking:
                identity_reprojection_losses = []
                for frame_id in self.opt.frame_ids[1:]:
                    if ty == "_local":
                        pred = inputs[("color"+ty, frame_id, source_scale)]
                    else:
                        pred = inputs[("color", frame_id, source_scale)]
                    identity_reprojection_losses.append(
                        self.compute_reprojection_loss(pred, target))

                identity_reprojection_losses = torch.cat(identity_reprojection_losses, 1)

                if self.opt.avg_reprojection:
                    identity_reprojection_loss = identity_reprojection_losses.mean(1, keepdim=True)
                else:
                    # save both images, and do min all at once below
                    identity_reprojection_loss = identity_reprojection_losses

            elif self.opt.predictive_mask:
                # use the predicted mask
                mask = outputs["predictive_mask"]["disp", scale]
                if not self.opt.v1_multiscale:
                    mask = F.interpolate(
                        mask, [self.opt.height, self.opt.width],
                        mode="bilinear", align_corners=False)

                reprojection_losses *= mask

                # add a loss pushing mask to 1 (using nn.BCELoss for stability)
                weighting_loss = 0.2 * nn.BCELoss()(mask, torch.ones(mask.shape).cuda())
                loss += weighting_loss.mean()

            if self.opt.avg_reprojection:
                reprojection_loss = reprojection_losses.mean(1, keepdim=True)
            else:
                reprojection_loss = reprojection_losses

            if not self.opt.disable_automasking:
                # add random numbers to break ties
                identity_reprojection_loss += torch.randn(
                    identity_reprojection_loss.shape, device=self.device) * 0.00001

                combined = torch.cat((identity_reprojection_loss, reprojection_loss), dim=1)
            else:
                combined = reprojection_loss

            if combined.shape[1] == 1:
                to_optimise = combined
            else:
                to_optimise, idxs = torch.min(combined, dim=1)

            if not self.opt.disable_automasking:
                outputs["identity_selection"+ty+"/{}".format(scale)] = (
                    idxs > identity_reprojection_loss.shape[1] - 1).float()

            loss += to_optimise.mean()
            mean_disp = disp.mean(2, True).mean(3, True)
            norm_disp = disp / (mean_disp + 1e-7)

            smooth_loss = get_smooth_loss(norm_disp, color)

            loss += self.opt.disparity_smoothness * smooth_loss / (2 ** scale)     
            loss_all += loss
        loss_all /= self.opt.num_scales

        return loss_all


    def compute_losses(self, inputs, outputs):
        """Compute the reprojection and smoothness losses for a minibatch
        """
        losses = {}

        loss_base = 0
        n = 0
        loss_base_ori = self.compute_losses_base(inputs, outputs, ty="")
        losses["loss_base_ori"] = loss_base_ori
        loss_base += loss_base_ori
        if self.opt.use_local_crop:
            loss_base_local = self.compute_losses_base(inputs, outputs, ty="_local")
            losses["loss_base_local"] = loss_base_local
            loss_base += loss_base_local
            n += 1
        if self.opt.use_patch_reshuffle:
            loss_base_reshuffle = self.compute_losses_base(inputs, outputs, ty="_reshuffle")
            losses["loss_base_reshuffle"] = loss_base_reshuffle
            loss_base += loss_base_reshuffle
            n += 1
        loss_base /= n+1
        losses["loss_base"] = loss_base


        # depth consistency loss
        loss_dc = torch.tensor(0.0).to(self.device)
        if self.opt.use_local_crop:
            loss_dc_local = 0
            for i in range(self.opt.num_scales):
                disp = outputs[("disp", i)]
                disp = F.interpolate(disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
                loss_dc_i = 0
                for b in range(self.opt.batch_size):
                    disp_local = outputs[("disp_local", i)][b].clone()
                    x0 = round(self.opt.width * (inputs[("grid_local")][b, 0, 0, 0].item() - (-1)) / 2.)
                    y0 = round(self.opt.height * (inputs[("grid_local")][b, 1, 0, 0].item() - (-1)) / 2.)        
                    w = round(self.opt.width/inputs[("ratio_local")][b, 0].item())
                    h = round(self.opt.height/inputs[("ratio_local")][b, 0].item())
                    disp_local = F.interpolate(disp_local.unsqueeze(0), [h, w], mode="bilinear", align_corners=False)
                    _, depth_local = disp_to_depth(disp_local, self.opt.min_depth, self.opt.max_depth)
                    depth_local *= inputs[("ratio_local")][b, 0]
                    _, depth_from_ori = disp_to_depth(disp[b, :, y0:y0+h, x0:x0+w].clone().unsqueeze(0), self.opt.min_depth, self.opt.max_depth)
                    
                    loss_dc_i += self.compute_SI_log_depth_loss(depth_local, depth_from_ori)
                loss_dc_i /= self.opt.batch_size
                loss_dc_local += loss_dc_i
            loss_dc_local /= self.opt.num_scales
            losses["loss_dc_local"] = loss_dc_local
            loss_dc += loss_dc_local

        if self.opt.use_patch_reshuffle:
            loss_dc_reshuffle = 0
            for i in range(self.opt.num_scales):
                _, depth = disp_to_depth(outputs[("disp", i)].clone(), self.opt.min_depth, self.opt.max_depth)
                disp_restore = outputs[("disp_reshuffle", i)]
                _, depth_restore = disp_to_depth(disp_restore, self.opt.min_depth, self.opt.max_depth)
                loss_dc_reshuffle += self.compute_SI_log_depth_loss(depth_restore, depth)   
            loss_dc_reshuffle /= self.opt.num_scales
            losses["loss_dc_reshuffle"] = loss_dc_reshuffle     
            loss_dc += loss_dc_reshuffle

        losses["loss_dc"] = loss_dc
        total_loss = loss_base + 0.07 * loss_dc
        # total_loss = loss_base
        losses["loss"] = total_loss

        return losses

    def compute_depth_losses(self, inputs, outputs, losses):
        """Compute depth metrics, to allow monitoring during training

        This isn't particularly accurate as it averages over the entire batch,
        so is only used to give an indication of validation performance
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = torch.clamp(F.interpolate(
            depth_pred, [375, 1242], mode="bilinear", align_corners=False), 1e-3, 80)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = depth_gt > 0

        # garg/eigen crop
        crop_mask = torch.zeros_like(mask)
        crop_mask[:, :, 153:371, 44:1197] = 1
        mask = mask * crop_mask

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)

        depth_pred = torch.clamp(depth_pred, min=1e-3, max=80)

        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
 
    def log_time(self, batch_idx, data_time, fp_time, bp_time, losses):
        """Print a logging statement to the terminal
        """
        batch_time = data_time + fp_time + bp_time
        # time_sofar = time.time() - self.start_time
        training_time_left = (self.num_total_steps - self.step) * batch_time if self.step > 1 else 0
        print_string = "epoch: {:>2}/{} | batch: {:>4}/{} | data time: {:.4f}" + " | batch time: {:.3f} | loss_base: {:.4f} | loss_dc: {:.4f} | lr: {:.2e} | time left: {}"
        logging.info(print_string.format(self.epoch, self.opt.num_epochs-1,batch_idx, self.num_steps_per_epoch, data_time, batch_time, losses["loss_base"].cpu().data, losses["loss_dc"].cpu().data, self.model_optimizer.state_dict()['param_groups'][0]['lr'], sec_to_hm_str(training_time_left)))

    def log_tensorboard(self, mode, inputs, outputs, losses):
        """Write an event to the tensorboard events file
        """
        writer = self.writers[mode]
        for l, v in losses.items():
            writer.add_scalar("{}".format(l), v, self.step)

        for j in range(min(3, self.opt.batch_size)):  # write a maxmimum of four images
            for s in range(self.opt.num_scales):
                for frame_id in self.opt.frame_ids:
                    writer.add_image(
                        "color_{}_{}/{}".format(frame_id, s, j),
                        inputs[("color", frame_id, s)][j].data, self.step)
                    if s == 0 and frame_id != 0:
                        writer.add_image(
                            "color_pred_{}_{}/{}".format(frame_id, s, j),
                            outputs[("color", frame_id, s)][j].data, self.step)

                writer.add_image(
                    "disp_{}/{}".format(s, j),
                    normalize_image(outputs[("disp", s)][j]), self.step)

                if self.opt.predictive_mask:
                    for f_idx, frame_id in enumerate(self.opt.frame_ids[1:]):
                        writer.add_image(
                            "predictive_mask_{}_{}/{}".format(frame_id, s, j),
                            outputs["predictive_mask"][("disp", s)][j, f_idx][None, ...],
                            self.step)

                elif not self.opt.disable_automasking:
                    writer.add_image(
                        "automask_{}/{}".format(s, j),
                        outputs["identity_selection/{}".format(s)][j][None, ...], self.step)

                if self.opt.use_local_crop:
                    for frame_id in self.opt.frame_ids:
                        writer.add_image(
                            "color_local_{}_{}/{}".format(frame_id, s, j),
                            inputs[("color_local", frame_id, s)][j].data, self.step)
                        if s == 0 and frame_id != 0:
                            writer.add_image(
                                "color_pred_local_{}_{}/{}".format(frame_id, s, j),
                                outputs[("color_local", frame_id, s)][j].data, self.step)

                    writer.add_image(
                        "disp_local_{}/{}".format(s, j),
                        normalize_image(outputs[("disp_local", s)][j]), self.step)

                    # writer.add_image(
                    #     "automask_local_{}/{}".format(s, j),
                    #     outputs["identity_selection_local/{}".format(s)][j][None, ...], self.step)

                if self.opt.use_patch_reshuffle:
                    writer.add_image(
                        "disp_reshuffle_restore_{}/{}".format(s, j),
                        normalize_image(outputs[("disp_reshuffle", s)][j]), self.step)

    def save_opts(self):
        """Save options to disk so we know what we ran this experiment with
        """
        to_save = self.opt.__dict__.copy()
        with open(os.path.join(self.log_path, 'opt.json'), 'w') as f:
            json.dump(to_save, f, indent=2)
        s = os.path.split(os.path.realpath(__file__))[0]+'/'
        s = 'ls {} | grep -v splits | xargs'.format(s)
        t = os.path.join(self.log_path, 'codes')
        os.system("rm -rf {}".format(t))
        os.mkdir(t)
        os.system("cp -rf `{}` {}".format(s, t))


    def save_model(self, ep_end=False, batch_idx=0):
        """Save model weights to disk
        """

        models_dir = os.path.join(self.log_path, "ckpts_each_epoch")
        if not os.path.exists(models_dir):
            os.makedirs(models_dir)
        to_save = {}
        for model_name, model in self.models.items():
            if self.opt.world_size == 1:
                to_save[model_name] = model.state_dict()
            else:
                to_save[model_name] = model.module.state_dict()
        to_save['height'] = self.opt.height
        to_save['width'] = self.opt.width
        to_save['use_stereo'] = self.opt.use_stereo     
        to_save['step_in_total'] = self.step
        to_save["batch_idx"] = batch_idx
        to_save['optimizer'] = self.model_optimizer.state_dict()
        to_save['lr_scheduler'] = self.model_lr_scheduler.state_dict()

        save_ep_path = os.path.join(models_dir, "ckpt{}.pth".format(self.epoch))
        save_path = os.path.join(self.log_path, "ckpt.pth")
        if ep_end:
            to_save["epoch"] = self.epoch + 1
            torch.save(to_save, save_ep_path)
        else:   
            to_save["epoch"] = self.epoch 

        torch.save(to_save, save_path)


    def load_ckpt(self):
        """Load checkpoint to resume a training, used in training process.
        """
        logging.info(" ")
        load_path = os.path.join(self.log_path, "ckpt.pth")
        if not os.path.exists(load_path):
            logging.info("No checkpoint to resume, train from epoch 0.")
            return None

        logging.info("Resume checkpoint from {}".format(os.path.abspath(load_path)))
        checkpoint = torch.load(load_path, map_location='cpu')
        for model_name, model in self.models.items():
            model_dict = model.state_dict()
            pretrained_dict = checkpoint[model_name]
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            model_dict.update(pretrained_dict)
            model.load_state_dict(model_dict)
        self.ep_start = checkpoint['epoch']
        self.batch_start = checkpoint['batch_idx']
        self.step = checkpoint['step_in_total']
        logging.info("Start at eopch {}, batch index {}".format(self.ep_start, self.batch_start))
        return checkpoint

    def load_pretrained_model(self):
        """Load pretrained model(s) from disk, used for initializing.
        """
        self.opt.pretrained_path = os.path.abspath(self.opt.pretrained_path)

        assert os.path.exists(self.opt.pretrained_path), \
            "Cannot find folder {}".format(self.opt.pretrained_path)
        logging.info("Loading pretrained model from folder {}".format(self.opt.pretrained_path))

        checkpoint = torch.load(self.opt.pretrained_path, map_location='cpu')
        for model_name, model in self.models.items():
            model_dict = model.state_dict()
            pretrained_dict = checkpoint[model_name]
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            model_dict.update(pretrained_dict)
            model.load_state_dict(model_dict)



if __name__ == "__main__":
    opts.world_size = torch.cuda.device_count()
    if opts.world_size > 1:
        dist.init_process_group(backend='nccl', init_method='env://')
        torch.cuda.set_device(opts.local_rank)
        opts.global_rank = torch.distributed.get_rank()
    trainer = Trainer(opts)
    trainer.train()

train.py
Sauf4896/BDEdepth
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function
import os
import hashlib
import zipfile
from six.moves import urllib
import logging.config
import numpy as np
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR
from thop import clever_format
from thop import profile


def profile_once(encoder, decoder, x):
    x_e = x[0, :, :, :].unsqueeze(0)
    x_d = encoder(x_e)
    flops_e, params_e = profile(encoder, inputs=(x_e, ), verbose=False)
    flops_d, params_d = profile(decoder, inputs=(x_d, ), verbose=False)

    flops, params = clever_format([flops_e + flops_d, params_e + params_d], "%.3f")
    flops_e, params_e = clever_format([flops_e, params_e], "%.3f")
    flops_d, params_d = clever_format([flops_d, params_d], "%.3f")

    return flops, params, flops_e, params_e, flops_d, params_d


def time_sync():
    # PyTorch-accurate time
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    return time.time()


def setup_logging(log_file='log.txt',filemode='w', rank=0):
    """Setup logging configuration
    """
    logging.basicConfig(level=logging.INFO if rank == 0 else logging.WARN,
                        format="%(asctime)s - %(levelname)s - %(message)s",
                        datefmt="%Y-%m-%d %H:%M:%S",
                        filename=log_file,
                        filemode=filemode)
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter('%(message)s')
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)

def cosine_schedule_with_warmup(optimizer: Optimizer,
                                num_training_steps: int,
                                num_warmup_steps: int = 2000,
                                num_cycles: float = 0.5,
                                eta_min: float = 0.0,
                                last_epoch: int = -1):
    assert num_training_steps > num_warmup_steps
    def lr_lambda(current_step):
        current_step += 1
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float( max(1, num_training_steps - num_warmup_steps + 1))
        return max(0.0, 0.5 * (1.0 + np.cos(np.pi * float(num_cycles) * 2.0 * progress))*(1-eta_min) + eta_min)

    return LambdaLR(optimizer, lr_lambda, last_epoch)

def readlines(filename):
    """Read all the lines in a text file and return as a list
    """
    with open(filename, 'r') as f:
        lines = f.read().splitlines()
    return lines


def normalize_image(x):
    """Rescale image pixels to span range [0, 1]
    """
    ma = float(x.max().cpu().data)
    mi = float(x.min().cpu().data)
    d = ma - mi if ma != mi else 1e5
    return (x - mi) / d


def sec_to_hm(t):
    """Convert time in seconds to time in hours, minutes and seconds
    e.g. 10239 -> (2, 50, 39)
    """
    t = int(t)
    s = t % 60
    t //= 60
    m = t % 60
    t //= 60
    return t, m, s


def sec_to_hm_str(t):
    """Convert time in seconds to a nice string
    e.g. 10239 -> '02h50m39s'
    """
    h, m, s = sec_to_hm(t)
    return "{:02d}h{:02d}m{:02d}s".format(h, m, s)


def download_model_if_doesnt_exist(model_name):
    """If pretrained kitti model doesn't exist, download and unzip it
    """
    # values are tuples of (<google cloud URL>, <md5 checksum>)
    download_paths = {
        "mono_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_640x192.zip",
             "a964b8356e08a02d009609d9e3928f7c"),
        "stereo_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_640x192.zip",
             "3dfb76bcff0786e4ec07ac00f658dd07"),
        "mono+stereo_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_640x192.zip",
             "c024d69012485ed05d7eaa9617a96b81"),
        "mono_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_no_pt_640x192.zip",
             "9c2f071e35027c895a4728358ffc913a"),
        "stereo_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_no_pt_640x192.zip",
             "41ec2de112905f85541ac33a854742d1"),
        "mono+stereo_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_no_pt_640x192.zip",
             "46c3b824f541d143a45c37df65fbab0a"),
        "mono_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_1024x320.zip",
             "0ab0766efdfeea89a0d9ea8ba90e1e63"),
        "stereo_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_1024x320.zip",
             "afc2f2126d70cf3fdf26b550898b501a"),
        "mono+stereo_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_1024x320.zip",
             "cdc5fc9b23513c07d5b19235d9ef08f7"),
        }

    if not os.path.exists("models"):
        os.makedirs("models")

    model_path = os.path.join("models", model_name)

    def check_file_matches_md5(checksum, fpath):
        if not os.path.exists(fpath):
            return False
        with open(fpath, 'rb') as f:
            current_md5checksum = hashlib.md5(f.read()).hexdigest()
        return current_md5checksum == checksum

    # see if we have the model already downloaded...
    if not os.path.exists(os.path.join(model_path, "encoder.pth")):

        model_url, required_md5checksum = download_paths[model_name]

        if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
            print("-> Downloading pretrained model to {}".format(model_path + ".zip"))
            urllib.request.urlretrieve(model_url, model_path + ".zip")

        if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
            print("   Failed to download a file which matches the checksum - quitting")
            quit()

        print("   Unzipping model...")
        with zipfile.ZipFile(model_path + ".zip", 'r') as f:
            f.extractall(model_path)

        print("   Model unzipped to {}".format(model_path))

utils.py
Sauf4896/BDEdepth