Dataset loader that was added to Monodepth2

dji_dataset.py

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 DjiDataset(MonoDataset):
    def __init__(self, *args, **kwargs):
        super(DjiDataset, self).__init__(*args, **kwargs)

        self.K = np.array([
            [1.54686054e+03, 0.00000000e+00, 9.52234286e+02, 0],
            [0.00000000e+00, 1.55059362e+03, 5.23351473e+02, 0],
            [0.00000000e+00, 0.00000000e+00, 1.00000000e+00, 0],
            [0,              0,              0,              1]], dtype=np.float32)
        self.K[0, :] /= 1920
        self.K[1, :] /= 1080

        self.full_res_shape = (1920, 1080)


    def check_depth(self):
        return False


    def get_color(self, folder, frame_index, side, do_flip):
        assert side is None

        color = self.loader(self.get_image_path(folder, frame_index))

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

        return color


    def get_image_path(self, folder, frame_index):
        return os.path.join(self.data_path, folder, f'out-{frame_index}{self.img_ext}')