Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/__init__.py

@@ -0,0 +1 @@
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/augmentation.py

@@ -0,0 +1,58 @@
+import torch
+import torchvision.transforms as transforms
+import numpy as np
+import torch.nn.functional as F
+
+
+class RGBDAugmentor:
+    """ perform augmentation on RGB-D video """
+
+    def __init__(self, crop_size):
+        self.crop_size = crop_size
+        self.augcolor = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.4/3.14),
+            transforms.RandomGrayscale(p=0.1),
+            transforms.ToTensor()])
+
+        self.max_scale = 0.25
+
+    def spatial_transform(self, images, depths, poses, intrinsics):
+        """ cropping and resizing """
+        ht, wd = images.shape[2:]
+
+        max_scale = self.max_scale
+        min_scale = np.log2(np.maximum(
+            (self.crop_size[0] + 1) / float(ht),
+            (self.crop_size[1] + 1) / float(wd)))
+
+        scale = 2 ** np.random.uniform(min_scale, max_scale)
+        intrinsics = scale * intrinsics
+        depths = depths.unsqueeze(dim=1)
+
+        images = F.interpolate(images, scale_factor=scale, mode='bilinear', 
+            align_corners=False, recompute_scale_factor=True)
+        
+        depths = F.interpolate(depths, scale_factor=scale, recompute_scale_factor=True)
+
+        # always perform center crop (TODO: try non-center crops)
+        y0 = (images.shape[2] - self.crop_size[0]) // 2
+        x0 = (images.shape[3] - self.crop_size[1]) // 2
+
+        intrinsics = intrinsics - torch.tensor([0.0, 0.0, x0, y0])
+        images = images[:, :, y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        depths = depths[:, :, y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+
+        depths = depths.squeeze(dim=1)
+        return images, poses, depths, intrinsics
+
+    def color_transform(self, images):
+        """ color jittering """
+        num, ch, ht, wd = images.shape
+        images = images.permute(1, 2, 3, 0).reshape(ch, ht, wd*num)
+        images = 255 * self.augcolor(images[[2,1,0]] / 255.0)
+        return images[[2,1,0]].reshape(ch, ht, wd, num).permute(3,0,1,2).contiguous()
+
+    def __call__(self, images, poses, depths, intrinsics):
+        images = self.color_transform(images)
+        return self.spatial_transform(images, depths, poses, intrinsics)

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/base.py

@@ -0,0 +1,157 @@
+
+import numpy as np
+import torch
+import torch.utils.data as data
+import torch.nn.functional as F
+
+import csv
+import os
+import cv2
+import math
+import random
+import json
+import pickle
+import os.path as osp
+
+from .augmentation import RGBDAugmentor
+from .rgbd_utils import *
+
+class RGBDDataset(data.Dataset):
+    def __init__(self, name, datapath, n_frames=4, crop_size=[384,512], fmin=8.0, fmax=75.0, do_aug=True):
+        """ Base class for RGBD dataset """
+        self.aug = None
+        self.root = datapath
+        self.name = name
+
+        self.n_frames = n_frames
+        self.fmin = fmin # exclude very easy examples
+        self.fmax = fmax # exclude very hard examples
+        
+        if do_aug:
+            self.aug = RGBDAugmentor(crop_size=crop_size)
+
+        # building dataset is expensive, cache so only needs to be performed once
+        cur_path = osp.dirname(osp.abspath(__file__))
+        if not os.path.isdir(osp.join(cur_path, 'cache')):
+            os.mkdir(osp.join(cur_path, 'cache'))
+        
+        cache_path = osp.join(cur_path, 'cache', '{}.pickle'.format(self.name))
+
+        if osp.isfile(cache_path):
+            scene_info = pickle.load(open(cache_path, 'rb'))[0]
+        else:
+            scene_info = self._build_dataset()
+            with open(cache_path, 'wb') as cachefile:
+                pickle.dump((scene_info,), cachefile)
+
+        self.scene_info = scene_info
+        self._build_dataset_index()
+                
+    def _build_dataset_index(self):
+        self.dataset_index = []
+        for scene in self.scene_info:
+            if not self.__class__.is_test_scene(scene):
+                graph = self.scene_info[scene]['graph']
+                for i in graph:
+                    if len(graph[i][0]) > self.n_frames:
+                        self.dataset_index.append((scene, i))
+            else:
+                print("Reserving {} for validation".format(scene))
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        return np.load(depth_file)
+
+    def build_frame_graph(self, poses, depths, intrinsics, f=16, max_flow=256):
+        """ compute optical flow distance between all pairs of frames """
+        def read_disp(fn):
+            depth = self.__class__.depth_read(fn)[f//2::f, f//2::f]
+            depth[depth < 0.01] = np.mean(depth)
+            return 1.0 / depth
+
+        poses = np.array(poses)
+        intrinsics = np.array(intrinsics) / f
+        
+        disps = np.stack(list(map(read_disp, depths)), 0)
+        d = f * compute_distance_matrix_flow(poses, disps, intrinsics)
+
+        # uncomment for nice visualization
+        # import matplotlib.pyplot as plt
+        # plt.imshow(d)
+        # plt.show()
+
+        graph = {}
+        for i in range(d.shape[0]):
+            j, = np.where(d[i] < max_flow)
+            graph[i] = (j, d[i,j])
+
+        return graph
+
+    def __getitem__(self, index):
+        """ return training video """
+
+        index = index % len(self.dataset_index)
+        scene_id, ix = self.dataset_index[index]
+
+        frame_graph = self.scene_info[scene_id]['graph']
+        images_list = self.scene_info[scene_id]['images']
+        depths_list = self.scene_info[scene_id]['depths']
+        poses_list = self.scene_info[scene_id]['poses']
+        intrinsics_list = self.scene_info[scene_id]['intrinsics']
+
+        inds = [ ix ]
+        while len(inds) < self.n_frames:
+            # get other frames within flow threshold
+            k = (frame_graph[ix][1] > self.fmin) & (frame_graph[ix][1] < self.fmax)
+            frames = frame_graph[ix][0][k]
+
+            # prefer frames forward in time
+            if np.count_nonzero(frames[frames > ix]):
+                ix = np.random.choice(frames[frames > ix])
+            
+            elif np.count_nonzero(frames):
+                ix = np.random.choice(frames)
+
+            inds += [ ix ]
+
+        images, depths, poses, intrinsics = [], [], [], []
+        for i in inds:
+            images.append(self.__class__.image_read(images_list[i]))
+            depths.append(self.__class__.depth_read(depths_list[i]))
+            poses.append(poses_list[i])
+            intrinsics.append(intrinsics_list[i])
+
+        images = np.stack(images).astype(np.float32)
+        depths = np.stack(depths).astype(np.float32)
+        poses = np.stack(poses).astype(np.float32)
+        intrinsics = np.stack(intrinsics).astype(np.float32)
+
+        images = torch.from_numpy(images).float()
+        images = images.permute(0, 3, 1, 2)
+
+        disps = torch.from_numpy(1.0 / depths)
+        poses = torch.from_numpy(poses)
+        intrinsics = torch.from_numpy(intrinsics)
+
+        if self.aug is not None:
+            images, poses, disps, intrinsics = \
+                self.aug(images, poses, disps, intrinsics)
+
+        # scale scene
+        if len(disps[disps>0.01]) > 0:
+            s = disps[disps>0.01].mean()
+            disps = disps / s
+            poses[...,:3] *= s
+
+        return images, poses, disps, intrinsics 
+
+    def __len__(self):
+        return len(self.dataset_index)
+
+    def __imul__(self, x):
+        self.dataset_index *= x
+        return self

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/factory.py

@@ -0,0 +1,82 @@
+
+import pickle
+import os
+import os.path as osp
+
+# RGBD-Dataset
+from .tartan import TartanAir
+
+from .stream import ImageStream
+from .stream import StereoStream
+from .stream import RGBDStream
+
+# streaming datasets for inference
+from .tartan import TartanAirStream
+from .tartan import TartanAirTestStream
+
+def dataset_factory(dataset_list, **kwargs):
+    """ create a combined dataset """
+
+    from torch.utils.data import ConcatDataset
+
+    dataset_map = { 'tartan': (TartanAir, ) }
+    db_list = []
+    for key in dataset_list:
+        # cache datasets for faster future loading
+        db = dataset_map[key][0](**kwargs)
+
+        print("Dataset {} has {} images".format(key, len(db)))
+        db_list.append(db)
+
+    return ConcatDataset(db_list)
+            
+
+def create_datastream(dataset_path, **kwargs):
+    """ create data_loader to stream images 1 by 1 """
+
+    from torch.utils.data import DataLoader
+
+    if osp.isfile(osp.join(dataset_path, 'calibration.txt')):
+        db = ETH3DStream(dataset_path, **kwargs)
+
+    elif osp.isdir(osp.join(dataset_path, 'image_left')):
+        db = TartanAirStream(dataset_path, **kwargs)
+
+    elif osp.isfile(osp.join(dataset_path, 'rgb.txt')):
+        db = TUMStream(dataset_path, **kwargs)
+
+    elif osp.isdir(osp.join(dataset_path, 'mav0')):
+        db = EurocStream(dataset_path, **kwargs)
+
+    elif osp.isfile(osp.join(dataset_path, 'calib.txt')):
+        db = KITTIStream(dataset_path, **kwargs)
+
+    else:
+        # db = TartanAirStream(dataset_path, **kwargs)
+        db = TartanAirTestStream(dataset_path, **kwargs)
+    
+    stream = DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
+    return stream
+
+
+def create_imagestream(dataset_path, **kwargs):
+    """ create data_loader to stream images 1 by 1 """
+    from torch.utils.data import DataLoader
+
+    db = ImageStream(dataset_path, **kwargs)
+    return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
+
+def create_stereostream(dataset_path, **kwargs):
+    """ create data_loader to stream images 1 by 1 """
+    from torch.utils.data import DataLoader
+
+    db = StereoStream(dataset_path, **kwargs)
+    return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
+
+def create_rgbdstream(dataset_path, **kwargs):
+    """ create data_loader to stream images 1 by 1 """
+    from torch.utils.data import DataLoader
+
+    db = RGBDStream(dataset_path, **kwargs)
+    return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/rgbd_utils.py

@@ -0,0 +1,190 @@
+import numpy as np
+import os.path as osp
+
+import torch
+from lietorch import SE3
+
+import geom.projective_ops as pops
+from scipy.spatial.transform import Rotation
+
+
+def parse_list(filepath, skiprows=0):
+    """ read list data """
+    data = np.loadtxt(filepath, delimiter=' ', dtype=np.unicode_, skiprows=skiprows)
+    return data
+
+def associate_frames(tstamp_image, tstamp_depth, tstamp_pose, max_dt=1.0):
+    """ pair images, depths, and poses """
+    associations = []
+    for i, t in enumerate(tstamp_image):
+        if tstamp_pose is None:
+            j = np.argmin(np.abs(tstamp_depth - t))
+            if (np.abs(tstamp_depth[j] - t) < max_dt):
+                associations.append((i, j))
+
+        else:
+            j = np.argmin(np.abs(tstamp_depth - t))
+            k = np.argmin(np.abs(tstamp_pose - t))
+        
+            if (np.abs(tstamp_depth[j] - t) < max_dt) and \
+                    (np.abs(tstamp_pose[k] - t) < max_dt):
+                associations.append((i, j, k))
+            
+    return associations
+
+def loadtum(datapath, frame_rate=-1):
+    """ read video data in tum-rgbd format """
+    if osp.isfile(osp.join(datapath, 'groundtruth.txt')):
+        pose_list = osp.join(datapath, 'groundtruth.txt')
+    
+    elif osp.isfile(osp.join(datapath, 'pose.txt')):
+        pose_list = osp.join(datapath, 'pose.txt')
+
+    else:
+        return None, None, None, None
+
+    image_list = osp.join(datapath, 'rgb.txt')
+    depth_list = osp.join(datapath, 'depth.txt')
+
+    calib_path = osp.join(datapath, 'calibration.txt')
+    intrinsic = None
+    if osp.isfile(calib_path):
+        intrinsic = np.loadtxt(calib_path, delimiter=' ')
+        intrinsic = intrinsic.astype(np.float64)
+
+    image_data = parse_list(image_list)
+    depth_data = parse_list(depth_list)
+    pose_data = parse_list(pose_list, skiprows=1)
+    pose_vecs = pose_data[:,1:].astype(np.float64)
+
+    tstamp_image = image_data[:,0].astype(np.float64)
+    tstamp_depth = depth_data[:,0].astype(np.float64)
+    tstamp_pose = pose_data[:,0].astype(np.float64)
+    associations = associate_frames(tstamp_image, tstamp_depth, tstamp_pose)
+
+    # print(len(tstamp_image))
+    # print(len(associations))
+
+    indicies = range(len(associations))[::5]
+
+    # indicies = [ 0 ]
+    # for i in range(1, len(associations)):
+    #     t0 = tstamp_image[associations[indicies[-1]][0]]
+    #     t1 = tstamp_image[associations[i][0]]
+    #     if t1 - t0 > 1.0 / frame_rate:
+    #         indicies += [ i ]
+
+    images, poses, depths, intrinsics, tstamps = [], [], [], [], []
+    for ix in indicies:
+        (i, j, k) = associations[ix]
+        images += [ osp.join(datapath, image_data[i,1]) ]
+        depths += [ osp.join(datapath, depth_data[j,1]) ]
+        poses += [ pose_vecs[k] ]
+        tstamps += [ tstamp_image[i] ]
+        
+        if intrinsic is not None:
+            intrinsics += [ intrinsic ]
+
+    return images, depths, poses, intrinsics, tstamps
+
+
+def all_pairs_distance_matrix(poses, beta=2.5):
+    """ compute distance matrix between all pairs of poses """
+    poses = np.array(poses, dtype=np.float32)
+    poses[:,:3] *= beta # scale to balence rot + trans
+    poses = SE3(torch.from_numpy(poses))
+
+    r = (poses[:,None].inv() * poses[None,:]).log()
+    return r.norm(dim=-1).cpu().numpy()
+
+def pose_matrix_to_quaternion(pose):
+    """ convert 4x4 pose matrix to (t, q) """
+    q = Rotation.from_matrix(pose[:3, :3]).as_quat()
+    return np.concatenate([pose[:3, 3], q], axis=0)
+
+def compute_distance_matrix_flow(poses, disps, intrinsics):
+    """ compute flow magnitude between all pairs of frames """
+    if not isinstance(poses, SE3):
+        poses = torch.from_numpy(poses).float().cuda()[None]
+        poses = SE3(poses).inv()
+
+        disps = torch.from_numpy(disps).float().cuda()[None]
+        intrinsics = torch.from_numpy(intrinsics).float().cuda()[None]
+
+    N = poses.shape[1]
+    
+    ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
+    ii = ii.reshape(-1).cuda()
+    jj = jj.reshape(-1).cuda()
+
+    MAX_FLOW = 100.0
+    matrix = np.zeros((N, N), dtype=np.float32)
+
+    s = 2048
+    for i in range(0, ii.shape[0], s):
+        flow1, val1 = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
+        flow2, val2 = pops.induced_flow(poses, disps, intrinsics, jj[i:i+s], ii[i:i+s])
+        
+        flow = torch.stack([flow1, flow2], dim=2)
+        val = torch.stack([val1, val2], dim=2)
+        
+        mag = flow.norm(dim=-1).clamp(max=MAX_FLOW)
+        mag = mag.view(mag.shape[1], -1)
+        val = val.view(val.shape[1], -1)
+
+        mag = (mag * val).mean(-1) / val.mean(-1)
+        mag[val.mean(-1) < 0.7] = np.inf
+
+        i1 = ii[i:i+s].cpu().numpy()
+        j1 = jj[i:i+s].cpu().numpy()
+        matrix[i1, j1] = mag.cpu().numpy()
+
+    return matrix
+
+
+def compute_distance_matrix_flow2(poses, disps, intrinsics, beta=0.4):
+    """ compute flow magnitude between all pairs of frames """
+    # if not isinstance(poses, SE3):
+    #     poses = torch.from_numpy(poses).float().cuda()[None]
+    #     poses = SE3(poses).inv()
+
+    #     disps = torch.from_numpy(disps).float().cuda()[None]
+    #     intrinsics = torch.from_numpy(intrinsics).float().cuda()[None]
+
+    N = poses.shape[1]
+    
+    ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
+    ii = ii.reshape(-1)
+    jj = jj.reshape(-1)
+
+    MAX_FLOW = 128.0
+    matrix = np.zeros((N, N), dtype=np.float32)
+
+    s = 2048
+    for i in range(0, ii.shape[0], s):
+        flow1a, val1a = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s], tonly=True)
+        flow1b, val1b = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
+        flow2a, val2a = pops.induced_flow(poses, disps, intrinsics, jj[i:i+s], ii[i:i+s], tonly=True)
+        flow2b, val2b = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
+
+        flow1 = flow1a + beta * flow1b
+        val1 = val1a * val2b
+
+        flow2 = flow2a + beta * flow2b
+        val2 = val2a * val2b
+        
+        flow = torch.stack([flow1, flow2], dim=2)
+        val = torch.stack([val1, val2], dim=2)
+        
+        mag = flow.norm(dim=-1).clamp(max=MAX_FLOW)
+        mag = mag.view(mag.shape[1], -1)
+        val = val.view(val.shape[1], -1)
+
+        mag = (mag * val).mean(-1) / val.mean(-1)
+        mag[val.mean(-1) < 0.8] = np.inf
+
+        i1 = ii[i:i+s].cpu().numpy()
+        j1 = jj[i:i+s].cpu().numpy()
+        matrix[i1, j1] = mag.cpu().numpy()
+
+    return matrix

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/stream.py

@@ -0,0 +1,234 @@
+
+import numpy as np
+import torch
+import torch.utils.data as data
+import torch.nn.functional as F
+
+import csv
+import os
+import cv2
+import math
+import random
+import json
+import pickle
+import os.path as osp
+
+from .rgbd_utils import *
+
+class RGBDStream(data.Dataset):
+    def __init__(self, datapath, frame_rate=-1, image_size=[384,512], crop_size=[0,0]):
+        self.datapath = datapath
+        self.frame_rate = frame_rate
+        self.image_size = image_size
+        self.crop_size = crop_size
+        self._build_dataset_index()
+    
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        return np.load(depth_file)
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        """ return training video """
+        image = self.__class__.image_read(self.images[index])
+        image = torch.from_numpy(image).float()
+        image = image.permute(2, 0, 1)
+
+        try:
+            tstamp = self.tstamps[index]
+        except:
+            tstamp = index
+
+        pose = torch.from_numpy(self.poses[index]).float()
+        intrinsic = torch.from_numpy(self.intrinsics[index]).float()
+
+        # resize image
+        sx = self.image_size[1] / image.shape[2]
+        sy = self.image_size[0] / image.shape[1]
+
+        image = F.interpolate(image[None], self.image_size, mode='bilinear', align_corners=False)[0]
+
+        fx, fy, cx, cy = intrinsic.unbind(dim=0)
+        fx, cx = sx * fx, sx * cx
+        fy, cy = sy * fy, sy * cy
+        
+        # crop image
+        if self.crop_size[0] > 0:
+            cy = cy - self.crop_size[0]
+            image = image[:,self.crop_size[0]:-self.crop_size[0],:]
+
+        if self.crop_size[1] > 0:
+            cx = cx - self.crop_size[1]
+            image = image[:,:,self.crop_size[1]:-self.crop_size[1]]
+
+        intrinsic = torch.stack([fx, fy, cx, cy])
+
+        return tstamp, image, pose, intrinsic
+
+
+class ImageStream(data.Dataset):
+    def __init__(self, datapath, intrinsics, rate=1, image_size=[384,512]):
+        rgb_list = osp.join(datapath, 'rgb.txt')
+        if os.path.isfile(rgb_list):
+            rgb_list = np.loadtxt(rgb_list, delimiter=' ', dtype=np.unicode_)
+            self.timestamps = rgb_list[:,0].astype(np.float)
+            self.images = [os.path.join(datapath, x) for x in rgb_list[:,1]]
+            self.images = self.images[::rate]
+            self.timestamps = self.timestamps[::rate]
+
+        else:
+            import glob
+            self.images = sorted(glob.glob(osp.join(datapath, '*.jpg'))) +  sorted(glob.glob(osp.join(datapath, '*.png')))
+            self.images = self.images[::rate]
+
+        self.intrinsics = intrinsics
+        self.image_size = image_size
+
+    def __len__(self):
+        return len(self.images)
+
+    @staticmethod
+    def image_read(imfile):
+        return cv2.imread(imfile)
+
+    def __getitem__(self, index):
+        """ return training video """
+        image = self.__class__.image_read(self.images[index])
+
+        try:
+            tstamp = self.timestamps[index]
+        except:
+            tstamp = index
+
+        ht0, wd0 = image.shape[:2]
+        ht1, wd1 = self.image_size
+
+        intrinsics = torch.as_tensor(self.intrinsics)
+        intrinsics[0] *= wd1 / wd0
+        intrinsics[1] *= ht1 / ht0
+        intrinsics[2] *= wd1 / wd0
+        intrinsics[3] *= ht1 / ht0
+
+        # resize image
+        ikwargs = {'mode': 'bilinear', 'align_corners': True}
+        image = torch.from_numpy(image).float().permute(2, 0, 1)
+        image = F.interpolate(image[None], self.image_size, **ikwargs)[0]
+
+        return tstamp, image, intrinsics
+
+
+
+class StereoStream(data.Dataset):
+    def __init__(self, datapath, intrinsics, rate=1, image_size=[384,512], 
+            map_left=None, map_right=None, left_root='image_left', right_root='image_right'):
+        import glob
+        self.intrinsics = intrinsics
+        self.image_size = image_size
+        
+        imgs = sorted(glob.glob(osp.join(datapath, left_root, '*.png')))[::rate]
+        self.images_l = []
+        self.images_r = []
+        self.tstamps = []
+
+        for img_l in imgs:
+            img_r = img_l.replace(left_root, right_root)
+            if os.path.isfile(img_r):
+                t = np.float(img_l.split('/')[-1].replace('.png', ''))
+                self.tstamps.append(t)
+                self.images_l += [ img_l ]
+                self.images_r += [ img_r ]
+
+        self.map_left = map_left
+        self.map_right = map_right
+
+    def __len__(self):
+        return len(self.images_l)
+
+    @staticmethod
+    def image_read(imfile, imap=None):
+        image = cv2.imread(imfile)
+        if imap is not None:
+            image = cv2.remap(image, imap[0], imap[1], interpolation=cv2.INTER_LINEAR)
+        return image
+
+    def __getitem__(self, index):
+        """ return training video """
+        tstamp = self.tstamps[index]
+        image_l = self.__class__.image_read(self.images_l[index], self.map_left)
+        image_r = self.__class__.image_read(self.images_r[index], self.map_right)
+
+        ht0, wd0 = image_l.shape[:2]
+        ht1, wd1 = self.image_size
+
+        intrinsics = torch.as_tensor(self.intrinsics)
+        intrinsics[0] *= wd1 / wd0
+        intrinsics[1] *= ht1 / ht0
+        intrinsics[2] *= wd1 / wd0
+        intrinsics[3] *= ht1 / ht0
+
+        image_l = torch.from_numpy(image_l).float().permute(2, 0, 1)
+        image_r = torch.from_numpy(image_r).float().permute(2, 0, 1)
+
+        # resize image
+        ikwargs = {'mode': 'bilinear', 'align_corners': True}
+        image_l = F.interpolate(image_l[None], self.image_size, **ikwargs)[0]
+        image_r = F.interpolate(image_r[None], self.image_size, **ikwargs)[0]
+
+        return tstamp, image_l, image_r, intrinsics
+
+
+
+# class RGBDStream(data.Dataset):
+#     def __init__(self, datapath, intrinsics=None, rate=1, image_size=[384,512]):
+#         assoc_file = osp.join(datapath, 'associated.txt')
+#         assoc_list = np.loadtxt(assoc_file, delimiter=' ', dtype=np.unicode_)
+        
+#         self.intrinsics = intrinsics
+#         self.image_size = image_size
+        
+#         self.timestamps = assoc_list[:,0].astype(np.float)[::rate]
+#         self.images = [os.path.join(datapath, x) for x in assoc_list[:,1]][::rate]
+#         self.depths = [os.path.join(datapath, x) for x in assoc_list[:,3]][::rate]
+
+#     def __len__(self):
+#         return len(self.images)
+
+#     @staticmethod
+#     def image_read(imfile):
+#         return cv2.imread(imfile)
+
+#     @staticmethod
+#     def depth_read(depth_file):
+#         depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+#         return depth.astype(np.float32) / 5000.0
+
+#     def __getitem__(self, index):
+#         """ return training video """
+#         tstamp = self.timestamps[index]
+#         image = self.__class__.image_read(self.images[index])
+#         depth = self.__class__.depth_read(self.depths[index])
+
+#         ht0, wd0 = image.shape[:2]
+#         ht1, wd1 = self.image_size
+
+#         intrinsics = torch.as_tensor(self.intrinsics)
+#         intrinsics[0] *= wd1 / wd0
+#         intrinsics[1] *= ht1 / ht0
+#         intrinsics[2] *= wd1 / wd0
+#         intrinsics[3] *= ht1 / ht0
+
+#         # resize image
+#         ikwargs = {'mode': 'bilinear', 'align_corners': True}
+#         image = torch.from_numpy(image).float().permute(2, 0, 1)
+#         image = F.interpolate(image[None], self.image_size, **ikwargs)[0]
+
+#         depth = torch.from_numpy(depth).float()[None,None]
+#         depth = F.interpolate(depth, self.image_size, mode='nearest').squeeze()
+
+#         return tstamp, image, depth, intrinsics

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/tartan.py

@@ -0,0 +1,138 @@
+
+import numpy as np
+import torch
+import glob
+import cv2
+import os
+import os.path as osp
+
+from lietorch import SE3
+from .base import RGBDDataset
+from .stream import RGBDStream
+
+cur_path = osp.dirname(osp.abspath(__file__))
+test_split = osp.join(cur_path, 'tartan_test.txt')
+test_split = open(test_split).read().split()
+
+
+class TartanAir(RGBDDataset):
+
+    # scale depths to balance rot & trans
+    DEPTH_SCALE = 5.0
+
+    def __init__(self, mode='training', **kwargs):
+        self.mode = mode
+        self.n_frames = 2
+        super(TartanAir, self).__init__(name='TartanAir', **kwargs)
+
+    @staticmethod 
+    def is_test_scene(scene):
+        # print(scene, any(x in scene for x in test_split))
+        return any(x in scene for x in test_split)
+
+    def _build_dataset(self):
+        from tqdm import tqdm
+        print("Building TartanAir dataset")
+
+        scene_info = {}
+        scenes = glob.glob(osp.join(self.root, '*/*/*/*'))
+        for scene in tqdm(sorted(scenes)):
+            images = sorted(glob.glob(osp.join(scene, 'image_left/*.png')))
+            depths = sorted(glob.glob(osp.join(scene, 'depth_left/*.npy')))
+            
+            poses = np.loadtxt(osp.join(scene, 'pose_left.txt'), delimiter=' ')
+            poses = poses[:, [1, 2, 0, 4, 5, 3, 6]]
+            poses[:,:3] /= TartanAir.DEPTH_SCALE
+            intrinsics = [TartanAir.calib_read()] * len(images)
+
+            # graph of co-visible frames based on flow
+            graph = self.build_frame_graph(poses, depths, intrinsics)
+
+            scene = '/'.join(scene.split('/'))
+            scene_info[scene] = {'images': images, 'depths': depths, 
+                'poses': poses, 'intrinsics': intrinsics, 'graph': graph}
+
+        return scene_info
+
+    @staticmethod
+    def calib_read():
+        return np.array([320.0, 320.0, 320.0, 240.0])
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = np.load(depth_file) / TartanAir.DEPTH_SCALE
+        depth[depth==np.nan] = 1.0
+        depth[depth==np.inf] = 1.0
+        return depth
+
+
+class TartanAirStream(RGBDStream):
+    def __init__(self, datapath, **kwargs):
+        super(TartanAirStream, self).__init__(datapath=datapath, **kwargs)
+
+    def _build_dataset_index(self):
+        """ build list of images, poses, depths, and intrinsics """
+        self.root = 'datasets/TartanAir'
+
+        scene = osp.join(self.root, self.datapath)
+        image_glob = osp.join(scene, 'image_left/*.png')
+        images = sorted(glob.glob(image_glob))
+
+        poses = np.loadtxt(osp.join(scene, 'pose_left.txt'), delimiter=' ')
+        poses = poses[:, [1, 2, 0, 4, 5, 3, 6]]
+
+        poses = SE3(torch.as_tensor(poses))
+        poses = poses[[0]].inv() * poses
+        poses = poses.data.cpu().numpy()
+
+        intrinsic = self.calib_read(self.datapath)
+        intrinsics = np.tile(intrinsic[None], (len(images), 1))
+
+        self.images = images[::int(self.frame_rate)]
+        self.poses = poses[::int(self.frame_rate)]
+        self.intrinsics = intrinsics[::int(self.frame_rate)]
+
+    @staticmethod
+    def calib_read(datapath):
+        return np.array([320.0, 320.0, 320.0, 240.0])
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+
+class TartanAirTestStream(RGBDStream):
+    def __init__(self, datapath, **kwargs):
+        super(TartanAirTestStream, self).__init__(datapath=datapath, **kwargs)
+
+    def _build_dataset_index(self):
+        """ build list of images, poses, depths, and intrinsics """
+        self.root = 'datasets/mono'
+        image_glob = osp.join(self.root, self.datapath, '*.png')
+        images = sorted(glob.glob(image_glob))
+
+        poses = np.loadtxt(osp.join(self.root, 'mono_gt', self.datapath + '.txt'), delimiter=' ')
+        poses = poses[:, [1, 2, 0, 4, 5, 3, 6]]
+
+        poses = SE3(torch.as_tensor(poses))
+        poses = poses[[0]].inv() * poses
+        poses = poses.data.cpu().numpy()
+
+        intrinsic = self.calib_read(self.datapath)
+        intrinsics = np.tile(intrinsic[None], (len(images), 1))
+
+        self.images = images[::int(self.frame_rate)]
+        self.poses = poses[::int(self.frame_rate)]
+        self.intrinsics = intrinsics[::int(self.frame_rate)]
+
+    @staticmethod
+    def calib_read(datapath):
+        return np.array([320.0, 320.0, 320.0, 240.0])
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/data_readers/tartan_test.txt

@@ -0,0 +1,32 @@
+abandonedfactory/abandonedfactory/Easy/P011
+abandonedfactory/abandonedfactory/Hard/P011
+abandonedfactory_night/abandonedfactory_night/Easy/P013
+abandonedfactory_night/abandonedfactory_night/Hard/P014
+amusement/amusement/Easy/P008
+amusement/amusement/Hard/P007
+carwelding/carwelding/Easy/P007
+endofworld/endofworld/Easy/P009
+gascola/gascola/Easy/P008
+gascola/gascola/Hard/P009
+hospital/hospital/Easy/P036
+hospital/hospital/Hard/P049
+japanesealley/japanesealley/Easy/P007
+japanesealley/japanesealley/Hard/P005
+neighborhood/neighborhood/Easy/P021
+neighborhood/neighborhood/Hard/P017
+ocean/ocean/Easy/P013
+ocean/ocean/Hard/P009
+office2/office2/Easy/P011
+office2/office2/Hard/P010
+office/office/Hard/P007
+oldtown/oldtown/Easy/P007
+oldtown/oldtown/Hard/P008
+seasidetown/seasidetown/Easy/P009
+seasonsforest/seasonsforest/Easy/P011
+seasonsforest/seasonsforest/Hard/P006
+seasonsforest_winter/seasonsforest_winter/Easy/P009
+seasonsforest_winter/seasonsforest_winter/Hard/P018
+soulcity/soulcity/Easy/P012
+soulcity/soulcity/Hard/P009
+westerndesert/westerndesert/Easy/P013
+westerndesert/westerndesert/Hard/P007