Spaces:
Running
Running
import os | |
import random | |
from PIL import Image | |
import cv2 | |
import h5py | |
import numpy as np | |
import torch | |
from torch.utils.data import Dataset, DataLoader, ConcatDataset | |
import torchvision.transforms.functional as tvf | |
import kornia.augmentation as K | |
import os.path as osp | |
import matplotlib.pyplot as plt | |
from dkm.utils import get_depth_tuple_transform_ops, get_tuple_transform_ops | |
from dkm.utils.transforms import GeometricSequential | |
from tqdm import tqdm | |
class ScanNetScene: | |
def __init__( | |
self, | |
data_root, | |
scene_info, | |
ht=384, | |
wt=512, | |
min_overlap=0.0, | |
shake_t=0, | |
rot_prob=0.0, | |
) -> None: | |
self.scene_root = osp.join(data_root, "scans", "scans_train") | |
self.data_names = scene_info["name"] | |
self.overlaps = scene_info["score"] | |
# Only sample 10s | |
valid = (self.data_names[:, -2:] % 10).sum(axis=-1) == 0 | |
self.overlaps = self.overlaps[valid] | |
self.data_names = self.data_names[valid] | |
if len(self.data_names) > 10000: | |
pairinds = np.random.choice( | |
np.arange(0, len(self.data_names)), 10000, replace=False | |
) | |
self.data_names = self.data_names[pairinds] | |
self.overlaps = self.overlaps[pairinds] | |
self.im_transform_ops = get_tuple_transform_ops(resize=(ht, wt), normalize=True) | |
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 = cv2.imread(str(depth_ref), cv2.IMREAD_UNCHANGED) | |
depth = depth / 1000 | |
depth = torch.from_numpy(depth).float() # (h, w) | |
return depth | |
def __len__(self): | |
return len(self.data_names) | |
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 read_scannet_pose(self, path): | |
"""Read ScanNet's Camera2World pose and transform it to World2Camera. | |
Returns: | |
pose_w2c (np.ndarray): (4, 4) | |
""" | |
cam2world = np.loadtxt(path, delimiter=" ") | |
world2cam = np.linalg.inv(cam2world) | |
return world2cam | |
def read_scannet_intrinsic(self, path): | |
"""Read ScanNet's intrinsic matrix and return the 3x3 matrix.""" | |
intrinsic = np.loadtxt(path, delimiter=" ") | |
return intrinsic[:-1, :-1] | |
def __getitem__(self, pair_idx): | |
# read intrinsics of original size | |
data_name = self.data_names[pair_idx] | |
scene_name, scene_sub_name, stem_name_1, stem_name_2 = data_name | |
scene_name = f"scene{scene_name:04d}_{scene_sub_name:02d}" | |
# read the intrinsic of depthmap | |
K1 = K2 = self.read_scannet_intrinsic( | |
osp.join(self.scene_root, scene_name, "intrinsic", "intrinsic_color.txt") | |
) # the depth K is not the same, but doesnt really matter | |
# read and compute relative poses | |
T1 = self.read_scannet_pose( | |
osp.join(self.scene_root, scene_name, "pose", f"{stem_name_1}.txt") | |
) | |
T2 = self.read_scannet_pose( | |
osp.join(self.scene_root, scene_name, "pose", f"{stem_name_2}.txt") | |
) | |
T_1to2 = torch.tensor(np.matmul(T2, np.linalg.inv(T1)), dtype=torch.float)[ | |
:4, :4 | |
] # (4, 4) | |
# Load positive pair data | |
im_src_ref = os.path.join( | |
self.scene_root, scene_name, "color", f"{stem_name_1}.jpg" | |
) | |
im_pos_ref = os.path.join( | |
self.scene_root, scene_name, "color", f"{stem_name_2}.jpg" | |
) | |
depth_src_ref = os.path.join( | |
self.scene_root, scene_name, "depth", f"{stem_name_1}.png" | |
) | |
depth_pos_ref = os.path.join( | |
self.scene_root, scene_name, "depth", f"{stem_name_2}.png" | |
) | |
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]) | |
) | |
data_dict = { | |
"query": im_src, | |
"support": im_pos, | |
"query_depth": depth_src[0, 0], | |
"support_depth": depth_pos[0, 0], | |
"K1": K1, | |
"K2": K2, | |
"T_1to2": T_1to2, | |
} | |
return data_dict | |
class ScanNetBuilder: | |
def __init__(self, data_root="data/scannet") -> None: | |
self.data_root = data_root | |
self.scene_info_root = os.path.join(data_root, "scannet_indices") | |
self.all_scenes = os.listdir(self.scene_info_root) | |
def build_scenes(self, split="train", min_overlap=0.0, **kwargs): | |
# Note: split doesn't matter here as we always use same scannet_train scenes | |
scene_names = self.all_scenes | |
scenes = [] | |
for scene_name in tqdm(scene_names): | |
scene_info = np.load( | |
os.path.join(self.scene_info_root, scene_name), allow_pickle=True | |
) | |
scenes.append( | |
ScanNetScene( | |
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( | |
ScanNetBuilder("data/scannet").build_scenes(split="train") | |
) | |
mega_test[0] | |