third_party/gim/dkm/datasets/megadepth.py

import os
import random
from PIL import Image
import h5py
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, ConcatDataset

from dkm.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops
import torchvision.transforms.functional as tvf
from dkm.utils.transforms import GeometricSequential
import kornia.augmentation as K


class MegadepthScene:
    def __init__(
        self,
        data_root,
        scene_info,
        ht=384,
        wt=512,
        min_overlap=0.0,
        shake_t=0,
        rot_prob=0.0,
        normalize=True,
    ) -> None:
        self.data_root = data_root
        self.image_paths = scene_info["image_paths"]
        self.depth_paths = scene_info["depth_paths"]
        self.intrinsics = scene_info["intrinsics"]
        self.poses = scene_info["poses"]
        self.pairs = scene_info["pairs"]
        self.overlaps = scene_info["overlaps"]
        threshold = self.overlaps > min_overlap
        self.pairs = self.pairs[threshold]
        self.overlaps = self.overlaps[threshold]
        if len(self.pairs) > 100000:
            pairinds = np.random.choice(
                np.arange(0, len(self.pairs)), 100000, replace=False
            )
            self.pairs = self.pairs[pairinds]
            self.overlaps = self.overlaps[pairinds]
        # counts, bins = np.histogram(self.overlaps,20)
        # print(counts)
        self.im_transform_ops = get_tuple_transform_ops(
            resize=(ht, wt), normalize=normalize
        )
        self.depth_transform_ops = get_depth_tuple_transform_ops(
            resize=(ht, wt), normalize=False
        )
        self.wt, self.ht = wt, ht
        self.shake_t = shake_t
        self.H_generator = GeometricSequential(K.RandomAffine(degrees=90, p=rot_prob))

    def load_im(self, im_ref, crop=None):
        im = Image.open(im_ref)
        return im

    def load_depth(self, depth_ref, crop=None):
        depth = np.array(h5py.File(depth_ref, "r")["depth"])
        return torch.from_numpy(depth)

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

    def scale_intrinsic(self, K, wi, hi):
        sx, sy = self.wt / wi, self.ht / hi
        sK = torch.tensor([[sx, 0, 0], [0, sy, 0], [0, 0, 1]])
        return sK @ K

    def rand_shake(self, *things):
        t = np.random.choice(range(-self.shake_t, self.shake_t + 1), size=2)
        return [
            tvf.affine(thing, angle=0.0, translate=list(t), scale=1.0, shear=[0.0, 0.0])
            for thing in things
        ], t

    def __getitem__(self, pair_idx):
        # read intrinsics of original size
        idx1, idx2 = self.pairs[pair_idx]
        K1 = torch.tensor(self.intrinsics[idx1].copy(), dtype=torch.float).reshape(3, 3)
        K2 = torch.tensor(self.intrinsics[idx2].copy(), dtype=torch.float).reshape(3, 3)

        # read and compute relative poses
        T1 = self.poses[idx1]
        T2 = self.poses[idx2]
        T_1to2 = torch.tensor(np.matmul(T2, np.linalg.inv(T1)), dtype=torch.float)[
            :4, :4
        ]  # (4, 4)

        # Load positive pair data
        im1, im2 = self.image_paths[idx1], self.image_paths[idx2]
        depth1, depth2 = self.depth_paths[idx1], self.depth_paths[idx2]
        im_src_ref = os.path.join(self.data_root, im1)
        im_pos_ref = os.path.join(self.data_root, im2)
        depth_src_ref = os.path.join(self.data_root, depth1)
        depth_pos_ref = os.path.join(self.data_root, depth2)
        # return torch.randn((1000,1000))
        im_src = self.load_im(im_src_ref)
        im_pos = self.load_im(im_pos_ref)
        depth_src = self.load_depth(depth_src_ref)
        depth_pos = self.load_depth(depth_pos_ref)

        # Recompute camera intrinsic matrix due to the resize
        K1 = self.scale_intrinsic(K1, im_src.width, im_src.height)
        K2 = self.scale_intrinsic(K2, im_pos.width, im_pos.height)
        # Process images
        im_src, im_pos = self.im_transform_ops((im_src, im_pos))
        depth_src, depth_pos = self.depth_transform_ops(
            (depth_src[None, None], depth_pos[None, None])
        )
        [im_src, im_pos, depth_src, depth_pos], t = self.rand_shake(
            im_src, im_pos, depth_src, depth_pos
        )
        im_src, Hq = self.H_generator(im_src[None])
        depth_src = self.H_generator.apply_transform(depth_src, Hq)
        K1[:2, 2] += t
        K2[:2, 2] += t
        K1 = Hq[0] @ K1
        data_dict = {
            "query": im_src[0],
            "query_identifier": self.image_paths[idx1].split("/")[-1].split(".jpg")[0],
            "support": im_pos,
            "support_identifier": self.image_paths[idx2]
            .split("/")[-1]
            .split(".jpg")[0],
            "query_depth": depth_src[0, 0],
            "support_depth": depth_pos[0, 0],
            "K1": K1,
            "K2": K2,
            "T_1to2": T_1to2,
        }
        return data_dict


class MegadepthBuilder:
    def __init__(self, data_root="data/megadepth") -> None:
        self.data_root = data_root
        self.scene_info_root = os.path.join(data_root, "prep_scene_info")
        self.all_scenes = os.listdir(self.scene_info_root)
        self.test_scenes = ["0017.npy", "0004.npy", "0048.npy", "0013.npy"]
        self.test_scenes_loftr = ["0015.npy", "0022.npy"]

    def build_scenes(self, split="train", min_overlap=0.0, **kwargs):
        if split == "train":
            scene_names = set(self.all_scenes) - set(self.test_scenes)
        elif split == "train_loftr":
            scene_names = set(self.all_scenes) - set(self.test_scenes_loftr)
        elif split == "test":
            scene_names = self.test_scenes
        elif split == "test_loftr":
            scene_names = self.test_scenes_loftr
        else:
            raise ValueError(f"Split {split} not available")
        scenes = []
        for scene_name in scene_names:
            scene_info = np.load(
                os.path.join(self.scene_info_root, scene_name), allow_pickle=True
            ).item()
            scenes.append(
                MegadepthScene(
                    self.data_root, scene_info, min_overlap=min_overlap, **kwargs
                )
            )
        return scenes

    def weight_scenes(self, concat_dataset, alpha=0.5):
        ns = []
        for d in concat_dataset.datasets:
            ns.append(len(d))
        ws = torch.cat([torch.ones(n) / n**alpha for n in ns])
        return ws


if __name__ == "__main__":
    mega_test = ConcatDataset(MegadepthBuilder().build_scenes(split="train"))
    mega_test[0]