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/thirdparty/lietorch/examples/core/data_readers/__init__.py

@@ -0,0 +1 @@
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/augmentation.py

@@ -0,0 +1,57 @@
+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.2, contrast=0.2, saturation=0.2, hue=0.4/3.14),
+            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=True, 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/thirdparty/lietorch/examples/core/data_readers/base.py

@@ -0,0 +1,149 @@
+
+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, root, name, 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 = root
+        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__))
+        cache_path = osp.join(cur_path, 'cache', '{}.pickle'.format(self.name))
+
+        if not osp.isdir(osp.join(cur_path, 'cache')):
+            os.mkdir(osp.join(cur_path, 'cache'))
+
+        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))
+
+    @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)
+
+        depths = torch.from_numpy(depths)
+        poses = torch.from_numpy(poses)
+        intrinsics = torch.from_numpy(intrinsics)
+
+        if self.aug is not None:
+            images, poses, depths, intrinsics = \
+                self.aug(images, poses, depths, intrinsics)
+
+        return images, poses, depths, 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/thirdparty/lietorch/examples/core/data_readers/eth3d.py

@@ -0,0 +1,93 @@
+
+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 lietorch import SE3
+from .base import RGBDDataset
+from .stream import RGBDStream
+from .augmentation import RGBDAugmentor
+from .rgbd_utils import loadtum, all_pairs_distance_matrix
+
+class ETH3D(RGBDDataset):
+    def __init__(self, **kwargs):
+        super(ETH3D, self).__init__(root='datasets/ETH3D', name='ETH3D', **kwargs)
+
+    @staticmethod 
+    def is_test_scene(scene):
+        return False
+
+    def _build_dataset(self):
+        from tqdm import tqdm
+        print("Building ETH3D dataset")
+
+        scene_info = {}
+        dataset_index = []
+
+        for scene in tqdm(os.listdir(self.root)):
+            scene_path = osp.join(self.root, scene)
+            
+            if not osp.isdir(scene_path):
+                continue
+            
+            # don't use scenes with no rgb info
+            if 'dark' in scene or 'kidnap' in scene:
+                continue
+
+            scene_data, graph = {}, {}
+            images, depths, poses, intrinsics = loadtum(scene_path, skip=2)
+
+            # graph of co-visible frames based on flow
+            graph = self.build_frame_graph(poses, depths, intrinsics)
+
+            scene_info[scene] = {'images': images, 'depths': depths, 
+                'poses': poses, 'intrinsics': intrinsics, 'graph': graph}
+
+        return scene_info
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 5000.0
+
+
+class ETH3DStream(RGBDStream):
+    def __init__(self, datapath, **kwargs):
+        super(ETH3DStream, self).__init__(datapath=datapath, **kwargs)
+
+    def _build_dataset_index(self):
+        """ build list of images, poses, depths, and intrinsics """
+        images, depths, poses, intrinsics = loadtum(self.datapath, self.frame_rate)
+
+        # set first pose to identity
+        poses = SE3(torch.as_tensor(poses))
+        poses = poses[[0]].inv() * poses
+        poses = poses.data.cpu().numpy()
+
+        self.images = images
+        self.poses = poses
+        self.depths = depths
+        self.intrinsics = intrinsics
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 5000.0

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/factory.py

@@ -0,0 +1,59 @@
+
+import pickle
+import os
+import os.path as osp
+
+# RGBD-Dataset
+from .tartan import TartanAir
+from .nyu2 import NYUv2
+from .eth3d import ETH3D
+from .scannet import ScanNet
+
+# streaming datasets for inference
+from .eth3d import ETH3DStream
+from .tum import TUMStream
+from .tartan import TartanAirStream
+
+
+def dataset_factory(dataset_list, **kwargs):
+    """ create a combined dataset """
+
+    from torch.utils.data import ConcatDataset
+
+    dataset_map = {
+        'tartan': (TartanAir, 1),
+        'nyu': (NYUv2, 2),
+        'eth': (ETH3D, 5),
+        'scannet': (ScanNet, 1)}
+
+    db_list = []
+    for key in dataset_list:
+        # cache datasets for faster future loading
+        db = dataset_map[key][0](**kwargs)
+        db *= dataset_map[key][1]
+
+        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.isfile(osp.join(dataset_path, 'rgb.txt')):
+        db = TUMStream(dataset_path, **kwargs)
+    
+    elif osp.isdir(osp.join(dataset_path, 'image_left')):
+        db = TartanStream(dataset_path, **kwargs)
+
+    stream = DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
+    return stream
+
+
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/nyu2.py

@@ -0,0 +1,65 @@
+
+import numpy as np
+import torch
+import glob
+import cv2
+import os
+import os.path as osp
+
+from .base import RGBDDataset
+from .augmentation import RGBDAugmentor
+from .rgbd_utils import all_pairs_distance_matrix, loadtum
+
+class NYUv2(RGBDDataset):
+    def __init__(self, **kwargs):
+        super(NYUv2, self).__init__(root='datasets/NYUv2', name='NYUv2', **kwargs)
+
+    @staticmethod 
+    def is_test_scene(scene):
+        return False
+
+    def _build_dataset(self):
+
+        from tqdm import tqdm
+        print("Building NYUv2 dataset")
+
+        scene_info = {}
+        dataset_index = []
+
+        scenes = os.listdir(self.root)
+        for scene in tqdm(scenes):
+            scene_path = osp.join(self.root, scene)
+            images, depths, poses, intrinsics = loadtum(scene_path, frame_rate=10)
+
+            # filter out some errors in dataset
+            if images is None or len(images) < 8:
+                continue
+
+            intrinsic = NYUv2.calib_read()
+            intrinsics = [intrinsic] * len(images)
+
+            # graph of co-visible frames based on flow
+            graph = self.build_frame_graph(poses, depths, intrinsics)
+
+            scene_info[scene] = {'images': images, 'depths': depths, 
+                'poses': poses, 'intrinsics': intrinsics, 'graph': graph}
+
+        return scene_info
+
+    @staticmethod
+    def calib_read():
+        fx = 5.1885790117450188e+02
+        fy = 5.1946961112127485e+02
+        cx = 3.2558244941119034e+02
+        cy = 2.5373616633400465e+02
+        return np.array([fx, fy, cx, cy])
+
+    @staticmethod
+    def image_read(image_file):
+        return cv2.imread(image_file)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 5000.0
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/rgbd_utils.py

@@ -0,0 +1,132 @@
+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=0.08):
+    """ 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')
+
+    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)
+
+    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 = [], [], [], []
+    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] ]
+        
+        if intrinsic is not None:
+            intrinsics += [ intrinsic ]
+
+    return images, depths, poses, intrinsics
+
+
+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

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/scannet.py

@@ -0,0 +1,91 @@
+
+import numpy as np
+import pickle
+import torch
+import glob
+import cv2
+import os
+import os.path as osp
+
+from .base import RGBDDataset
+from .augmentation import RGBDAugmentor
+from .rgbd_utils import pose_matrix_to_quaternion
+from .rgbd_utils import  all_pairs_distance_matrix
+
+
+class ScanNet(RGBDDataset):
+    def __init__(self, **kwargs):
+        super(ScanNet, self).__init__(root='datasets/ScanNet', name='ScanNet', **kwargs)
+
+    @staticmethod
+    def is_test_scene(scene):
+        scanid = int(re.findall(r'scene(.+?)_', scene)[0])
+        return scanid > 660
+
+    def _build_dataset_index(self):
+        """ construct scene_info and dataset_index objects """
+
+        from tqdm import tqdm
+        print("Building ScanNet dataset")
+
+        scene_info = {}
+        dataset_index = []
+
+        for scene in tqdm(os.listdir(self.root)):
+            scene_path = osp.join(self.root, scene)
+            depth_glob = osp.join(scene_path, 'depth', '*.png')
+            depth_list = glob.glob(depth_glob)
+
+            get_indicies = lambda x: int(osp.basename(x).split('.')[0])
+            get_images = lambda i: osp.join(scene_path, 'color', '%d.jpg' % i)
+            get_depths = lambda i: osp.join(scene_path, 'depth', '%d.png' % i)
+            get_poses = lambda i: osp.join(scene_path, 'pose', '%d.txt' % i)
+
+            indicies = sorted(map(get_indicies, depth_list))[::2]
+            image_list = list(map(get_images, indicies))
+            depth_list = list(map(get_depths, indicies))
+
+            pose_list = map(get_poses, indicies)
+            pose_list = list(map(ScanNet.pose_read, pose_list))
+
+            # remove nan poses
+            pvecs = np.stack(pose_list, 0)
+            keep, = np.where(~np.any(np.isnan(pvecs) | np.isinf(pvecs), axis=1))
+            images = [image_list[i] for i in keep]
+            depths = [depth_list[i] for i in keep]
+            poses = [pose_list[i] for i in keep]
+
+            intrinsic = ScanNet.calib_read(scene_path)
+            intrinsics = [intrinsic] * len(images)
+
+            graph = self.build_frame_graph(poses, depths, intrinsics)
+
+            scene_info[scene] = {'images': images, 'depths': depths, 
+                'poses': poses, 'intrinsics': intrinsics, 'graph': graph}
+
+            for i in range(len(images)):
+                if len(graph[i][0]) > 1:
+                    dataset_index.append((scene, i))
+
+        return scene_info, dataset_index
+            
+    @staticmethod
+    def calib_read(scene_path):
+        intrinsic_file = osp.join(scene_path, 'intrinsic', 'intrinsic_depth.txt')
+        K = np.loadtxt(intrinsic_file, delimiter=' ')
+        return np.array([K[0,0], K[1,1], K[0,2], K[1,2]])
+
+    @staticmethod
+    def pose_read(pose_file):
+        pose = np.loadtxt(pose_file, delimiter=' ').astype(np.float64)
+        return pose_matrix_to_quaternion(pose)
+
+    @staticmethod
+    def image_read(image_file):
+        image = cv2.imread(image_file)
+        return cv2.resize(image, (640, 480))
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 1000.0

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/stream.py

@@ -0,0 +1,64 @@
+
+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, crop_size=[384,512]):
+        self.datapath = datapath
+        self.frame_rate = frame_rate
+        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)
+
+        depth = self.__class__.depth_read(self.depths[index])
+        depth = torch.from_numpy(depth).float()
+
+        pose = torch.from_numpy(self.poses[index]).float()
+        intrinsic = torch.from_numpy(self.intrinsics[index]).float()
+
+        sx = self.crop_size[1] / depth.shape[1]
+        sy = self.crop_size[0] / depth.shape[0]
+        image = F.interpolate(image[None], self.crop_size, mode='bilinear', align_corners=True)[0]
+        depth = F.interpolate(depth[None,None], self.crop_size, mode='nearest')[0,0]
+
+        image = image[..., 8:-8, 8:-8]
+        depth = depth[..., 8:-8, 8:-8]
+
+        fx, fy, cx, cy = intrinsic.unbind(dim=0)
+        intrinsic = torch.stack([sx*fx, sy*fy, sx*cx - 8, sy*cy - 8])
+
+        # intrinsic *= torch.as_tensor([sx, sy, sx, sy])
+        return index, image, depth, pose, intrinsic
+
+
+
+                

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/tartan.py

@@ -0,0 +1,153 @@
+
+import numpy as np
+import torch
+import glob
+import cv2
+import os
+import os.path as osp
+
+from .base import RGBDDataset
+from .stream import RGBDStream
+
+class TartanAir(RGBDDataset):
+
+    # scale depths to balance rot & trans
+    DEPTH_SCALE = 5.0
+    TEST_SET = ['westerndesert', 'seasidetown', 'seasonsforest_winter', 'office2', 'gascola']
+
+    def __init__(self, mode='training', **kwargs):
+        self.mode = mode
+        self.n_frames = 2
+        super(TartanAir, self).__init__(root='datasets/TartanAir', name='TartanAir', **kwargs)
+
+    @staticmethod
+    def is_test_scene(scene):
+        return scene.split('/')[-3] in TartanAir.TEST_SET
+
+    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 TartanAirTest(torch.utils.data.Dataset):
+    def __init__(self, root='datasets/Tartan'):
+        self.root = root
+        self.dataset_index = []
+
+        self.scene_info = {}
+        scenes = glob.glob(osp.join(self.root, '*/*/*/*'))
+        
+        for scene in sorted(scenes):
+            image_glob = osp.join(scene, 'image_left/*.png')
+            depth_glob = osp.join(scene, 'depth_left/*.npy')            
+            images = sorted(glob.glob(image_glob))
+            depths = sorted(glob.glob(depth_glob))
+
+            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)
+            
+            self.scene_info[scene] = {'images': images, 
+                'depths': depths, 'poses': poses, 'intrinsics': intrinsics}
+        
+        with open('assets/tartan_test.txt') as f:
+            self.dataset_index = f.readlines()
+        
+    def __getitem__(self, index):
+        """ load test example """
+
+        scene_id, ix1, ix2 = self.dataset_index[index].split()
+        inds = [int(ix1), int(ix2)]
+
+        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']
+
+        images, depths, poses, intrinsics = [], [], [], []
+        for i in inds:
+            images.append(TartanAir.image_read(images_list[i]))
+            depths.append(TartanAir.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)
+
+        depths = torch.from_numpy(depths)
+        poses = torch.from_numpy(poses)
+        intrinsics = torch.from_numpy(intrinsics)
+
+        return images, poses, depths, intrinsics 
+
+    def __len__(self):
+        return len(self.dataset_index)
+
+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 """
+        images, poses, depths, intrinsics = loadtum(self.datapath)
+        intrinsic = NYUv2.TUMStream(self.datapath)
+        intrinsics = np.tile(intrinsic[None], (len(images), 1))
+
+        self.images = images
+        self.poses = poses
+        self.depths = depths
+        self.intrinsics = intrinsics
+
+    @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)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 5000.0

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/data_readers/tum.py

@@ -0,0 +1,85 @@
+
+import numpy as np
+import torch
+
+import csv
+import os
+import cv2
+import math
+import random
+import json
+import pickle
+import os.path as osp
+
+from lietorch import SE3
+from .stream import RGBDStream
+from .rgbd_utils import loadtum
+
+intrinsics_dict = {
+    'freiburg1': [517.3, 516.5, 318.6, 255.3],
+    'freiburg2': [520.9, 521.0, 325.1, 249.7],
+    'freiburg3': [535.4, 539.2, 320.1, 247.6], 
+}
+
+distortion_dict = {
+    'freiburg1': [0.2624, -0.9531, -0.0054, 0.0026, 1.1633],
+    'freiburg2': [0.2312, -0.7849, -0.0033, -0.0001, 0.9172],
+    'freiburg3': [0, 0, 0, 0, 0],
+}
+
+def as_intrinsics_matrix(intrinsics):
+    K = np.eye(3)
+    K[0,0] = intrinsics[0]
+    K[1,1] = intrinsics[1]
+    K[0,2] = intrinsics[2]
+    K[1,2] = intrinsics[3]
+    return K
+
+
+class TUMStream(RGBDStream):
+    def __init__(self, datapath, **kwargs):
+        super(TUMStream, self).__init__(datapath=datapath, **kwargs)
+
+    def _build_dataset_index(self):
+        """ build list of images, poses, depths, and intrinsics """
+        images, depths, poses, intrinsics = loadtum(self.datapath, self.frame_rate)
+        intrinsic, _ = TUMStream.calib_read(self.datapath)
+        intrinsics = np.tile(intrinsic[None], (len(images), 1))
+
+        # set first pose to identity
+        poses = SE3(torch.as_tensor(poses))
+        poses = poses[[0]].inv() * poses
+        poses = poses.data.cpu().numpy()
+
+        self.images = images
+        self.poses = poses
+        self.depths = depths
+        self.intrinsics = intrinsics
+
+    @staticmethod
+    def calib_read(datapath):
+        if 'freiburg1' in datapath:
+            intrinsic = intrinsics_dict['freiburg1']
+            d_coef = distortion_dict['freiburg1']
+
+        elif 'freiburg2' in datapath:
+            intrinsic = intrinsics_dict['freiburg2']
+            d_coef = distortion_dict['freiburg2']
+
+        elif 'freiburg3' in datapath:
+            intrinsic = intrinsics_dict['freiburg3']
+            d_coef = distortion_dict['freiburg3']
+
+        return np.array(intrinsic), np.array(d_coef)
+
+    @staticmethod
+    def image_read(image_file):
+        intrinsics, d_coef = TUMStream.calib_read(image_file)
+        K = as_intrinsics_matrix(intrinsics)
+        image = cv2.imread(image_file)
+        return cv2.undistort(image, K, d_coef)
+
+    @staticmethod
+    def depth_read(depth_file):
+        depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)
+        return depth.astype(np.float32) / 5000.0

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/ba.py

@@ -0,0 +1,77 @@
+import lietorch
+import torch
+import torch.nn.functional as F
+
+from .chol import block_solve
+import geom.projective_ops as pops
+
+# utility functions for scattering ops
+def safe_scatter_add_mat(H, data, ii, jj, B, M, D):
+    v = (ii >= 0) & (jj >= 0)
+    H.scatter_add_(1, (ii[v]*M + jj[v]).view(1,-1,1,1).repeat(B,1,D,D), data[:,v])
+
+def safe_scatter_add_vec(b, data, ii, B, M, D):
+    v = ii >= 0
+    b.scatter_add_(1, ii[v].view(1,-1,1).repeat(B,1,D), data[:,v])
+
+def MoBA(target, weight, poses, disps, intrinsics, ii, jj, fixedp=1, lm=0.0001, ep=0.1):
+    """ MoBA: Motion Only Bundle Adjustment """
+
+    B, M = poses.shape[:2]
+    D = poses.manifold_dim
+    N = ii.shape[0]
+
+    ### 1: commpute jacobians and residuals ###
+    coords, valid, (Ji, Jj) = pops.projective_transform(
+        poses, disps, intrinsics, ii, jj, jacobian=True)
+
+    r = (target - coords).view(B, N, -1, 1)
+    w = (valid * weight).view(B, N, -1, 1)
+
+    ### 2: construct linear system ###
+    Ji = Ji.view(B, N, -1, D)
+    Jj = Jj.view(B, N, -1, D)
+    wJiT = (.001 * w * Ji).transpose(2,3)
+    wJjT = (.001 * w * Jj).transpose(2,3)
+
+    Hii = torch.matmul(wJiT, Ji)
+    Hij = torch.matmul(wJiT, Jj)
+    Hji = torch.matmul(wJjT, Ji)
+    Hjj = torch.matmul(wJjT, Jj)
+
+    vi = torch.matmul(wJiT, r).squeeze(-1)
+    vj = torch.matmul(wJjT, r).squeeze(-1)
+
+    # only optimize keyframe poses
+    M = M - fixedp
+    ii = ii - fixedp
+    jj = jj - fixedp
+
+    H = torch.zeros(B, M*M, D, D, device=target.device)
+    safe_scatter_add_mat(H, Hii, ii, ii, B, M, D)
+    safe_scatter_add_mat(H, Hij, ii, jj, B, M, D)
+    safe_scatter_add_mat(H, Hji, jj, ii, B, M, D)
+    safe_scatter_add_mat(H, Hjj, jj, jj, B, M, D)
+    H = H.reshape(B, M, M, D, D)
+
+    v = torch.zeros(B, M, D, device=target.device)
+    safe_scatter_add_vec(v, vi, ii, B, M, D)
+    safe_scatter_add_vec(v, vj, jj, B, M, D)
+
+    ### 3: solve the system + apply retraction ###
+    dx = block_solve(H, v, ep=ep, lm=lm)
+    
+    poses1, poses2 = poses[:,:fixedp], poses[:,fixedp:]
+    poses2 = poses2.retr(dx)
+    
+    poses = lietorch.cat([poses1, poses2], dim=1)
+    return poses
+
+def SLessBA(target, weight, poses, disps, intrinsics, ii, jj, fixedp=1):
+    """ Structureless Bundle Adjustment """
+    pass
+
+
+def BA(target, weight, poses, disps, intrinsics, ii, jj, fixedp=1):
+    """ Full Bundle Adjustment """
+    pass

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/chol.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn.functional as F
+import geom.projective_ops as pops
+
+class CholeskySolver(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, H, b):
+        # don't crash training if cholesky decomp fails
+        try:
+            U = torch.cholesky(H)
+            xs = torch.cholesky_solve(b, U)
+            ctx.save_for_backward(U, xs)
+            ctx.failed = False
+        except Exception as e:
+            print(e)
+            ctx.failed = True
+            xs = torch.zeros_like(b)
+
+        return xs
+
+    @staticmethod
+    def backward(ctx, grad_x):
+        if ctx.failed:
+            return None, None
+
+        U, xs = ctx.saved_tensors
+        dz = torch.cholesky_solve(grad_x, U)
+        dH = -torch.matmul(xs, dz.transpose(-1,-2))
+
+        return dH, dz
+
+def block_solve(H, b, ep=0.1, lm=0.0001):
+    """ solve normal equations """
+    B, N, _, D, _ = H.shape
+    I = torch.eye(D).to(H.device)
+    H = H + (ep + lm*H) * I
+
+    H = H.permute(0,1,3,2,4)
+    H = H.reshape(B, N*D, N*D)
+    b = b.reshape(B, N*D, 1)
+
+    x = CholeskySolver.apply(H,b)
+    return x.reshape(B, N, D)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/graph_utils.py

@@ -0,0 +1,106 @@
+
+import torch
+import numpy as np
+from collections import OrderedDict
+
+import lietorch
+from data_readers.rgbd_utils import compute_distance_matrix_flow
+
+
+def graph_to_edge_list(graph):
+    ii, jj, kk = [], [], []
+    for s, u in enumerate(graph):
+        for v in graph[u]:
+            ii.append(u)
+            jj.append(v)
+            kk.append(s)
+
+    ii = torch.as_tensor(ii).cuda()
+    jj = torch.as_tensor(jj).cuda()
+    kk = torch.as_tensor(kk).cuda()
+    return ii, jj, kk
+
+def keyframe_indicies(graph):
+    return torch.as_tensor([u for u in graph]).cuda()
+
+
+def meshgrid(m, n, device='cuda'):
+    ii, jj = torch.meshgrid(torch.arange(m), torch.arange(n))
+    return ii.reshape(-1).to(device), jj.reshape(-1).to(device)
+
+
+class KeyframeGraph:
+    def __init__(self, images, poses, depths, intrinsics):
+        self.images = images.cpu()
+        self.depths = depths.cpu()
+        self.poses = poses
+        self.intrinsics = intrinsics
+
+        depths = depths[..., 3::8, 3::8].float().cuda()
+        disps = torch.where(depths>0.1, 1.0/depths, depths)
+
+        N = poses.shape[1]
+        d = compute_distance_matrix_flow(poses, disps, intrinsics / 8.0)
+
+        i, j = 0, 0
+        ixs = [ i ]
+
+        while j < N-1:
+            if d[i, j+1] > 7.5:
+                ixs += [ j ]
+                i = j
+            j += 1
+
+        # indicies of keyframes
+        self.distance_matrix = d[ixs][:,ixs]
+        self.ixs = np.array(ixs)
+        self.frame_graph = {}
+
+        for i in range(N):
+            k = np.argmin(np.abs(i - self.ixs))
+            j = self.ixs[k]
+            self.frame_graph[i] = (k, poses[:,i] * poses[:,j].inv())
+
+    def get_keyframes(self):
+        ix = torch.as_tensor(self.ixs).cuda()
+        return self.images[:,ix], self.poses[:,ix], self.depths[:,ix], self.intrinsics[:,ix]
+
+    def get_graph(self, num=-1, thresh=24.0, r=2):
+        d = self.distance_matrix.copy()
+        
+        N = d.shape[0]
+        if num < 0:
+            num = N
+
+        graph = OrderedDict()
+        for i in range(N):
+            graph[i] = [j for j in range(N) if i!=j and abs(i-j) <= 2]
+
+        for i in range(N):
+            for j in range(i-r, i+r+1):
+                if j >= 0 and j < N:
+                    d[i,j] = np.inf
+
+        for _ in range(num):
+            ix = np.argmin(d)
+            i, j = ix // N, ix % N
+
+            if d[i,j] < thresh:
+                graph[i].append(j)
+                for ii in range(i-r, i+r+1):
+                    for jj in range(j-r, j+r+1):
+                        if ii>=0 and jj>=0 and ii<N and jj<N:
+                            d[ii,jj] = np.inf
+            else:
+                break
+        
+        return graph
+
+    def get_poses(self, keyframe_poses):
+        
+        poses_list = []
+        for i in range(self.poses.shape[1]):
+            k, dP = self.frame_graph[i]
+            poses_list += [dP * keyframe_poses[:,k]]
+
+        return lietorch.stack(poses_list, 1)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/losses.py

@@ -0,0 +1,68 @@
+import numpy as np
+import torch
+from lietorch import SO3, SE3, Sim3
+from .graph_utils import graph_to_edge_list
+
+def pose_metrics(dE):
+    """ Translation/Rotation/Scaling metrics from Sim3 """
+    t, q, s = dE.data.split([3, 4, 1], -1)
+    ang = SO3(q).log().norm(dim=-1)
+
+    # convert radians to degrees
+    r_err = (180 / np.pi) * ang
+    t_err = t.norm(dim=-1)
+    s_err = (s - 1.0).abs()
+    return r_err, t_err, s_err
+
+def geodesic_loss(Ps, Gs, graph, gamma=0.9):
+    """ Loss function for training network """
+
+    # relative pose
+    ii, jj, kk = graph_to_edge_list(graph)
+    dP = Ps[:,jj] * Ps[:,ii].inv()
+
+    n = len(Gs)
+    geodesic_loss = 0.0
+
+    for i in range(n):
+        w = gamma ** (n - i - 1)
+        dG = Gs[i][:,jj] * Gs[i][:,ii].inv()
+        
+        # pose error
+        d = (dG * dP.inv()).log()
+
+        if isinstance(dG, SE3):
+            tau, phi = d.split([3,3], dim=-1)
+            geodesic_loss += w * (
+                tau.norm(dim=-1).mean() + 
+                phi.norm(dim=-1).mean())
+
+        elif isinstance(dG, Sim3):
+            tau, phi, sig = d.split([3,3,1], dim=-1)
+            geodesic_loss += w * (
+                tau.norm(dim=-1).mean() + 
+                phi.norm(dim=-1).mean() + 
+                0.05 * sig.norm(dim=-1).mean())
+            
+        dE = Sim3(dG * dP.inv()).detach()
+        r_err, t_err, s_err = pose_metrics(dE)
+
+    metrics = {
+        'r_error': r_err.mean().item(),
+        't_error': t_err.mean().item(),
+        's_error': s_err.mean().item(),
+    }
+
+    return geodesic_loss, metrics
+
+
+def residual_loss(residuals, gamma=0.9):
+    """ loss on system residuals """
+    residual_loss = 0.0
+    n = len(residuals)
+
+    for i in range(n):
+        w = gamma ** (n - i - 1)
+        residual_loss += w * residuals[i].abs().mean()
+
+    return residual_loss, {'residual': residual_loss.item()}

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/projective_ops.py

@@ -0,0 +1,114 @@
+import torch
+import torch.nn.functional as F
+
+from lietorch import SE3, Sim3
+
+MIN_DEPTH = 0.1
+
+def extract_intrinsics(intrinsics):
+    return intrinsics[...,None,None,:].unbind(dim=-1)
+
+def iproj(disps, intrinsics):
+    """ pinhole camera inverse projection """
+    ht, wd = disps.shape[2:]
+    fx, fy, cx, cy = extract_intrinsics(intrinsics)
+    
+    y, x = torch.meshgrid(
+        torch.arange(ht).to(disps.device).float(),
+        torch.arange(wd).to(disps.device).float())
+
+    i = torch.ones_like(disps)
+    X = (x - cx) / fx
+    Y = (y - cy) / fy
+    return torch.stack([X, Y, i, disps], dim=-1)
+
+def proj(Xs, intrinsics, jacobian=False):
+    """ pinhole camera projection """
+    fx, fy, cx, cy = extract_intrinsics(intrinsics)
+    X, Y, Z, D = Xs.unbind(dim=-1)
+    d = torch.where(Z.abs() < 0.001, torch.zeros_like(Z), 1.0/Z)
+
+    x = fx * (X * d) + cx
+    y = fy * (Y * d) + cy
+    coords = torch.stack([x,y, D*d], dim=-1)
+
+    if jacobian:
+        B, N, H, W = d.shape
+        o = torch.zeros_like(d)
+        proj_jac = torch.stack([
+             fx*d,     o, -fx*X*d*d,  o,
+                o,  fy*d, -fy*Y*d*d,  o,
+                o,     o,    -D*d*d,  d,
+        ], dim=-1).view(B, N, H, W, 3, 4)
+
+        return coords, proj_jac
+
+    return coords, None
+
+def actp(Gij, X0, jacobian=False):
+    """ action on point cloud """
+    X1 = Gij[:,:,None,None] * X0
+    
+    if jacobian:
+        X, Y, Z, d = X1.unbind(dim=-1)
+        o = torch.zeros_like(d)
+        B, N, H, W = d.shape
+
+        if isinstance(Gij, SE3):
+            Ja = torch.stack([
+                d,  o,  o,  o,  Z, -Y,
+                o,  d,  o, -Z,  o,  X, 
+                o,  o,  d,  Y, -X,  o,
+                o,  o,  o,  o,  o,  o,
+            ], dim=-1).view(B, N, H, W, 4, 6)
+
+        elif isinstance(Gij, Sim3):
+            Ja = torch.stack([
+                d,  o,  o,  o,  Z, -Y,  X,
+                o,  d,  o, -Z,  o,  X,  Y,
+                o,  o,  d,  Y, -X,  o,  Z,
+                o,  o,  o,  o,  o,  o,  o
+            ], dim=-1).view(B, N, H, W, 4, 7)
+
+        return X1, Ja
+
+    return X1, None
+
+def projective_transform(poses, depths, intrinsics, ii, jj, jacobian=False):
+    """ map points from ii->jj """
+
+    # inverse project (pinhole)
+    X0 = iproj(depths[:,ii], intrinsics[:,ii])
+    
+    # transform
+    Gij = poses[:,jj] * poses[:,ii].inv()
+    X1, Ja = actp(Gij, X0, jacobian=jacobian)
+    
+    # project (pinhole)
+    x1, Jp = proj(X1, intrinsics[:,jj], jacobian=jacobian)
+
+    # exclude points too close to camera
+    valid = ((X1[...,2] > MIN_DEPTH) & (X0[...,2] > MIN_DEPTH)).float()
+    valid = valid.unsqueeze(-1)
+
+    if jacobian:
+        Jj = torch.matmul(Jp, Ja)
+        Ji = -Gij[:,:,None,None,None].adjT(Jj)
+        return x1, valid, (Ji, Jj)
+
+    return x1, valid
+
+
+def induced_flow(poses, disps, intrinsics, ii, jj):
+    """ optical flow induced by camera motion """
+
+    ht, wd = disps.shape[2:]
+    y, x = torch.meshgrid(
+        torch.arange(ht).to(disps.device).float(),
+        torch.arange(wd).to(disps.device).float())
+
+    coords0 = torch.stack([x, y], dim=-1)
+    coords1, valid = projective_transform(poses, disps, intrinsics, ii, jj)
+
+    return coords1[...,:2] - coords0, valid
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/geom/sampler_utils.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn.functional as F
+
+def _bilinear_sampler(img, coords, mode='bilinear', mask=False):
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1,1], dim=-1)
+    xgrid = 2*xgrid/(W-1) - 1
+    ygrid = 2*ygrid/(H-1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True)
+
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+
+    return img
+
+def bilinear_sampler(img, coords):
+    """ Wrapper for bilinear sampler for inputs with extra batch dimensions """
+    unflatten = False
+    if len(img.shape) == 5:
+        unflatten = True
+        b, n, c, h, w = img.shape
+        img = img.view(b*n, c, h, w)
+        coords = coords.view(b*n, h, w, 2)
+
+    img1 = _bilinear_sampler(img, coords)
+    
+    if unflatten:
+        return img1.view(b, n, c, h, w)
+
+    return img1
+
+def sample_depths(depths, coords):
+    batch, num, ht, wd = depths.shape
+    depths = depths.view(batch, num, 1, ht, wd)
+    coords = coords.view(batch, num, ht, wd, 2)
+    
+    depths_proj = bilinear_sampler(depths, coords)
+    return depths_proj.view(batch, num, ht, wd, 1)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/logger.py

@@ -0,0 +1,46 @@
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+
+
+SUM_FREQ = 100
+
+class Logger:
+    def __init__(self, name, scheduler):
+        self.total_steps = 0
+        self.running_loss = {}
+        self.writer = None
+        self.name = name
+        self.scheduler = scheduler
+
+    def _print_training_status(self):
+        if self.writer is None:
+            self.writer = SummaryWriter('runs/%s' % self.name)
+            print([k for k in self.running_loss])
+
+        lr = self.scheduler.get_lr().pop()
+        metrics_data = [self.running_loss[k]/SUM_FREQ for k in self.running_loss.keys()]
+        training_str = "[{:6d}, {:10.7f}] ".format(self.total_steps+1, lr)
+        metrics_str = ("{:10.4f}, "*len(metrics_data)).format(*metrics_data)
+        
+        # print the training status
+        print(training_str + metrics_str)
+
+        for key in self.running_loss:
+            val = self.running_loss[key] / SUM_FREQ
+            self.writer.add_scalar(key, val, self.total_steps)
+            self.running_loss[key] = 0.0
+
+    def push(self, metrics):
+
+        for key in metrics:
+            if key not in self.running_loss:
+                self.running_loss[key] = 0.0
+
+            self.running_loss[key] += metrics[key]
+
+        if self.total_steps % SUM_FREQ == SUM_FREQ-1:
+            self._print_training_status()
+            self.running_loss = {}
+
+        self.total_steps += 1

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/modules/clipping.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+GRAD_CLIP = .01
+
+class GradClip(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_x):
+        o = torch.zeros_like(grad_x)
+        grad_x = torch.where(grad_x.abs()>GRAD_CLIP, o, grad_x)
+        grad_x = torch.where(torch.isnan(grad_x), o, grad_x)
+        return grad_x
+
+class GradientClip(nn.Module):
+    def __init__(self):
+        super(GradientClip, self).__init__()
+
+    def forward(self, x):
+        return GradClip.apply(x)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/modules/corr.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn.functional as F
+
+from geom.sampler_utils import bilinear_sampler
+import lietorch_extras
+
+
+class CorrSampler(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, volume, coords, radius):
+        ctx.save_for_backward(volume,coords)
+        ctx.radius = radius
+        corr, = lietorch_extras.corr_index_forward(volume, coords, radius)
+        return corr
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        volume, coords = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_volume, = lietorch_extras.corr_index_backward(volume, coords, grad_output, ctx.radius)
+        return grad_volume, None, None
+
+
+class CorrBlock:
+    def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
+        self.num_levels = num_levels
+        self.radius = radius
+        self.corr_pyramid = []
+
+        # all pairs correlation
+        corr = CorrBlock.corr(fmap1, fmap2)
+
+        batch, num, h1, w1, h2, w2 = corr.shape
+        corr = corr.reshape(batch*num*h1*w1, 1, h2, w2)
+        
+        for i in range(self.num_levels):
+            self.corr_pyramid.append(
+                corr.view(batch*num, h1, w1, h2//2**i, w2//2**i))
+            corr = F.avg_pool2d(corr, 2, stride=2)
+            
+    def __call__(self, coords):
+        out_pyramid = []
+        batch, num, ht, wd, _ = coords.shape
+        coords = coords.permute(0,1,4,2,3)
+        coords = coords.contiguous().view(batch*num, 2, ht, wd)
+        
+        for i in range(self.num_levels):
+            corr = CorrSampler.apply(self.corr_pyramid[i], coords/2**i, self.radius)
+            out_pyramid.append(corr.view(batch, num, -1, ht, wd))
+
+        return torch.cat(out_pyramid, dim=2)
+
+    def append(self, other):
+        for i in range(self.num_levels):
+            self.corr_pyramid[i] = torch.cat([self.corr_pyramid[i], other.corr_pyramid[i]], 0)
+
+    def remove(self, ix):
+        for i in range(self.num_levels):
+            self.corr_pyramid[i] = self.corr_pyramid[i][ix].contiguous()
+
+    @staticmethod
+    def corr(fmap1, fmap2):
+        """ all-pairs correlation """
+        batch, num, dim, ht, wd = fmap1.shape
+        fmap1 = fmap1.reshape(batch*num, dim, ht*wd) / 4.0
+        fmap2 = fmap2.reshape(batch*num, dim, ht*wd) / 4.0
+        
+        corr = torch.matmul(fmap1.transpose(1,2), fmap2)
+        return corr.view(batch, num, ht, wd, ht, wd)
+
+
+class CorrLayer(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, fmap1, fmap2, coords, r):
+        ctx.r = r
+        fmap1 = fmap1.contiguous()
+        fmap2 = fmap2.contiguous()
+        coords = coords.contiguous()
+        ctx.save_for_backward(fmap1, fmap2, coords)
+        corr, = lietorch_extras.altcorr_forward(fmap1, fmap2, coords, ctx.r)
+        return corr
+
+    @staticmethod
+    def backward(ctx, grad_corr):
+        fmap1, fmap2, coords = ctx.saved_tensors
+        grad_corr = grad_corr.contiguous()
+        fmap1_grad, fmap2_grad, coords_grad = \
+            lietorch_extras.altcorr_backward(fmap1, fmap2, coords, grad_corr, ctx.r)
+        return fmap1_grad, fmap2_grad, coords_grad, None
+
+
+
+class AltCorrBlock:
+    def __init__(self, fmaps, inds, num_levels=4, radius=3):
+        self.num_levels = num_levels
+        self.radius = radius
+        self.inds = inds
+
+        B, N, C, H, W = fmaps.shape
+        fmaps = fmaps.view(B*N, C, H, W)
+        
+        self.pyramid = []
+        for i in range(self.num_levels):
+            sz = (B, N, H//2**i, W//2**i, C)
+            fmap_lvl = fmaps.permute(0, 2, 3, 1)
+            self.pyramid.append(fmap_lvl.reshape(*sz))
+            fmaps = F.avg_pool2d(fmaps, 2, stride=2)
+  
+    def corr_fn(self, coords, ii, jj):
+        B, N, H, W, S, _ = coords.shape
+        coords = coords.permute(0, 1, 4, 2, 3, 5)
+
+        corr_list = []
+        for i in range(self.num_levels):
+            r = self.radius
+            fmap1_i = self.pyramid[0][:, ii]
+            fmap2_i = self.pyramid[i][:, jj]
+
+            coords_i = (coords / 2**i).reshape(B*N, S, H, W, 2).contiguous()
+            fmap1_i = fmap1_i.reshape((B*N,) + fmap1_i.shape[2:])
+            fmap2_i = fmap2_i.reshape((B*N,) + fmap2_i.shape[2:])
+
+            corr = CorrLayer.apply(fmap1_i, fmap2_i, coords_i, self.radius)
+            corr = corr.view(B, N, S, -1, H, W).permute(0, 1, 3, 4, 5, 2)
+            corr_list.append(corr)
+
+        corr = torch.cat(corr_list, dim=2)
+        return corr / 16.0
+
+
+    def __call__(self, coords, ii, jj):
+        squeeze_output = False
+        if len(coords.shape) == 5:
+            coords = coords.unsqueeze(dim=-2)
+            squeeze_output = True
+
+        corr = self.corr_fn(coords, ii, jj)
+        
+        if squeeze_output:
+            corr = corr.squeeze(dim=-1)
+
+        return corr.contiguous()
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/modules/extractor.py

@@ -0,0 +1,271 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(ResidualBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(BottleneckBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
+        self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
+        self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes//4)
+            self.norm2 = nn.BatchNorm2d(planes//4)
+            self.norm3 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes//4)
+            self.norm2 = nn.InstanceNorm2d(planes//4)
+            self.norm3 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            self.norm3 = nn.Sequential()
+            if not stride == 1:
+                self.norm4 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+        y = self.relu(self.norm3(self.conv3(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+DIM=32
+
+class BasicEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0, multidim=False):
+        super(BasicEncoder, self).__init__()
+        self.norm_fn = norm_fn
+        self.multidim = multidim
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=DIM)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(DIM)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(DIM)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, DIM, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = DIM
+        self.layer1 = self._make_layer(DIM,  stride=1)
+        self.layer2 = self._make_layer(2*DIM, stride=2)
+        self.layer3 = self._make_layer(4*DIM, stride=2)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(4*DIM, output_dim, kernel_size=1)
+
+        if self.multidim:
+            self.layer4 = self._make_layer(256, stride=2)
+            self.layer5 = self._make_layer(512, stride=2)
+
+            self.in_planes = 256
+            self.layer6 = self._make_layer(256, stride=1)
+
+            self.in_planes = 128
+            self.layer7 = self._make_layer(128, stride=1)
+
+            self.up1 = nn.Conv2d(512, 256, 1)
+            self.up2 = nn.Conv2d(256, 128, 1)
+            self.conv3 = nn.Conv2d(128, output_dim, kernel_size=1)
+
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+        else:
+            self.dropout = None
+
+        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.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        b, n, c1, h1, w1 = x.shape
+        x = x.view(b*n, c1, h1, w1)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.conv2(x)
+
+        _, c2, h2, w2 = x.shape
+        return x.view(b, n, c2, h2, w2)
+
+
+
+class BasicEncoder16(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0, multidim=False):
+        super(BasicEncoder16, self).__init__()
+        self.norm_fn = norm_fn
+        self.multidim = multidim
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=DIM)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(DIM)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(DIM)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, DIM, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = DIM
+        self.layer1 = self._make_layer(DIM,  stride=1)
+        self.layer2 = self._make_layer(2*DIM, stride=2)
+        self.layer3 = self._make_layer(4*DIM, stride=2)
+        self.layer4 = self._make_layer(4*DIM, stride=2)
+
+
+        # output convolution
+        self.conv2 = nn.Conv2d(4*DIM, output_dim, kernel_size=1)
+
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+        else:
+            self.dropout = None
+
+        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.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        b, n, c1, h1, w1 = x.shape
+        x = x.view(b*n, c1, h1, w1)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.conv2(x)
+
+        _, c2, h2, w2 = x.shape
+        return x.view(b, n, c2, h2, w2)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/modules/gru.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+
+class ConvGRU(nn.Module):
+    def __init__(self, h_planes=128, i_planes=128):
+        super(ConvGRU, self).__init__()
+        self.do_checkpoint = False
+        self.convz = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
+        self.convr = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
+        self.convq = nn.Conv2d(h_planes+i_planes, h_planes, 3, padding=1)
+
+    def forward(self, net, *inputs):
+        inp = torch.cat(inputs, dim=1)
+        net_inp = torch.cat([net, inp], dim=1)
+
+        z = torch.sigmoid(self.convz(net_inp))
+        r = torch.sigmoid(self.convr(net_inp))
+        q = torch.tanh(self.convq(torch.cat([r*net, inp], dim=1)))
+
+        net = (1-z) * net + z * q
+        return net

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/modules/unet.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+
+# Unet model from https://github.com/usuyama/pytorch-unet
+
+
+
+GRAD_CLIP = .01
+
+class GradClip(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_x):
+        o = torch.zeros_like(grad_x)
+        grad_x = torch.where(grad_x.abs()>GRAD_CLIP, o, grad_x)
+        grad_x = torch.where(torch.isnan(grad_x), o, grad_x)
+        return grad_x
+
+
+def double_conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 5, padding=2),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(out_channels, out_channels, 5, padding=2),
+        nn.ReLU(inplace=True)
+    )
+
+
+class UNet(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+                
+        self.dconv_down1 = double_conv(128, 128)
+        self.dconv_down2 = double_conv(128, 256)
+        self.dconv_down3 = double_conv(256, 256)
+        # self.dconv_down4 = double_conv(256, 512)        
+
+        self.maxpool = nn.AvgPool2d(2)
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)        
+        
+        self.dconv_up3 = double_conv(256 + 256, 256)
+        self.dconv_up2 = double_conv(256 + 256, 128)
+        self.dconv_up1 = double_conv(128 + 128, 128)
+
+        self.conv_r = nn.Conv2d(128, 3, 1)
+        self.conv_w = nn.Conv2d(128, 3, 1)
+
+
+    def forward(self, x):
+        b, n, c, ht, wd = x.shape
+        x = x.view(b*n, c, ht, wd)
+
+        conv1 = self.dconv_down1(x)
+        x = self.maxpool(conv1)
+
+        conv2 = self.dconv_down2(x)
+        x = self.maxpool(conv2)
+        
+        conv3 = self.dconv_down3(x)
+        x = torch.cat([x, conv3], dim=1)
+        
+        x = self.dconv_up3(x)
+        x = self.upsample(x)        
+        x = torch.cat([x, conv2], dim=1)       
+
+        x = self.dconv_up2(x)
+        x = self.upsample(x)        
+        x = torch.cat([x, conv1], dim=1)   
+        
+        x = self.dconv_up1(x)
+        r = self.conv_r(x)
+        w = self.conv_w(x)
+
+        w = torch.sigmoid(w)
+        w = w.view(b, n, 3, ht, wd).permute(0,1,3,4,2)
+        r = r.view(b, n, 3, ht, wd).permute(0,1,3,4,2)
+
+        # w = GradClip.apply(w)
+        # r = GradClip.apply(r)
+        return r, w

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/rslam.py

@@ -0,0 +1,143 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+from networks.modules.extractor import BasicEncoder
+from networks.modules.corr import CorrBlock
+from networks.modules.gru import ConvGRU
+from networks.modules.clipping import GradientClip
+
+from lietorch import SE3, Sim3
+from geom.ba import MoBA
+
+import geom.projective_ops as pops
+from geom.sampler_utils import bilinear_sampler, sample_depths
+from geom.graph_utils import graph_to_edge_list, keyframe_indicies
+
+
+class UpdateModule(nn.Module):
+    def __init__(self, args):
+        super(UpdateModule, self).__init__()
+        self.args = args
+
+        cor_planes = 4 * (2*3 + 1)**2 + 1
+
+        self.corr_encoder = nn.Sequential(
+            nn.Conv2d(cor_planes, 128, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True))
+
+        self.flow_encoder = nn.Sequential(
+            nn.Conv2d(3, 128, 7, padding=3),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 64, 3, padding=1),
+            nn.ReLU(inplace=True))
+
+        self.weight = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 3, 3, padding=1),
+            GradientClip(),
+            nn.Sigmoid())
+
+        self.delta = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 3, 3, padding=1),
+            GradientClip())
+
+        self.gru = ConvGRU(128, 128+128+64)
+
+    def forward(self, net, inp, corr, flow):
+        """ RaftSLAM update operator """
+
+        batch, num, ch, ht, wd = net.shape
+        output_dim = (batch, num, -1, ht, wd)
+        net = net.view(batch*num, -1, ht, wd)
+        inp = inp.view(batch*num, -1, ht, wd)        
+        corr = corr.view(batch*num, -1, ht, wd)
+        flow = flow.view(batch*num, -1, ht, wd)
+
+        corr = self.corr_encoder(corr)
+        flow = self.flow_encoder(flow)
+        net = self.gru(net, inp, corr, flow)
+
+        ### update variables ###
+        delta = self.delta(net).view(*output_dim)
+        weight = self.weight(net).view(*output_dim)
+
+        delta = delta.permute(0,1,3,4,2).contiguous()
+        weight = weight.permute(0,1,3,4,2).contiguous()
+
+        net = net.view(*output_dim)
+        return net, delta, weight
+
+
+class RaftSLAM(nn.Module):
+    def __init__(self, args):
+        super(RaftSLAM, self).__init__()
+        self.args = args
+        self.fnet = BasicEncoder(output_dim=128, norm_fn='instance')
+        self.cnet = BasicEncoder(output_dim=256, norm_fn='none')
+        self.update = UpdateModule(args)
+
+    def extract_features(self, images):
+        """ run feeature extraction networks """
+        fmaps = self.fnet(images)
+        net = self.cnet(images)
+        
+        net, inp = net.split([128,128], dim=2)
+        net = torch.tanh(net)
+        inp = torch.relu(inp)
+        return fmaps, net, inp
+
+    def forward(self, Gs, images, depths, intrinsics, graph=None, num_steps=12):
+        """ Estimates SE3 or Sim3 between pair of frames """
+
+        u = keyframe_indicies(graph)
+        ii, jj, kk = graph_to_edge_list(graph)
+
+        depths = depths[:, :, 3::8, 3::8]
+        intrinsics = intrinsics / 8
+        mask = (depths > 0.1).float()
+        disps = torch.where(depths>0.1, 1.0/depths, depths)
+
+        fmaps, net, inp = self.extract_features(images)
+        net, inp = net[:,ii], inp[:,ii]
+        corr_fn = CorrBlock(fmaps[:,ii], fmaps[:,jj], num_levels=4, radius=3)
+
+        coords, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+        residual = torch.zeros_like(coords[...,:2])
+
+        Gs_list, coords_list, residual_list = [], [], []
+        for step in range(num_steps):
+            Gs = Gs.detach()
+            coords = coords.detach()
+            residual = residual.detach()
+
+            corr = corr_fn(coords[...,:2])
+            flow = residual.permute(0,1,4,2,3).clamp(-32.0, 32.0)
+            
+            corr = torch.cat([corr, mask[:,ii,None]], dim=2)
+            flow = torch.cat([flow, mask[:,ii,None]], dim=2)
+            net, delta, weight = self.update(net, inp, corr, flow)
+
+            target = coords + delta
+            weight[...,2] = 0.0
+
+            for i in range(3):
+                Gs = MoBA(target, weight, Gs, disps, intrinsics, ii, jj)
+
+            coords, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+            residual = (target - coords)[...,:2]
+
+            Gs_list.append(Gs)
+            coords_list.append(target)
+
+            valid_mask = valid_mask * mask[:,ii].unsqueeze(-1)
+            residual_list.append(valid_mask * residual)
+
+        return Gs_list, residual_list

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/sim3_net.py

@@ -0,0 +1,135 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+from networks.modules.extractor import BasicEncoder
+from networks.modules.corr import CorrBlock
+from networks.modules.gru import ConvGRU
+from networks.modules.clipping import GradientClip
+
+from lietorch import SE3, Sim3
+from geom.ba import MoBA
+
+import geom.projective_ops as pops
+from geom.sampler_utils import bilinear_sampler, sample_depths
+from geom.graph_utils import graph_to_edge_list, keyframe_indicies
+
+
+class UpdateModule(nn.Module):
+    def __init__(self, args):
+        super(UpdateModule, self).__init__()
+        self.args = args
+
+        cor_planes = 4 * (2*3 + 1)**2
+
+        self.encoder = nn.Sequential(
+            nn.Conv2d(cor_planes, 128, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True))
+
+        self.weight = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 3, 3, padding=1),
+            GradientClip(),
+            nn.Sigmoid())
+
+        self.delta = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 3, 3, padding=1),
+            GradientClip())
+
+        self.gru = ConvGRU(128, 128+128+1)
+
+    def forward(self, net, inp, corr, dz):
+        """ update operator """
+
+        batch, num, ch, ht, wd = net.shape
+        output_dim = (batch, num, -1, ht, wd)
+        net = net.view(batch*num, -1, ht, wd)
+        inp = inp.view(batch*num, -1, ht, wd)        
+
+        corr = corr.view(batch*num, -1, ht, wd)
+        dz = dz.view(batch*num, 1, ht, wd)
+        corr = self.encoder(corr)
+        net = self.gru(net, inp, corr, dz)
+
+        ### update variables ###
+        delta = self.delta(net).view(*output_dim)
+        weight = self.weight(net).view(*output_dim)
+
+        delta = delta.permute(0,1,3,4,2).contiguous()
+        weight = weight.permute(0,1,3,4,2).contiguous()
+
+        net = net.view(*output_dim)
+        return net, delta, weight
+
+
+class Sim3Net(nn.Module):
+    def __init__(self, args):
+        super(Sim3Net, self).__init__()
+        self.args = args
+        self.fnet = BasicEncoder(output_dim=128, norm_fn='instance')
+        self.cnet = BasicEncoder(output_dim=256, norm_fn='none')
+        self.update = UpdateModule(args)
+
+    def extract_features(self, images):
+        """ run feeature extraction networks """
+        fmaps = self.fnet(images)
+        net = self.cnet(images)
+        
+        net, inp = net.split([128,128], dim=2)
+        net = torch.tanh(net)
+        inp = torch.relu(inp)
+        return fmaps, net, inp
+
+    def forward(self, Gs, images, depths, intrinsics, graph=None, num_steps=12):
+        """ Estimates SE3 or Sim3 between pair of frames """
+
+        if graph is None:
+            graph = OrderedDict()
+            graph[0] = [1]
+            graph[1] = [0]
+
+        u = keyframe_indicies(graph)
+        ii, jj, kk = graph_to_edge_list(graph)
+
+        # use inverse depth parameterization
+        depths = depths.clamp(min=0.1, max=1000.0)
+        disps = 1.0 / depths[:, :, 3::8, 3::8]
+        intrinsics = intrinsics / 8.0
+
+        fmaps, net, inp = self.extract_features(images)
+        corr_fn = CorrBlock(fmaps[:,ii], fmaps[:,jj], num_levels=4, radius=3)
+
+        Gs_list, coords_list, residual_list = [], [], []
+        for step in range(num_steps):
+            Gs = Gs.detach()
+            coords1_xyz, _ = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+
+            coords1, zinv_proj = coords1_xyz.split([2,1], dim=-1)
+            zinv = sample_depths(disps[:,jj], coords1)
+            dz = (zinv - zinv_proj).clamp(-1.0, 1.0)
+
+            corr = corr_fn(coords1)
+            net, delta, weight = self.update(net, inp, corr, dz)
+
+            target = coords1_xyz + delta
+            for i in range(3):
+                Gs = MoBA(target, weight, Gs, disps, intrinsics, ii, jj)
+
+            coords1_xyz, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+            residual = valid_mask * (target - coords1_xyz)
+
+            Gs_list.append(Gs)
+            coords_list.append(target)
+            residual_list.append(residual)
+
+        return Gs_list, residual_list
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/core/networks/slam_system.py

@@ -0,0 +1,340 @@
+import numpy as np
+import torch
+from collections import OrderedDict
+from torch.cuda.amp import autocast
+
+import matplotlib.pyplot as plt
+import lietorch
+from lietorch import SE3, Sim3
+
+from geom.ba import MoBA
+from .modules.corr import CorrBlock, AltCorrBlock
+from .rslam import RaftSLAM
+
+import geom.projective_ops as pops
+from geom.sampler_utils import bilinear_sampler
+from geom.graph_utils import KeyframeGraph, graph_to_edge_list
+
+def meshgrid(m, n, device='cuda'):
+    ii, jj = torch.meshgrid(torch.arange(m), torch.arange(n))
+    return ii.reshape(-1).to(device), jj.reshape(-1).to(device)
+
+def normalize_images(images):
+    images = images[:, :, [2,1,0]]
+    mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device)
+    std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device)
+    return (images/255.0).sub_(mean[:, None, None]).div_(std[:, None, None])
+
+class FactorGraph:
+    def __init__(self, hidden=None, inputs=None, residu=None, ii=None, jj=None):
+        self.hidden = hidden
+        self.inputs = inputs
+        self.residu = residu
+        self.ii = ii
+        self.jj = jj
+
+    def __iadd__(self, other):
+        if self.hidden is None:
+            self.hidden = other.hidden
+            self.inputs = other.inputs
+            self.residu = other.residu
+            self.ii = other.ii
+            self.jj = other.jj
+        else:
+            self.hidden = torch.cat([self.hidden, other.hidden], 1)
+            self.inputs = torch.cat([self.inputs, other.inputs], 1)
+            self.residu = torch.cat([self.residu, other.residu], 1)
+            self.ii = torch.cat([self.ii, other.ii], 0)
+            self.jj = torch.cat([self.jj, other.jj], 0)
+
+        return self
+
+    def rm(self, keep):
+        self.hidden = self.hidden[:,keep]
+        self.inputs = self.inputs[:,keep]
+        self.residu = self.residu[:,keep]
+        self.ii = self.ii[keep]
+        self.jj = self.jj[keep]
+
+
+class SLAMSystem(RaftSLAM):
+    def __init__(self, args):
+        super(SLAMSystem, self).__init__(args)
+        self.mem = 32
+        self.num_keyframes = 5
+
+        self.frontend = None
+        self.factors = FactorGraph()
+        self.count = 0
+        self.fixed_poses = 1
+
+        self.images_list = []
+        self.depths_list = []
+        self.intrinsics_list = []
+        
+    def initialize(self, ht, wd):
+        """ initialize slam buffers """
+
+        self.ht, self.wd = ht, wd
+        ht, wd = ht // 8, wd // 8
+
+        self.fmaps = torch.zeros(1, self.mem, 128, ht, wd, device='cuda', dtype=torch.half)
+        self.nets = torch.zeros(1, self.mem, 128, ht, wd, device='cuda', dtype=torch.half)
+        self.inps = torch.zeros(1, self.mem, 128, ht, wd, device='cuda', dtype=torch.half)
+
+        self.poses = SE3.Identity(1, 2048, device='cuda')
+        self.disps = torch.ones(1, 2048, ht, wd, device='cuda')
+        self.intrinsics = torch.zeros(1, 2048, 4, device='cuda')
+        self.tstamps = torch.zeros(2048, dtype=torch.long)
+
+    def set_frontend(self, frontend):
+        self.frontend = frontend
+
+    def add_point_cloud(self, index, image, pose, depth, intrinsics, s=8):
+        """ add point cloud to visualization """
+
+        if self.frontend is None:
+            return -1
+
+        image = image[...,s//2::s,s//2::s]
+        depth = depth[...,s//2::s,s//2::s]
+        intrinsics = intrinsics / s
+
+        # backproject
+        points = pops.iproj(1.0/depth[None], intrinsics[None])
+        points = points[...,:3] / points[...,[3]]
+        
+        points = points.reshape(-1, 3)
+        valid = (depth > 0).reshape(-1)
+        colors = image.reshape(3,-1).t() / 255.0
+        
+        point_data = points[valid].cpu().numpy()
+        color_data = colors[valid].cpu().numpy()
+        color_data = color_data[:, [2,1,0]]
+
+        pose_data = pose.inv()[0].data
+        self.frontend.update_pose(index, pose_data)
+        self.frontend.update_points(index, point_data, color_data)
+
+    def get_keyframes(self):
+        """ return keyframe poses and timestamps """
+        return self.poses[0, :self.count], self.tstamps[:self.count]
+
+    def raw_poses(self):
+        return self.poses[0, :self.count].inv().data
+
+    def add_keyframe(self, tstamp, image, depth, intrinsics):
+        """ add keyframe to factor graph """
+
+        if self.count == 0:
+            ht, wd = image.shape[3:]
+            self.initialize(ht, wd)
+
+        inputs = normalize_images(image)
+        with autocast(enabled=True):
+            fmaps, net, inp = self.extract_features(inputs)
+
+        ix = self.count % self.mem
+        self.fmaps[:, ix] = fmaps.squeeze(1)
+        self.nets[:, ix] = net.squeeze(1)
+        self.inps[:, ix] = inp.squeeze(1)
+
+        self.tstamps[self.count] = tstamp
+        self.intrinsics[:, self.count] = intrinsics / 8.0
+
+        disp = torch.where(depth > 0, 1.0/depth, depth)
+        self.disps[:, self.count] = disp[:,3::8,3::8]
+
+        pose = self.poses[:, self.count-1]
+        self.add_point_cloud(self.count, image, pose, depth, intrinsics)
+        self.count += 1
+        
+    def get_node_attributes(self, index):
+        index = index % self.mem
+        return self.fmaps[:, index], self.nets[:, index], self.inps[:, index]
+
+    def add_factors(self, ii, jj):
+        """ add factors to slam graph """
+        fmaps, hidden, inputs = self.get_node_attributes(ii)
+        residu_shape = (1, ii.shape[0], self.ht//8, self.wd//8, 2)
+        residu = torch.zeros(*residu_shape).cuda()
+        self.factors += FactorGraph(hidden, inputs, residu, ii, jj)
+
+    def transform_project(self, ii, jj, **kwargs):
+        """ helper function, compute project transform """
+        return pops.projective_transform(self.poses, self.disps, self.intrinsics, ii, jj, **kwargs)
+
+    def moba(self, num_steps=5, is_init=False):
+        """ motion only bundle adjustment """
+
+        ii, jj = self.factors.ii, self.factors.jj
+        ixs = torch.cat([ii, jj], 0)
+
+        with autocast(enabled=True):
+            fmap1 = self.fmaps[:, ii % self.mem]
+            fmap2 = self.fmaps[:, jj % self.mem]
+            poses = self.poses[:, :jj.max()+1]
+
+            corr_fn = CorrBlock(fmap1, fmap2, num_levels=4, radius=3)
+            mask = (self.disps[:,ii] > 0.01).float()
+
+            with autocast(enabled=False):
+                coords, valid_mask = pops.projective_transform(poses, self.disps, self.intrinsics, ii, jj)
+
+            for i in range(num_steps):
+                corr = corr_fn(coords[...,:2])
+                corr = torch.cat([corr, mask[:,:,None]], dim=2)
+
+                with autocast(enabled=False):
+                    flow = self.factors.residu.permute(0,1,4,2,3).clamp(-32.0, 32.0)
+                    flow = torch.cat([flow, mask[:,:,None]], dim=2)
+                
+                self.factors.hidden, delta, weight = \
+                    self.update(self.factors.hidden, self.factors.inputs, corr, flow)
+
+                with autocast(enabled=False):
+                    target = coords + delta
+                    weight[...,2] = 0.0
+
+                    for i in range(3):
+                        poses = MoBA(target, weight, poses, self.disps, 
+                            self.intrinsics, ii, jj, self.fixed_poses)
+
+                    coords, valid_mask = pops.projective_transform(poses, self.disps, self.intrinsics, ii, jj)
+                    self.factors.residu = (target - coords)[...,:2]
+
+        self.poses[:, :jj.max()+1] = poses
+
+        # update visualization
+        if self.frontend is not None:
+            for ix in ixs.cpu().numpy():
+                self.frontend.update_pose(ix, self.poses[:,ix].inv()[0].data)
+
+
+    def track(self, tstamp, image, depth, intrinsics):
+        """ main thread """
+
+        self.images_list.append(image)
+        self.depths_list.append(depth)
+        self.intrinsics_list.append(intrinsics)
+
+        # collect frames for initialization
+        if self.count < self.num_keyframes:
+            self.add_keyframe(tstamp, image, depth, intrinsics)
+
+            if self.count == self.num_keyframes:
+                ii, jj = meshgrid(self.num_keyframes, self.num_keyframes)
+                keep = ((ii - jj).abs() > 0) & (((ii - jj).abs() <= 3))
+
+                self.add_factors(ii[keep], jj[keep])
+                self.moba(num_steps=8, is_init=True)
+
+        else:
+            self.poses[:,self.count] = self.poses[:,self.count-1]
+            self.add_keyframe(tstamp, image, depth, intrinsics)
+
+            N = self.count
+            ii = torch.as_tensor([N-3, N-2, N-1, N-1, N-1], device='cuda')
+            jj = torch.as_tensor([N-1, N-1, N-2, N-3, N-4], device='cuda')
+            
+            self.add_factors(ii, jj)
+            self.moba(num_steps=4)
+
+            self.fixed_poses += 1
+            self.factors.rm(self.factors.ii + 2 >= self.fixed_poses)
+
+
+    def forward(self, poses, images, depths, intrinsics, num_steps=12):
+        """ Estimates SE3 or Sim3 between pair of frames """
+
+        keyframe_graph = KeyframeGraph(images, poses, depths, intrinsics)
+        images, Gs, depths, intrinsics = keyframe_graph.get_keyframes()
+
+        images = images.cuda()
+        depths = depths.cuda()
+
+        if self.frontend is not None:
+            self.frontend.reset()
+            for i, ix in enumerate(keyframe_graph.ixs):
+                self.add_point_cloud(ix, images[:,i], Gs[:,i], depths[:,i], intrinsics[:,i], s=4)
+            for i in range(poses.shape[1]):
+                self.frontend.update_pose(i, poses[:,i].inv()[0].data)
+
+        graph = keyframe_graph.get_graph()
+        ii, jj, kk = graph_to_edge_list(graph)
+        ixs = torch.cat([ii, jj], 0)
+
+        images = normalize_images(images.cuda())
+        depths = depths[:, :, 3::8, 3::8].cuda()
+        mask = (depths > 0.1).float()
+        disps = torch.where(depths>0.1, 1.0/depths, depths)
+        intrinsics = intrinsics / 8
+
+        with autocast(True):
+
+            fmaps, net, inp = self.extract_features(images)
+            net = net[:,ii]
+            
+            # alternate corr implementation uses less memory but 4x slower
+            corr_fn = AltCorrBlock(fmaps.float(), (ii, jj), num_levels=4, radius=3)
+            # corr_fn = CorrBlock(fmaps[:,ii], fmaps[:,jj], num_levels=4, radius=3)
+
+            with autocast(False):
+                coords, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+                residual = torch.zeros_like(coords[...,:2])
+
+            for step in range(num_steps):
+                print("Global refinement iteration #{}".format(step))                
+                net_list = []
+                targets_list = []
+                weights_list = []
+
+                s = 64
+                for i in range(0, ii.shape[0], s):
+                    ii1 = ii[i:i+s]
+                    jj1 = jj[i:i+s]
+
+                    corr1 = corr_fn(coords[:,i:i+s,:,:,:2], ii1, jj1)
+                    flow1 = residual[:, i:i+s].permute(0,1,4,2,3).clamp(-32.0, 32.0)
+                    
+                    corr1 = torch.cat([corr1, mask[:,ii1,None]], dim=2)
+                    flow1 = torch.cat([flow1, mask[:,ii1,None]], dim=2)
+
+                    net1, delta, weight = self.update(net[:,i:i+s], inp[:,ii1], corr1, flow1)
+                    net[:,i:i+s] = net1
+
+                    targets_list += [ coords[:,i:i+s] + delta.float() ]
+                    weights_list += [ weight.float() * torch.as_tensor([1.0, 1.0, 0.0]).cuda() ]
+
+                target = torch.cat(targets_list, 1)
+                weight = torch.cat(weights_list, 1)
+
+                with autocast(False):
+                    for i in range(3):
+                        Gs = MoBA(target, weight, Gs, disps, intrinsics, ii, jj, lm=0.00001, ep=.01)
+
+                    coords, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+                    residual = (target - coords)[...,:2]
+
+                poses = keyframe_graph.get_poses(Gs)
+                if self.frontend is not None:
+                    for i in range(poses.shape[1]):
+                        self.frontend.update_pose(i, poses[:,i].inv()[0].data)
+
+        return poses
+        
+    def global_refinement(self):
+        """ run global refinement """
+
+        poses = self.poses[:, :self.count]
+        images = torch.cat(self.images_list, 1).cpu()
+        depths = torch.stack(self.depths_list, 1).cpu()
+        intrinsics = torch.stack(self.intrinsics_list, 1)
+
+        poses = self.forward(poses, images, depths, intrinsics, num_steps=16)
+        self.poses[:, :self.count] = poses
+
+
+
+
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/pgo/main.py

@@ -0,0 +1,130 @@
+import torch
+from lietorch import SO3, SE3, LieGroupParameter
+
+import argparse
+import numpy as np
+import time
+import torch.optim as optim
+import torch.nn.functional as F
+
+
+def draw(verticies):
+    """ draw pose graph """
+    import open3d as o3d
+
+    n = len(verticies)
+    points = np.array([x[1][:3] for x in verticies])
+    lines = np.stack([np.arange(0,n-1), np.arange(1,n)], 1)
+
+    line_set = o3d.geometry.LineSet(
+        points=o3d.utility.Vector3dVector(points),
+        lines=o3d.utility.Vector2iVector(lines),
+    )
+    o3d.visualization.draw_geometries([line_set])
+
+def info2mat(info):
+    mat = np.zeros((6,6))
+    ix = 0
+    for i in range(mat.shape[0]):
+        mat[i,i:] = info[ix:ix+(6-i)]
+        mat[i:,i] = info[ix:ix+(6-i)]
+        ix += (6-i)
+
+    return mat
+
+def read_g2o(fn):
+    verticies, edges = [], []
+    with open(fn) as f:
+        for line in f:
+            line = line.split()
+            if line[0] == 'VERTEX_SE3:QUAT':
+                v = int(line[1])
+                pose = np.array(line[2:], dtype=np.float32)
+                verticies.append([v, pose])
+
+            elif line[0] == 'EDGE_SE3:QUAT':
+                u = int(line[1])
+                v = int(line[2])
+                pose = np.array(line[3:10], dtype=np.float32)
+                info = np.array(line[10:], dtype=np.float32)
+
+                info = info2mat(info)
+                edges.append([u, v, pose, info, line])
+
+    return verticies, edges
+
+def write_g2o(pose_graph, fn):
+    import csv
+    verticies, edges = pose_graph
+    with open(fn, 'w') as f:
+        writer = csv.writer(f, delimiter=' ')
+        for (v, pose) in verticies:
+            row = ['VERTEX_SE3:QUAT', v] + pose.tolist()
+            writer.writerow(row)
+        for edge in edges:
+            writer.writerow(edge[-1])
+
+def reshaping_fn(dE, b=1.5):
+    """ Reshaping function from "Intrinsic consensus on SO(3), Tron et al."""
+    ang = dE.log().norm(dim=-1)
+    err = 1/b - (1/b + ang) * torch.exp(-b*ang)
+    return err.sum()
+
+def gradient_initializer(pose_graph, n_steps=500, lr_init=0.2):
+    """ Riemannian Gradient Descent """
+
+    verticies, edges = pose_graph
+
+    # edge indicies (ii, jj)
+    ii = np.array([x[0] for x in edges])
+    jj = np.array([x[1] for x in edges])
+    ii = torch.from_numpy(ii).cuda()
+    jj = torch.from_numpy(jj).cuda()
+
+    Eij = np.stack([x[2][3:] for x in edges])
+    Eij = SO3(torch.from_numpy(Eij).float().cuda())
+
+    R = np.stack([x[1][3:] for x in verticies])
+    R = SO3(torch.from_numpy(R).float().cuda())
+    R = LieGroupParameter(R)
+
+    # use gradient descent with momentum
+    optimizer = optim.SGD([R], lr=lr_init, momentum=0.5)
+
+    start = time.time()
+    for i in range(n_steps):
+        optimizer.zero_grad()
+
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr_init * .995**i
+
+        # rotation error
+        dE = (R[ii].inv() * R[jj]) * Eij.inv()
+        loss = reshaping_fn(dE)
+
+        loss.backward()
+        optimizer.step()
+
+        if i%25 == 0:
+            print(i, lr_init * .995**i, loss.item())
+
+    # convert rotations to pose3
+    quats = R.group.data.detach().cpu().numpy()
+
+    for i in range(len(verticies)):
+        verticies[i][1][3:] = quats[i]
+
+    return verticies, edges
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--problem', help="input pose graph optimization file (.g2o format)")
+    args = parser.parse_args()
+
+    output_path = args.problem.replace('.g2o', '_rotavg.g2o')
+    input_pose_graph = read_g2o(args.problem)
+
+    rot_pose_graph = gradient_initializer(input_pose_graph)
+    write_g2o(rot_pose_graph, output_path)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/pgo/readme.md

@@ -0,0 +1,16 @@
+## Pose Graph Optimization / Rotation Averaging
+
+Pose Graph Optimization (PGO) is the problem of estimating the global trajectory from a set of relative pose measurements. PGO is typically performed using nonlinear least-squares algorithms (e.g Levenberg-Marquardt) and requires a good initialization in order to converge.
+
+In this experiment, we implement Riemannian Gradient Descent with a reshaping function (Tron et al. 2012). The algorithm is implemented in the function `gradient_initializer` and runs on the GPU using lietorch.
+
+### Running on a .g2o file
+
+Download a 3D problem from [datasets](https://lucacarlone.mit.edu/datasets/) (our implementation currently only supports uniform information matricies in Sphere-A, Torus, Cube, and Garage).
+
+Then run the `gradient_initializer` on the problem
+```python
+python main.py --problem=torus3D.g2o --steps=500
+```
+
+The output graph, `torus3D_rotavg.g2o`, can then be used as the initialization for non-linear least squares optimizers such as `ceres`, `g2o`, and `gtsam`.

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/readme.md

@@ -0,0 +1,9 @@
+# Examples
+
+Instructions for running demos and experiments can be found in each of the example directories
+1. [Pose Graph Optimization](pgo/readme.md) -> `pgo`
+1. [Sim3 Registration](registration/readme.md) -> `registration`
+1. [RGBD-SLAM](rgbdslam/readme.md) -> `rgbdslam`
+2. [RAFT-3D (SceneFlow)]()
+
+`core` contains networks, data loaders, and other common utility functions.

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/registration/assets/renderoption.json

@@ -0,0 +1,40 @@
+{
+	"background_color" : [ 1, 1, 1 ],
+	"class_name" : "RenderOption",
+	"default_mesh_color" : [ 0.69999999999999996, 0.69999999999999996, 0.69999999999999996 ],
+	"image_max_depth" : 3000,
+	"image_stretch_option" : 0,
+	"interpolation_option" : 0,
+	"light0_color" : [ 1, 1, 1 ],
+	"light0_diffuse_power" : 0.66000000000000003,
+	"light0_position" : [ 0, 0, 2 ],
+	"light0_specular_power" : 0.20000000000000001,
+	"light0_specular_shininess" : 100,
+	"light1_color" : [ 1, 1, 1 ],
+	"light1_diffuse_power" : 0.66000000000000003,
+	"light1_position" : [ 0, 0, 2 ],
+	"light1_specular_power" : 0.20000000000000001,
+	"light1_specular_shininess" : 100,
+	"light2_color" : [ 1, 1, 1 ],
+	"light2_diffuse_power" : 0.66000000000000003,
+	"light2_position" : [ 0, 0, -2 ],
+	"light2_specular_power" : 0.20000000000000001,
+	"light2_specular_shininess" : 100,
+	"light3_color" : [ 1, 1, 1 ],
+	"light3_diffuse_power" : 0.66000000000000003,
+	"light3_position" : [ 0, 0, -2 ],
+	"light3_specular_power" : 0.20000000000000001,
+	"light3_specular_shininess" : 100,
+	"light_ambient_color" : [ 0, 0, 0 ],
+	"light_on" : true,
+	"mesh_color_option" : 1,
+	"mesh_shade_option" : 0,
+	"mesh_show_back_face" : false,
+	"mesh_show_wireframe" : false,
+	"point_color_option" : 9,
+	"point_show_normal" : false,
+	"point_size" : 2,
+	"show_coordinate_frame" : false,
+	"version_major" : 1,
+	"version_minor" : 0
+}

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/registration/demo.py

@@ -0,0 +1,86 @@
+import sys
+sys.path.append('../core')
+
+import argparse
+import torch
+import cv2
+import numpy as np
+
+from viz import sim3_visualization
+from lietorch import SO3, SE3, Sim3
+from networks.sim3_net import Sim3Net
+
+def normalize_images(images):
+    images = images[:, :, [2,1,0]]
+    mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device)
+    std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device)
+    return (images/255.0).sub_(mean[:, None, None]).div_(std[:, None, None])
+
+def load_example(i=0):
+    """ get demo example """
+    DEPTH_SCALE = 5.0
+    if i==0:
+        image1 = cv2.imread('assets/image1.png')
+        image2 = cv2.imread('assets/image2.png')
+        depth1 = np.load('assets/depth1.npy') / DEPTH_SCALE
+        depth2 = np.load('assets/depth2.npy') / DEPTH_SCALE
+    
+    elif i==1:
+        image1 = cv2.imread('assets/image3.png')
+        image2 = cv2.imread('assets/image4.png')
+        depth1 = np.load('assets/depth3.npy') / DEPTH_SCALE
+        depth2 = np.load('assets/depth4.npy') / DEPTH_SCALE
+
+    images = np.stack([image1, image2], 0)
+    images = torch.from_numpy(images).permute(0,3,1,2)
+
+    depths = np.stack([depth1, depth2], 0)
+    depths = torch.from_numpy(depths).float()
+
+    intrinsics = np.array([320.0, 320.0, 320.0, 240.0])
+    intrinsics = np.tile(intrinsics[None], (2,1))
+    intrinsics = torch.from_numpy(intrinsics).float()
+
+    return images[None].cuda(), depths[None].cuda(), intrinsics[None].cuda()
+
+
+@torch.no_grad()
+def demo(model, index=0):
+
+    images, depths, intrinsics = load_example(index)
+
+    # initial transformation estimate
+    if args.transformation == 'SE3':
+        Gs = SE3.Identity(1, 2, device='cuda')
+
+    elif args.transformation == 'Sim3':
+        Gs = Sim3.Identity(1, 2, device='cuda')
+        depths[:,0] *= 2**(2*torch.rand(1) - 1.0).cuda()
+
+    images1 = normalize_images(images)
+    ests, _ = model(Gs, images1, depths, intrinsics, num_steps=12)
+
+    # only care about last transformation
+    Gs = ests[-1] 
+    T = Gs[:,0] * Gs[:,1].inv()
+    
+    T = T[0].matrix().double().cpu().numpy()
+    sim3_visualization(T, images, depths, intrinsics)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--transformation', default='SE3', help='checkpoint to restore')
+    parser.add_argument('--ckpt', help='checkpoint to restore')
+    args = parser.parse_args()
+
+    model = Sim3Net(args)
+    model.load_state_dict(torch.load(args.ckpt))
+
+    model.cuda()
+    model.eval()
+
+    # run two demos
+    demo(model, 0)
+    demo(model, 1)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/registration/main.py

@@ -0,0 +1,214 @@
+import sys
+sys.path.append('../core')
+
+import argparse
+import torch
+import cv2
+import numpy as np
+from collections import OrderedDict
+
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from data_readers.tartan import TartanAir, TartanAirTest
+
+from lietorch import SO3, SE3, Sim3
+from geom.losses import *
+
+# network
+from networks.sim3_net import Sim3Net
+from logger import Logger
+
+
+def show_image(image):
+    image = image.permute(1, 2, 0).cpu().numpy()
+    cv2.imshow('image', image / 255.0)
+    cv2.waitKey()
+
+def normalize_images(images):
+    images = images[:, :, [2,1,0]]
+    mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device)
+    std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device)
+    return (images/255.0).sub_(mean[:, None, None]).div_(std[:, None, None])
+
+@torch.no_grad()
+def evaluate(model):
+    """ evaluate trained model """
+
+    model.cuda()
+    model.eval()
+
+    R_THRESHOLD = 0.1
+    T_THRESHOLD = 0.01
+    S_THRESHOLD = 0.01
+
+    model.eval()
+    db = TartanAirTest()
+    test_loader = DataLoader(db, batch_size=1, shuffle=False, num_workers=4)
+
+    # random scales, make sure they are the same every time
+    from numpy.random import default_rng
+    rng = default_rng(1234)
+    scales = 2 ** rng.uniform(-1.0, 1.0, 2000)
+    scales = scales.astype(np.float32)
+
+    metrics = {'t': [], 'r': [], 's': []}
+    for i_batch, item in enumerate(test_loader):
+        images, poses, depths, intrinsics = [x.to('cuda') for x in item]
+
+        # convert poses w2c -> c2w
+        Ps = SE3(poses).inv()
+        batch, num = images.shape[:2]
+
+        if args.transformation == 'SE3':
+            Gs = SE3.Identity(Ps.shape, device='cuda')
+
+        elif args.transformation == 'Sim3':
+            Ps = Sim3(Ps)
+            Gs = Sim3.Identity(Ps.shape, device='cuda')
+
+            s = torch.as_tensor(scales[i_batch]).cuda().unsqueeze(0)
+            phi = torch.zeros(batch, num, 7, device='cuda')
+            phi[:,0,6] = s.log()
+
+            Ps = Sim3.exp(phi) * Ps
+            depths[:,0] *= s[:,None,None]
+
+        images = normalize_images(images)
+        Gs, _ = model(Gs, images, depths, intrinsics, num_steps=16)
+
+        Gs = Gs[-1]
+        dP = Ps[:,1] * Ps[:,0].inv()
+        dG = Gs[:,1] * Gs[:,0].inv()
+
+        dE = Sim3(dP.inv() * dG)
+        r_err, t_err, s_err = pose_metrics(dE)
+
+        t_err = t_err * TartanAir.DEPTH_SCALE
+
+        metrics['t'].append(t_err.item())
+        metrics['r'].append(r_err.item())
+        metrics['s'].append(s_err.item())
+
+    rlist = np.array(metrics['r'])
+    tlist = np.array(metrics['t'])
+    slist = np.array(metrics['s'])
+    
+    r_all = np.count_nonzero(rlist < R_THRESHOLD) / len(metrics['r'])
+    t_all = np.count_nonzero(tlist < T_THRESHOLD) / len(metrics['t'])
+    s_all = np.count_nonzero(slist < S_THRESHOLD) / len(metrics['s'])
+
+    print("Rotation Acc: ", r_all)
+    print("Translation Acc: ", t_all)
+    print("Scale Acc: ", s_all)
+
+
+def train(args):
+    """ Test to make sure project transform correctly maps points """
+
+    model = Sim3Net(args)
+    model.cuda()
+    model.train()
+
+    if args.ckpt is not None:
+        model.load_state_dict(torch.load(args.ckpt))
+
+    db = TartanAir(mode='training', n_frames=2, do_aug=True, fmin=8.0, fmax=100.0)
+    train_loader = DataLoader(db, batch_size=args.batch, shuffle=True, num_workers=4)
+
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
+    scheduler = optim.lr_scheduler.OneCycleLR(optimizer, 
+        args.lr, 100000, pct_start=0.01, cycle_momentum=False)
+
+    from collections import OrderedDict
+    graph = OrderedDict()
+    graph[0] = [1]
+    graph[1] = [0]
+
+    logger = Logger(args.name, scheduler)
+    should_keep_training = True
+    total_steps = 0
+
+    while should_keep_training:
+        for i_batch, item in enumerate(train_loader):
+            optimizer.zero_grad()
+            images, poses, depths, intrinsics = [x.to('cuda') for x in item]
+            
+            # convert poses w2c -> c2w
+            Ps = SE3(poses).inv()
+            batch, num = images.shape[:2]
+
+            if args.transformation == 'SE3':
+                Gs = SE3.Identity(Ps.shape, device='cuda')
+
+            elif args.transformation == 'Sim3':
+                Ps = Sim3(Ps)
+                Gs = Sim3.Identity(Ps.shape, device='cuda')
+
+                s = 2**(2*torch.rand(batch) - 1.0).cuda()
+                phi = torch.zeros(batch, num, 7, device='cuda')
+                phi[:,0,6] = s.log()
+
+                Ps = Sim3.exp(phi) * Ps
+                depths[:,0] *= s[:,None,None]
+
+            images = normalize_images(images)
+            Gs, residuals = model(Gs, images, depths, intrinsics, num_steps=args.iters)
+
+            geo_loss, geo_metrics = geodesic_loss(Ps, Gs, graph)
+            res_loss, res_metrics = residual_loss(residuals)
+
+            metrics = {}
+            metrics.update(geo_metrics)
+            metrics.update(res_metrics)
+
+            loss = args.w1 * geo_loss + args.w2 * res_loss
+            loss.backward()
+            
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
+            optimizer.step()
+            scheduler.step()
+            
+            logger.push(metrics)
+            total_steps += 1
+
+            if total_steps % 5000 == 0:
+                PATH = 'checkpoints/%s_%06d.pth' % (args.name, total_steps)
+                torch.save(model.state_dict(), PATH)
+
+                model.train()
+
+    return model
+                
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--name', default='bla', help='name your experiment')
+    parser.add_argument('--transformation', default='SE3', help='checkpoint to restore')
+    parser.add_argument('--ckpt', help='checkpoint to restore')
+    parser.add_argument('--train', action='store_true')
+
+    parser.add_argument('--batch', type=int, default=4)
+    parser.add_argument('--iters', type=int, default=8)
+    parser.add_argument('--lr', type=float, default=0.00025)
+    parser.add_argument('--clip', type=float, default=2.5)
+
+    parser.add_argument('--w1', type=float, default=10.0)
+    parser.add_argument('--w2', type=float, default=0.1)
+
+
+    args = parser.parse_args()
+
+    if args.train:
+        import os
+        if not os.path.isdir('checkpoints'):
+            os.mkdir('checkpoints')
+
+        model = train(args)
+    
+    else:
+        model = Sim3Net(args)
+        model.load_state_dict(torch.load(args.ckpt))
+
+    evaluate(model)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/registration/readme.md

@@ -0,0 +1,34 @@
+## SE3 / Sim3 Registration
+Estimate the 3D transformation between two RGB-D frames
+
+<img src="assets/registration.gif" width="480">
+
+### Models
+| Model     | Rot. Acc. | Tr. Acc. | Scale Acc. |
+| --------- | --------- | -------- | ---------- |
+| [se3.pth](https://drive.google.com/file/d/17pgeY5m-GXnrY3oFLPRaIrTZYvae_l9u/view?usp=sharing)   | 91.90         | 77.70 | -     |
+| [sim3.pth](https://drive.google.com/file/d/1LMnKND_4DAmd9DMTSKdz_zAoCgja6X43/view?usp=sharing)  | 93.45         | 76.05 | 98.70 |
+
+for thresholds 0.1 deg. rotation error, 1cm translation error, and 1% scale error.
+
+### Demo 
+Download one of the models to run the demo (requres Open3D)
+```python
+python demo.py --transformation=SE3 --ckpt=se3.pth
+python demo.py --transformation=Sim3 --ckpt=sim3.pth
+```
+
+### Training and Evaluation
+Training and evaluation is performed on the [TartanAir](https://theairlab.org/tartanair-dataset/) (only depth_left and image_left need to be downloaded). Note: our dataloader computes the optical flow between every pair of frames which can take several hours on the first run. However, this result is cached so that future loads will only take a few seconds.
+
+The training script expects the dataset to be in the directory datasets/TartanAir.
+
+To train a Sim3 network:
+```python
+python main.py --train --transformation=Sim3 --name=sim3
+```
+A trained model can then be evaluated:
+```python
+python main.py --transformation=Sim3 --ckpt=sim3.pth
+```
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/registration/viz.py

@@ -0,0 +1,87 @@
+import sys
+sys.path.append('../core')
+
+import argparse
+import torch
+import scipy
+import numpy as np
+
+import geom.projective_ops as pops
+import open3d as o3d
+
+def make_point_cloud(image, depth, intrinsics, max_depth=5.0):
+    """ create a point cloud """
+    colors = image.permute(1,2,0).view(-1,3)
+    colors = colors[...,[2,1,0]] / 255.0
+    clr = colors.cpu().numpy()
+
+    inv_depth = 1.0 / depth[None,None]
+    points = pops.iproj(inv_depth, intrinsics[None,None])
+    points = (points[..., :3] / points[..., 3:]).view(-1,3)
+    pts = points.cpu().numpy()
+
+    # open3d point cloud
+    pc = o3d.geometry.PointCloud()
+
+    keep = pts[:,2] < max_depth
+    pc.points = o3d.utility.Vector3dVector(pts[keep])
+    pc.colors = o3d.utility.Vector3dVector(clr[keep])
+
+    return pc
+
+def set_camera_pose(vis):
+    """ set initial camera position """
+    cam = vis.get_view_control().convert_to_pinhole_camera_parameters()
+
+    cam.extrinsic = np.array(
+        [[ 0.91396544,  0.1462376,  -0.37852575, 0.94374719],
+         [-0.13923432,  0.98919177,  0.04597225,  1.01177687],
+         [ 0.38115743,  0.01068673,  0.92444838,  3.35964868],
+         [ 0.,          0.,          0.,          1.        ]])
+    
+    vis.get_view_control().convert_from_pinhole_camera_parameters(cam)
+
+
+def sim3_visualization(T, images, depths, intrinsics):
+    """ convert depth to open3d point clouds """
+
+    images = images.squeeze(0)
+    depths = depths.squeeze(0)
+    intrinsics = intrinsics.squeeze(0)
+
+    pc1 = make_point_cloud(images[0], depths[0], intrinsics[0])
+    pc2 = make_point_cloud(images[1], depths[1], intrinsics[1])
+
+    sim3_visualization.index = 1
+    sim3_visualization.pc2 = pc2
+
+    NUM_STEPS = 100
+    dt = scipy.linalg.logm(T) / NUM_STEPS
+    dT = scipy.linalg.expm(dt)
+    sim3_visualization.transform = dT
+
+    def animation_callback(vis):
+        sim3_visualization.index += 1
+
+        pc2 = sim3_visualization.pc2
+        if sim3_visualization.index >= NUM_STEPS and \
+                sim3_visualization.index < 2*NUM_STEPS:
+            pc2.transform(sim3_visualization.transform)
+
+            vis.update_geometry(pc2)
+            vis.poll_events()
+            vis.update_renderer()
+
+    vis = o3d.visualization.Visualizer()
+    vis.register_animation_callback(animation_callback)
+    vis.create_window(height=540, width=960)
+
+    vis.add_geometry(pc1)
+    vis.add_geometry(pc2)
+
+    vis.get_render_option().load_from_json("assets/renderoption.json")
+    set_camera_pose(vis)
+
+    print("Press q to move to next example")
+    vis.run()
+    vis.destroy_window()

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/assets/renderoption.json

@@ -0,0 +1,40 @@
+{
+	"background_color" : [ 1, 1, 1 ],
+	"class_name" : "RenderOption",
+	"default_mesh_color" : [ 0.69999999999999996, 0.69999999999999996, 0.69999999999999996 ],
+	"image_max_depth" : 3000,
+	"image_stretch_option" : 0,
+	"interpolation_option" : 0,
+	"light0_color" : [ 1, 1, 1 ],
+	"light0_diffuse_power" : 0.66000000000000003,
+	"light0_position" : [ 0, 0, 2 ],
+	"light0_specular_power" : 0.20000000000000001,
+	"light0_specular_shininess" : 100,
+	"light1_color" : [ 1, 1, 1 ],
+	"light1_diffuse_power" : 0.66000000000000003,
+	"light1_position" : [ 0, 0, 2 ],
+	"light1_specular_power" : 0.20000000000000001,
+	"light1_specular_shininess" : 100,
+	"light2_color" : [ 1, 1, 1 ],
+	"light2_diffuse_power" : 0.66000000000000003,
+	"light2_position" : [ 0, 0, -2 ],
+	"light2_specular_power" : 0.20000000000000001,
+	"light2_specular_shininess" : 100,
+	"light3_color" : [ 1, 1, 1 ],
+	"light3_diffuse_power" : 0.66000000000000003,
+	"light3_position" : [ 0, 0, -2 ],
+	"light3_specular_power" : 0.20000000000000001,
+	"light3_specular_shininess" : 100,
+	"light_ambient_color" : [ 0, 0, 0 ],
+	"light_on" : true,
+	"mesh_color_option" : 1,
+	"mesh_shade_option" : 0,
+	"mesh_show_back_face" : false,
+	"mesh_show_wireframe" : false,
+	"point_color_option" : 9,
+	"point_show_normal" : false,
+	"point_size" : 2,
+	"show_coordinate_frame" : false,
+	"version_major" : 1,
+	"version_minor" : 0
+}

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/demo.py

@@ -0,0 +1,83 @@
+import sys
+sys.path.append('../core')
+
+from tqdm import tqdm
+import numpy as np
+import torch
+import cv2
+import os
+
+from viz import SLAMFrontend
+from lietorch import SE3
+from networks.slam_system import SLAMSystem
+from data_readers import factory
+
+
+def show_image(image):
+    image = image.permute(1, 2, 0).cpu().numpy()
+    cv2.imshow('image', image / 255.0)
+    cv2.waitKey(10)
+
+def evaluate(poses_gt, poses_est):
+    from rgbd_benchmark.evaluate_ate import evaluate_ate
+
+    poses_gt = poses_gt.cpu().numpy()
+    poses_est = poses_est.cpu().numpy()
+
+    N = poses_gt.shape[0]
+    poses_gt = dict([(i, poses_gt[i]) for i in range(N)])
+    poses_est = dict([(i, poses_est[i]) for i in range(N)])
+
+    results = evaluate_ate(poses_gt, poses_est)
+    print(results)
+
+
+@torch.no_grad()
+def run_slam(tracker, datapath, frame_rate=8.0):
+    """ run slam over full sequence """
+
+    torch.multiprocessing.set_sharing_strategy('file_system')
+    stream = factory.create_datastream(args.datapath, frame_rate=frame_rate)
+
+    # start the frontend thread
+    if args.viz:
+        frontend = SLAMFrontend().start()
+        tracker.set_frontend(frontend)
+
+    # store groundtruth poses for evaluation
+    poses_gt = []
+
+    for (tstamp, image, depth, pose, intrinsics) in tqdm(stream):
+        tracker.track(tstamp, image[None].cuda(), depth.cuda(), intrinsics.cuda())
+        poses_gt.append(pose)
+
+        if args.viz:
+            show_image(image[0])
+            frontend.update_pose(tstamp, pose[0], gt=True)
+
+    # global optimization / loop closure
+    if args.go:
+        tracker.global_refinement()
+
+    poses_gt = torch.cat(poses_gt, 0)
+    poses_est = tracker.raw_poses()
+    evaluate(poses_gt, poses_est) 
+        
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--datapath', help='path to video for slam')
+    parser.add_argument('--ckpt', help='saved network weights')
+    parser.add_argument('--viz', action='store_true', help='run visualization frontent')
+    parser.add_argument('--go', action='store_true', help='use global optimization')
+    parser.add_argument('--frame_rate', type=float, default=8.0, help='frame rate')
+    args = parser.parse_args()
+
+    # initialize tracker / load weights
+    tracker = SLAMSystem(args)
+    tracker.load_state_dict(torch.load(args.ckpt))
+    tracker.eval()
+    tracker.cuda()
+
+    run_slam(tracker, args.datapath, args.frame_rate)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/evaluate.py

@@ -0,0 +1,89 @@
+import sys
+sys.path.append('../core')
+
+from tqdm import tqdm
+import numpy as np
+import torch
+import cv2
+import os
+
+from lietorch import SE3
+from networks.slam_system import SLAMSystem
+from data_readers import factory
+
+def evaluate(poses_gt, poses_est):
+    from rgbd_benchmark.evaluate_ate import evaluate_ate
+
+    poses_gt = poses_gt.cpu().numpy()
+    poses_est = poses_est.cpu().numpy()
+
+    N = poses_gt.shape[0]
+    poses_gt = dict([(i, poses_gt[i]) for i in range(N)])
+    poses_est = dict([(i, poses_est[i]) for i in range(N)])
+
+    results = evaluate_ate(poses_gt, poses_est)
+    print(results)
+    return results['absolute_translational_error.rmse']
+
+@torch.no_grad()
+def run_slam(tracker, datapath, global_optimization=False, frame_rate=3):
+    """ run slam over full sequence """
+
+    torch.multiprocessing.set_sharing_strategy('file_system')
+    stream = factory.create_datastream(datapath, frame_rate=frame_rate)
+
+    # store groundtruth poses for evaluatino
+    poses_gt = []
+    for (tstamp, image, depth, pose, intrinsics) in tqdm(stream):
+        tracker.track(tstamp, image[None].cuda(), depth.cuda(), intrinsics.cuda())
+        poses_gt.append(pose)
+
+    if global_optimization:
+        tracker.global_refinement()
+
+    poses_gt = torch.cat(poses_gt, 0)
+    poses_est = tracker.raw_poses()   
+
+    ate = evaluate(poses_gt, poses_est) 
+    return ate
+
+def run_evaluation(ckpt, frame_rate=8.0):
+    validation_scenes = [
+        'rgbd_dataset_freiburg1_360',
+        'rgbd_dataset_freiburg1_desk',
+        'rgbd_dataset_freiburg1_desk2',
+        'rgbd_dataset_freiburg1_floor',
+        'rgbd_dataset_freiburg1_plant',
+        'rgbd_dataset_freiburg1_room',
+        'rgbd_dataset_freiburg1_rpy',
+        'rgbd_dataset_freiburg1_teddy',
+        'rgbd_dataset_freiburg1_xyz',
+    ]
+
+    results = {}
+    for scene in validation_scenes:
+        # initialize tracker / load weights
+        tracker = SLAMSystem(None)
+        tracker.load_state_dict(torch.load(ckpt))
+        tracker.eval()
+        tracker.cuda()
+        
+        datapath = os.path.join('datasets/TUM-RGBD', scene)
+        results[scene] = run_slam(tracker, datapath, 
+            global_optimization=args.go, frame_rate=frame_rate)
+
+    print("Aggregate Results: ")
+    for scene in results:
+        print(scene, results[scene])
+
+    print("MEAN: ", np.mean([results[key] for key in results]))
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt', help='saved network weights')
+    parser.add_argument('--frame_rate', type=float, default=8.0, help='frame rate')
+    parser.add_argument('--go', action='store_true', help='use global optimization')
+    args = parser.parse_args()
+
+    run_evaluation(args.ckpt, frame_rate=args.frame_rate)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/readme.md

@@ -0,0 +1,43 @@
+## RGB-D SLAM / VO
+
+
+<img src="assets/floor.png" alt="floor" height="280"/> <img src="assets/room.png" alt="room" height="280"/>
+
+
+### Pretrained Model
+
+Absolute Trajectory Error (ATE) on all freiburg1 sequences. The default model acts as a visual odometry system (no loop closure). The model rgbdslam.pth + go performs global optimization at the end of tracking to correct for drift.
+
+| Model  | 360 | desk | desk2 | floor | plant | room | rpy | teddy | xyz | avg |
+| -----  | --- | ---- | ----- | ----- | ----- | ---- | --- | ----- | --- | --- |
+| DeepV2D | 0.072 | 0.069 | 0.074 | 0.317 | 0.046 | 0.213 | 0.082 | 0.114 | 0.028 | 0.113 |
+| [lietorch_rgbdslam.pth](https://drive.google.com/file/d/1SVQTFCchZuhFeSucS5jLeNbOWyff4BA8/view?usp=sharing) | 0.076 | 0.045 | 0.054 | 0.057 | 0.032 | 0.143 | 0.064 | 0.092 | 0.033 | 0.066 |
+| [lietorch_rgbdslam.pth](https://drive.google.com/file/d/1SVQTFCchZuhFeSucS5jLeNbOWyff4BA8/view?usp=sharing) + go | 0.047 | 0.018 | 0.023 | 0.017 | 0.015 | 0.029 | 0.019 | 0.030 | 0.009 | 0.023 |
+
+### Demo
+Requires a GPU with at least 8gb of memory. First download a sequence from the [TUM-RGBD dataset](https://vision.in.tum.de/data/datasets/rgbd-dataset/download), then run the demo. You can interact with the Open3D window during tracking.
+
+```python
+python demo.py --ckpt=lietorch_rgbdslam.pth --datapath=<sequence path> --frame_rate=8.0 --go --viz
+```
+
+The `--frame_rate` flag determines the rate images are subsampled from the video (e.g `--frame_rate=8.0` subsamples the video at a rate of 8 fps). With a RTX-3090 GPU and visualization disabled, `--frame_rate <= 8.0` gives real-time performance.
+
+
+### Evaluation
+Assuming all TUM-RGBD sequences have been download, a trained model can be evaluated on the TUM-RGBD dataset
+```
+python evaluate.py --ckpt=rgbdslam.pth --datapath=<tum-rgbd root> --go --frame_rate=8.0
+```
+
+### Training
+We provide data_loaders for [NYUv2](https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html), [ScanNet](http://www.scan-net.org/), [ETH3D-SLAM](https://www.eth3d.net/slam_datasets), and [TartanAir](https://theairlab.org/tartanair-dataset/). The dataloaders will work directly on ScanNet, ETH3D, and TartanAir. For NYUv2, you will need to first extract the depths and images from the raw format then run ORB-SLAM2 to generate psuedo-groundtruth poses. Send me an email (Zachary Teed) if you need a link to the preprocessed NYU data.
+
+You can train on any subset of the datasets by listing their keys {`nyu`, `scannet`, `eth`, `tartan`}. The provided models are trained on scannet and nyu. Note: our dataloader computes the optical flow between every pair of frames which can take several hours on the first run. However, this result is cached so that future loads will only take a few seconds. The default training setttings require a GPU with 24 Gb of memory.
+
+```
+python train.py --batch=3 --iters=12 --lr=0.00025 --name nyu_scannet_eth_v2 --datasets nyu scannet
+```
+
+#### Training on your own dataset
+Additional datasets can easily be added by subclassing `RGBDDataset`, see `nyu2.py` or `scannet.py` as examples. To verify the dataloading is correct, you can use the `reprojection_test.py` script to verify that the warped images align.

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/reprojection_test.py

@@ -0,0 +1,51 @@
+import sys
+sys.path.append('../core')
+
+import torch
+import cv2
+import numpy as np
+
+from torch.utils.data import DataLoader
+from data_readers.factory import dataset_factory
+
+from lietorch import SO3, SE3, Sim3
+import geom.projective_ops as pops
+from geom.sampler_utils import bilinear_sampler
+
+def show_image(image):
+    if len(image.shape) == 3:
+        image = image.permute(1, 2, 0)
+    image = image.cpu().numpy()
+    cv2.imshow('image', image / 255.0)
+    cv2.waitKey()
+
+def reproj_test(args, N=2):
+    """ Test to make sure project transform correctly maps points """
+
+    db = dataset_factory(args.datasets, n_frames=N)
+    train_loader = DataLoader(db, batch_size=1, shuffle=True, num_workers=0)
+
+    for item in train_loader:
+        images, poses, depths, intrinsics = [x.to('cuda') for x in item]        
+        poses = SE3(poses).inv()
+        disps = 1.0 / depths
+
+        coords, _ = pops.projective_transform(poses, disps, intrinsics, [0], [1])
+        imagew = bilinear_sampler(images[:,[1]], coords[...,[0,1]])
+
+        # these two image should show camera motion
+        show_image(images[0,0])
+        show_image(images[0,1])
+
+        # these two images should show the camera motion removed by reprojection / warping
+        show_image(images[0,0])
+        show_image(imagew[0,0])
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--datasets', nargs='+', help='lists of datasets for training')
+    args = parser.parse_args()
+
+    reproj_test(args)

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/rgbd_benchmark/associate.py

@@ -0,0 +1,128 @@
+#!/usr/bin/python
+# Software License Agreement (BSD License)
+#
+# Copyright (c) 2013, Juergen Sturm, TUM
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above
+#    copyright notice, this list of conditions and the following
+#    disclaimer in the documentation and/or other materials provided
+#    with the distribution.
+#  * Neither the name of TUM nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# Requirements: 
+# sudo apt-get install python-argparse
+
+"""
+The Kinect provides the color and depth images in an un-synchronized way. This means that the set of time stamps from the color images do not intersect with those of the depth images. Therefore, we need some way of associating color images to depth images.
+
+For this purpose, you can use the ''associate.py'' script. It reads the time stamps from the rgb.txt file and the depth.txt file, and joins them by finding the best matches.
+"""
+
+import argparse
+import sys
+import os
+import numpy
+
+
+def read_file_list(filename):
+    """
+    Reads a trajectory from a text file. 
+    
+    File format:
+    The file format is "stamp d1 d2 d3 ...", where stamp denotes the time stamp (to be matched)
+    and "d1 d2 d3.." is arbitary data (e.g., a 3D position and 3D orientation) associated to this timestamp. 
+    
+    Input:
+    filename -- File name
+    
+    Output:
+    dict -- dictionary of (stamp,data) tuples
+    
+    """
+    file = open(filename)
+    data = file.read()
+    lines = data.replace(","," ").replace("\t"," ").split("\n") 
+    list = [[v.strip() for v in line.split(" ") if v.strip()!=""] for line in lines if len(line)>0 and line[0]!="#"]
+    list = [(float(l[0]),l[1:]) for l in list if len(l)>1]
+    return dict(list)
+
+def associate(first_list, second_list,offset=0.0,max_difference=0.02):
+    """
+    Associate two dictionaries of (stamp,data). As the time stamps never match exactly, we aim 
+    to find the closest match for every input tuple.
+    
+    Input:
+    first_list -- first dictionary of (stamp,data) tuples
+    second_list -- second dictionary of (stamp,data) tuples
+    offset -- time offset between both dictionaries (e.g., to model the delay between the sensors)
+    max_difference -- search radius for candidate generation
+
+    Output:
+    matches -- list of matched tuples ((stamp1,data1),(stamp2,data2))
+    
+    """
+    first_keys = list(first_list.keys())
+    second_keys = list(second_list.keys())
+    potential_matches = [(abs(a - (b + offset)), a, b) 
+                         for a in first_keys 
+                         for b in second_keys 
+                         if abs(a - (b + offset)) < max_difference]
+    potential_matches.sort()
+    matches = []
+    for diff, a, b in potential_matches:
+        if a in first_keys and b in second_keys:
+            first_keys.remove(a)
+            second_keys.remove(b)
+            matches.append((a, b))
+    
+    matches.sort()
+    return matches
+
+if __name__ == '__main__':
+    
+    # parse command line
+    parser = argparse.ArgumentParser(description='''
+    This script takes two data files with timestamps and associates them   
+    ''')
+    parser.add_argument('first_file', help='first text file (format: timestamp data)')
+    parser.add_argument('second_file', help='second text file (format: timestamp data)')
+    parser.add_argument('--first_only', help='only output associated lines from first file', action='store_true')
+    parser.add_argument('--offset', help='time offset added to the timestamps of the second file (default: 0.0)',default=0.0)
+    parser.add_argument('--max_difference', help='maximally allowed time difference for matching entries (default: 0.02)',default=0.02)
+    args = parser.parse_args()
+
+    first_list = read_file_list(args.first_file)
+    second_list = read_file_list(args.second_file)
+
+    matches = associate(first_list, second_list,float(args.offset),float(args.max_difference))    
+
+    if args.first_only:
+        for a,b in matches:
+            print("%f %s"%(a," ".join(first_list[a])))
+    else:
+        for a,b in matches:
+            print("%f %s %f %s"%(a," ".join(first_list[a]),b-float(args.offset)," ".join(second_list[b])))
+            
+        

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/rgbd_benchmark/evaluate_ate.py

@@ -0,0 +1,206 @@
+#!/usr/bin/python
+# Software License Agreement (BSD License)
+#
+# Copyright (c) 2013, Juergen Sturm, TUM
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above
+#    copyright notice, this list of conditions and the following
+#    disclaimer in the documentation and/or other materials provided
+#    with the distribution.
+#  * Neither the name of TUM nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# Requirements: 
+# sudo apt-get install python-argparse
+
+"""
+This script computes the absolute trajectory error from the ground truth
+trajectory and the estimated trajectory.
+"""
+
+import sys
+import numpy
+import argparse
+if __name__=="__main__":
+    import associate
+else:
+    from . import associate
+
+def align(model,data):
+    """Align two trajectories using the method of Horn (closed-form).
+    
+    Input:
+    model -- first trajectory (3xn)
+    data -- second trajectory (3xn)
+    
+    Output:
+    rot -- rotation matrix (3x3)
+    trans -- translation vector (3x1)
+    trans_error -- translational error per point (1xn)
+    
+    """
+    numpy.set_printoptions(precision=3,suppress=True)
+    model_zerocentered = model - model.mean(1)
+    data_zerocentered = data - data.mean(1)
+    
+    W = numpy.zeros( (3,3) )
+    for column in range(model.shape[1]):
+        W += numpy.outer(model_zerocentered[:,column],data_zerocentered[:,column])
+    U,d,Vh = numpy.linalg.linalg.svd(W.transpose())
+    S = numpy.matrix(numpy.identity( 3 ))
+    if(numpy.linalg.det(U) * numpy.linalg.det(Vh)<0):
+        S[2,2] = -1
+    rot = U*S*Vh
+    trans = data.mean(1) - rot * model.mean(1)
+    
+    model_aligned = rot * model + trans
+    alignment_error = model_aligned - data
+    
+    trans_error = numpy.sqrt(numpy.sum(numpy.multiply(alignment_error,alignment_error),0)).A[0]
+        
+    return rot,trans,trans_error
+
+def plot_traj(ax,stamps,traj,style,color,label):
+    """
+    Plot a trajectory using matplotlib. 
+    
+    Input:
+    ax -- the plot
+    stamps -- time stamps (1xn)
+    traj -- trajectory (3xn)
+    style -- line style
+    color -- line color
+    label -- plot legend
+    
+    """
+    stamps.sort()
+    interval = numpy.median([s-t for s,t in zip(stamps[1:],stamps[:-1])])
+    x = []
+    y = []
+    last = stamps[0]
+    for i in range(len(stamps)):
+        if stamps[i]-last < 2*interval:
+            x.append(traj[i][0])
+            y.append(traj[i][1])
+        elif len(x)>0:
+            ax.plot(x,y,style,color=color,label=label)
+            label=""
+            x=[]
+            y=[]
+        last= stamps[i]
+    if len(x)>0:
+        ax.plot(x,y,style,color=color,label=label)
+            
+
+def evaluate_ate(first_list, second_list, _args=""):
+    # parse command line
+    parser = argparse.ArgumentParser(
+        description='This script computes the absolute trajectory error from the ground truth trajectory and the estimated trajectory.')
+    # parser.add_argument('first_file', help='ground truth trajectory (format: timestamp tx ty tz qx qy qz qw)')
+    # parser.add_argument('second_file', help='estimated trajectory (format: timestamp tx ty tz qx qy qz qw)')
+    parser.add_argument('--offset', help='time offset added to the timestamps of the second file (default: 0.0)',default=0.0)
+    parser.add_argument('--scale', help='scaling factor for the second trajectory (default: 1.0)',default=1.0)
+    parser.add_argument('--max_difference', help='maximally allowed time difference for matching entries (default: 0.02)',default=0.02)
+    parser.add_argument('--save', help='save aligned second trajectory to disk (format: stamp2 x2 y2 z2)')
+    parser.add_argument('--save_associations', help='save associated first and aligned second trajectory to disk (format: stamp1 x1 y1 z1 stamp2 x2 y2 z2)')
+    parser.add_argument('--plot', help='plot the first and the aligned second trajectory to an image (format: png)')
+    parser.add_argument('--verbose', help='print all evaluation data (otherwise, only the RMSE absolute translational error in meters after alignment will be printed)', action='store_true')
+    args = parser.parse_args(_args)
+
+    # first_list = associate.read_file_list(args.first_file)
+    # second_list = associate.read_file_list(args.second_file)
+
+    matches = associate.associate(first_list, second_list,float(args.offset),float(args.max_difference))    
+    if len(matches)<2:
+        raise ValueError("Couldn't find matching timestamp pairs between groundtruth and estimated trajectory! Did you choose the correct sequence?")
+
+    first_xyz = numpy.matrix([[float(value) for value in first_list[a][0:3]] for a,b in matches]).transpose()
+    second_xyz = numpy.matrix([[float(value)*float(args.scale) for value in second_list[b][0:3]] for a,b in matches]).transpose()
+    rot,trans,trans_error = align(second_xyz,first_xyz)
+    
+    second_xyz_aligned = rot * second_xyz + trans
+    
+    first_stamps = list(first_list.keys())
+    first_stamps.sort()
+    first_xyz_full = numpy.matrix([[float(value) for value in first_list[b][0:3]] for b in first_stamps]).transpose()
+    
+    second_stamps = list(second_list.keys())
+    second_stamps.sort()
+    second_xyz_full = numpy.matrix([[float(value)*float(args.scale) for value in second_list[b][0:3]] for b in second_stamps]).transpose()
+    second_xyz_full_aligned = rot * second_xyz_full + trans
+    
+    if args.verbose:
+        print( "compared_pose_pairs %d pairs"%(len(trans_error)))
+             
+        print( "absolute_translational_error.rmse %f m"%numpy.sqrt(numpy.dot(trans_error,trans_error) / len(trans_error)))
+        print( "absolute_translational_error.mean %f m"%numpy.mean(trans_error))
+        print( "absolute_translational_error.median %f m"%numpy.median(trans_error))
+        print( "absolute_translational_error.std %f m"%numpy.std(trans_error))
+        print( "absolute_translational_error.min %f m"%numpy.min(trans_error))
+        print( "absolute_translational_error.max %f m"%numpy.max(trans_error))
+    
+        
+    if args.save_associations:
+        file = open(args.save_associations,"w")
+        file.write("\n".join(["%f %f %f %f %f %f %f %f"%(a,x1,y1,z1,b,x2,y2,z2) for (a,b),(x1,y1,z1),(x2,y2,z2) in zip(matches,first_xyz.transpose().A,second_xyz_aligned.transpose().A)]))
+        file.close()
+        
+    if args.save:
+        file = open(args.save,"w")
+        file.write("\n".join(["%f "%stamp+" ".join(["%f"%d for d in line]) for stamp,line in zip(second_stamps,second_xyz_full_aligned.transpose().A)]))
+        file.close()
+
+    if args.plot:
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        import matplotlib.pylab as pylab
+        from matplotlib.patches import Ellipse
+        fig = plt.figure()
+        ax = fig.add_subplot(111)
+        plot_traj(ax,first_stamps,first_xyz_full.transpose().A,'-',"black","ground truth")
+        plot_traj(ax,second_stamps,second_xyz_full_aligned.transpose().A,'-',"blue","estimated")
+
+        label="difference"
+        for (a,b),(x1,y1,z1),(x2,y2,z2) in zip(matches,first_xyz.transpose().A,second_xyz_aligned.transpose().A):
+            ax.plot([x1,x2],[y1,y2],'-',color="red",label=label)
+            label=""
+            
+        ax.legend()
+            
+        ax.set_xlabel('x [m]')
+        ax.set_ylabel('y [m]')
+        plt.savefig(args.plot,dpi=90)
+        
+    return {
+        "compared_pose_pairs": (len(trans_error)),
+        "absolute_translational_error.rmse": numpy.sqrt(numpy.dot(trans_error,trans_error) / len(trans_error)),
+        "absolute_translational_error.mean": numpy.mean(trans_error),
+        "absolute_translational_error.median": numpy.median(trans_error),
+        "absolute_translational_error.std": numpy.std(trans_error),
+        "absolute_translational_error.min": numpy.min(trans_error),
+        "absolute_translational_error.max": numpy.max(trans_error),
+    }
+
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/rgbd_benchmark/evaluate_rpe.py

@@ -0,0 +1,476 @@
+#!/usr/bin/python
+# Software License Agreement (BSD License)
+#
+# Copyright (c) 2013, Juergen Sturm, TUM
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above
+#    copyright notice, this list of conditions and the following
+#    disclaimer in the documentation and/or other materials provided
+#    with the distribution.
+#  * Neither the name of TUM nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+This script computes the relative pose error from the ground truth trajectory
+and the estimated trajectory.
+"""
+
+import argparse
+import random
+import numpy
+import sys
+
+_EPS = numpy.finfo(float).eps * 4.0
+
+def transform44(l):
+    """
+    Generate a 4x4 homogeneous transformation matrix from a 3D point and unit quaternion.
+    
+    Input:
+    l -- tuple consisting of (stamp,tx,ty,tz,qx,qy,qz,qw) where
+         (tx,ty,tz) is the 3D position and (qx,qy,qz,qw) is the unit quaternion.
+         
+    Output:
+    matrix -- 4x4 homogeneous transformation matrix
+    """
+    t = l[1:4]
+    q = numpy.array(l[4:8], dtype=numpy.float64, copy=True)
+    nq = numpy.dot(q, q)
+    if nq < _EPS:
+        return numpy.array((
+        (                1.0,                 0.0,                 0.0, t[0])
+        (                0.0,                 1.0,                 0.0, t[1])
+        (                0.0,                 0.0,                 1.0, t[2])
+        (                0.0,                 0.0,                 0.0, 1.0)
+        ), dtype=numpy.float64)
+    q *= numpy.sqrt(2.0 / nq)
+    q = numpy.outer(q, q)
+    return numpy.array((
+        (1.0-q[1, 1]-q[2, 2],     q[0, 1]-q[2, 3],     q[0, 2]+q[1, 3], t[0]),
+        (    q[0, 1]+q[2, 3], 1.0-q[0, 0]-q[2, 2],     q[1, 2]-q[0, 3], t[1]),
+        (    q[0, 2]-q[1, 3],     q[1, 2]+q[0, 3], 1.0-q[0, 0]-q[1, 1], t[2]),
+        (                0.0,                 0.0,                 0.0, 1.0)
+        ), dtype=numpy.float64)
+
+def read_trajectory(filename, matrix=True):
+    """
+    Read a trajectory from a text file. 
+    
+    Input:
+    filename -- file to be read
+    matrix -- convert poses to 4x4 matrices
+    
+    Output:
+    dictionary of stamped 3D poses
+    """
+    file = open(filename)
+    data = file.read()
+    lines = data.replace(","," ").replace("\t"," ").split("\n") 
+    list = [[float(v.strip()) for v in line.split(" ") if v.strip()!=""] for line in lines if len(line)>0 and line[0]!="#"]
+    list_ok = []
+    for i,l in enumerate(list):
+        if l[4:8]==[0,0,0,0]:
+            continue
+        isnan = False
+        for v in l:
+            if numpy.isnan(v): 
+                isnan = True
+                break
+        if isnan:
+            sys.stderr.write("Warning: line %d of file '%s' has NaNs, skipping line\n"%(i,filename))
+            continue
+        list_ok.append(l)
+    if matrix :
+      traj = dict([(l[0],transform44(l[0:])) for l in list_ok])
+    else:
+      traj = dict([(l[0],l[1:8]) for l in list_ok])
+    return traj
+
+def find_closest_index(L,t):
+    """
+    Find the index of the closest value in a list.
+    
+    Input:
+    L -- the list
+    t -- value to be found
+    
+    Output:
+    index of the closest element
+    """
+    beginning = 0
+    difference = abs(L[0] - t)
+    best = 0
+    end = len(L)
+    while beginning < end:
+        middle = int((end+beginning)/2)
+        if abs(L[middle] - t) < difference:
+            difference = abs(L[middle] - t)
+            best = middle
+        if t == L[middle]:
+            return middle
+        elif L[middle] > t:
+            end = middle
+        else:
+            beginning = middle + 1
+    return best
+
+def ominus(a,b):
+    """
+    Compute the relative 3D transformation between a and b.
+    
+    Input:
+    a -- first pose (homogeneous 4x4 matrix)
+    b -- second pose (homogeneous 4x4 matrix)
+    
+    Output:
+    Relative 3D transformation from a to b.
+    """
+    return numpy.dot(numpy.linalg.inv(a),b)
+
+def scale(a,scalar):
+    """
+    Scale the translational components of a 4x4 homogeneous matrix by a scale factor.
+    """
+    return numpy.array(
+        [[a[0,0], a[0,1], a[0,2], a[0,3]*scalar],
+         [a[1,0], a[1,1], a[1,2], a[1,3]*scalar],
+         [a[2,0], a[2,1], a[2,2], a[2,3]*scalar],
+         [a[3,0], a[3,1], a[3,2], a[3,3]]]
+                       )
+
+def compute_distance(transform):
+    """
+    Compute the distance of the translational component of a 4x4 homogeneous matrix.
+    """
+    return numpy.linalg.norm(transform[0:3,3])
+
+def compute_angle(transform):
+    """
+    Compute the rotation angle from a 4x4 homogeneous matrix.
+    """
+    # an invitation to 3-d vision, p 27
+    return numpy.arccos( min(1,max(-1, (numpy.trace(transform[0:3,0:3]) - 1)/2) ))
+
+def distances_along_trajectory(traj):
+    """
+    Compute the translational distances along a trajectory. 
+    """
+    keys = traj.keys()
+    keys.sort()
+    motion = [ominus(traj[keys[i+1]],traj[keys[i]]) for i in range(len(keys)-1)]
+    distances = [0]
+    sum = 0
+    for t in motion:
+        sum += compute_distance(t)
+        distances.append(sum)
+    return distances
+    
+def rotations_along_trajectory(traj,scale):
+    """
+    Compute the angular rotations along a trajectory. 
+    """
+    keys = traj.keys()
+    keys.sort()
+    motion = [ominus(traj[keys[i+1]],traj[keys[i]]) for i in range(len(keys)-1)]
+    distances = [0]
+    sum = 0
+    for t in motion:
+        sum += compute_angle(t)*scale
+        distances.append(sum)
+    return distances
+    
+
+def evaluate_trajectory(traj_gt,traj_est,param_max_pairs=10000,param_fixed_delta=False,param_delta=1.00,param_delta_unit="s",param_offset=0.00,param_scale=1.00):
+    """
+    Compute the relative pose error between two trajectories.
+    
+    Input:
+    traj_gt -- the first trajectory (ground truth)
+    traj_est -- the second trajectory (estimated trajectory)
+    param_max_pairs -- number of relative poses to be evaluated
+    param_fixed_delta -- false: evaluate over all possible pairs
+                         true: only evaluate over pairs with a given distance (delta)
+    param_delta -- distance between the evaluated pairs
+    param_delta_unit -- unit for comparison:
+                        "s": seconds
+                        "m": meters
+                        "rad": radians
+                        "deg": degrees
+                        "f": frames
+    param_offset -- time offset between two trajectories (to model the delay)
+    param_scale -- scale to be applied to the second trajectory
+    
+    Output:
+    list of compared poses and the resulting translation and rotation error
+    """
+    stamps_gt = list(traj_gt.keys())
+    stamps_est = list(traj_est.keys())
+    stamps_gt.sort()
+    stamps_est.sort()
+    
+    stamps_est_return = []
+    for t_est in stamps_est:
+        t_gt = stamps_gt[find_closest_index(stamps_gt,t_est + param_offset)]
+        t_est_return = stamps_est[find_closest_index(stamps_est,t_gt - param_offset)]
+        t_gt_return = stamps_gt[find_closest_index(stamps_gt,t_est_return + param_offset)]
+        if not t_est_return in stamps_est_return:
+            stamps_est_return.append(t_est_return)
+    if(len(stamps_est_return)<2):
+        raise Exception("Number of overlap in the timestamps is too small. Did you run the evaluation on the right files?")
+
+    if param_delta_unit=="s":
+        index_est = list(traj_est.keys())
+        index_est.sort()
+    elif param_delta_unit=="m":
+        index_est = distances_along_trajectory(traj_est)
+    elif param_delta_unit=="rad":
+        index_est = rotations_along_trajectory(traj_est,1)
+    elif param_delta_unit=="deg":
+        index_est = rotations_along_trajectory(traj_est,180/numpy.pi)
+    elif param_delta_unit=="f":
+        index_est = range(len(traj_est))
+    else:
+        raise Exception("Unknown unit for delta: '%s'"%param_delta_unit)
+
+    if not param_fixed_delta:
+        if(param_max_pairs==0 or len(traj_est)<numpy.sqrt(param_max_pairs)):
+            pairs = [(i,j) for i in range(len(traj_est)) for j in range(len(traj_est))]
+        else:
+            pairs = [(random.randint(0,len(traj_est)-1),random.randint(0,len(traj_est)-1)) for i in range(param_max_pairs)]
+    else:
+        pairs = []
+        for i in range(len(traj_est)):
+            j = find_closest_index(index_est,index_est[i] + param_delta)
+            if j!=len(traj_est)-1: 
+                pairs.append((i,j))
+        if(param_max_pairs!=0 and len(pairs)>param_max_pairs):
+            pairs = random.sample(pairs,param_max_pairs)
+    
+    gt_interval = numpy.median([s-t for s,t in zip(stamps_gt[1:],stamps_gt[:-1])])
+    gt_max_time_difference = 2*gt_interval
+    
+    result = []
+    for i,j in pairs:
+        stamp_est_0 = stamps_est[i]
+        stamp_est_1 = stamps_est[j]
+
+        stamp_gt_0 = stamps_gt[ find_closest_index(stamps_gt,stamp_est_0 + param_offset) ]
+        stamp_gt_1 = stamps_gt[ find_closest_index(stamps_gt,stamp_est_1 + param_offset) ]
+        
+        if(abs(stamp_gt_0 - (stamp_est_0 + param_offset)) > gt_max_time_difference  or
+           abs(stamp_gt_1 - (stamp_est_1 + param_offset)) > gt_max_time_difference):
+            continue
+        
+        error44 = ominus(  scale(
+                           ominus( traj_est[stamp_est_1], traj_est[stamp_est_0] ),param_scale),
+                           ominus( traj_gt[stamp_gt_1], traj_gt[stamp_gt_0] ) )
+        
+        trans = compute_distance(error44)
+        rot = compute_angle(error44)
+        
+        result.append([stamp_est_0,stamp_est_1,stamp_gt_0,stamp_gt_1,trans,rot])
+        
+    if len(result)<2:
+        raise Exception("Couldn't find matching timestamp pairs between groundtruth and estimated trajectory!")
+        
+    return result
+
+def percentile(seq,q):
+    """
+    Return the q-percentile of a list
+    """
+    seq_sorted = list(seq)
+    seq_sorted.sort()
+    return seq_sorted[int((len(seq_sorted)-1)*q)]
+
+
+def evaluate_rpe(_args):
+    random.seed(0)
+
+    parser = argparse.ArgumentParser(description='''
+    This script computes the relative pose error from the ground truth trajectory and the estimated trajectory. 
+    ''')
+    parser.add_argument('groundtruth_file', help='ground-truth trajectory file (format: "timestamp tx ty tz qx qy qz qw")')
+    parser.add_argument('estimated_file', help='estimated trajectory file (format: "timestamp tx ty tz qx qy qz qw")')
+    parser.add_argument('--max_pairs', help='maximum number of pose comparisons (default: 10000, set to zero to disable downsampling)', default=10000)
+    parser.add_argument('--fixed_delta', help='only consider pose pairs that have a distance of delta delta_unit (e.g., for evaluating the drift per second/meter/radian)', action='store_true')
+    parser.add_argument('--delta', help='delta for evaluation (default: 1.0)',default=1.0)
+    parser.add_argument('--delta_unit', help='unit of delta (options: \'s\' for seconds, \'m\' for meters, \'rad\' for radians, \'f\' for frames; default: \'s\')',default='s')
+    parser.add_argument('--offset', help='time offset between ground-truth and estimated trajectory (default: 0.0)',default=0.0)
+    parser.add_argument('--scale', help='scaling factor for the estimated trajectory (default: 1.0)',default=1.0)
+    parser.add_argument('--save', help='text file to which the evaluation will be saved (format: stamp_est0 stamp_est1 stamp_gt0 stamp_gt1 trans_error rot_error)')
+    parser.add_argument('--plot', help='plot the result to a file (requires --fixed_delta, output format: png)')
+    parser.add_argument('--verbose', help='print all evaluation data (otherwise, only the mean translational error measured in meters will be printed)', action='store_true')
+    args = parser.parse_args(_args)
+    
+    if args.plot and not args.fixed_delta:
+        raise ValueError("The '--plot' option can only be used in combination with '--fixed_delta'")
+    
+    traj_gt = read_trajectory(args.groundtruth_file)
+    traj_est = read_trajectory(args.estimated_file)
+    
+    result = evaluate_trajectory(traj_gt,
+                                 traj_est,
+                                 int(args.max_pairs),
+                                 args.fixed_delta,
+                                 float(args.delta),
+                                 args.delta_unit,
+                                 float(args.offset),
+                                 float(args.scale))
+    
+    stamps = numpy.array(result)[:,0]
+    trans_error = numpy.array(result)[:,4]
+    rot_error = numpy.array(result)[:,5]
+    
+    if args.save:
+        f = open(args.save,"w")
+        f.write("\n".join([" ".join(["%f"%v for v in line]) for line in result]))
+        f.close()
+    
+    if args.verbose:
+        print( "compared_pose_pairs %d pairs"%(len(trans_error)))
+
+        print( "translational_error.rmse %f m"%numpy.sqrt(numpy.dot(trans_error,trans_error) / len(trans_error)))
+        print( "translational_error.mean %f m"%numpy.mean(trans_error))
+        print( "translational_error.median %f m"%numpy.median(trans_error))
+        print( "translational_error.std %f m"%numpy.std(trans_error))
+        print( "translational_error.min %f m"%numpy.min(trans_error))
+        print( "translational_error.max %f m"%numpy.max(trans_error))
+
+        print( "rotational_error.rmse %f deg"%(numpy.sqrt(numpy.dot(rot_error,rot_error) / len(rot_error)) * 180.0 / numpy.pi))
+        print( "rotational_error.mean %f deg"%(numpy.mean(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.median %f deg"%(numpy.median(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.std %f deg"%(numpy.std(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.min %f deg"%(numpy.min(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.max %f deg"%(numpy.max(rot_error) * 180.0 / numpy.pi))
+    else:
+        print( numpy.mean(trans_error))
+
+    if args.plot:    
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        import matplotlib.pylab as pylab
+        fig = plt.figure()
+        ax = fig.add_subplot(111)        
+        ax.plot(stamps - stamps[0],trans_error,'-',color="blue")
+        #ax.plot([t for t,e in err_rot],[e for t,e in err_rot],'-',color="red")
+        ax.set_xlabel('time [s]')
+        ax.set_ylabel('translational error [m]')
+        plt.savefig(args.plot,dpi=300)
+        
+    return {
+        "translational_error.rmse": numpy.sqrt(numpy.dot(trans_error,trans_error) / len(trans_error)),
+        "translational_error.mean": numpy.mean(trans_error),
+        "translational_error.median": numpy.median(trans_error),
+        "translational_error.std": numpy.std(trans_error),
+        "translational_error.min": numpy.min(trans_error),
+        "translational_error.max": numpy.max(trans_error),
+
+        "rotational_error.rmse": (numpy.sqrt(numpy.dot(rot_error,rot_error) / len(rot_error)) * 180.0 / numpy.pi),
+        "rotational_error.mean": (numpy.mean(rot_error) * 180.0 / numpy.pi),
+        "rotational_error.median": (numpy.median(rot_error) * 180.0 / numpy.pi),
+        "rotational_error.std": (numpy.std(rot_error) * 180.0 / numpy.pi),
+        "rotational_error.min": (numpy.min(rot_error) * 180.0 / numpy.pi),
+        "rotational_error.max": (numpy.max(rot_error) * 180.0 / numpy.pi),
+    }
+
+
+
+if __name__ == '__main__':
+    random.seed(0)
+
+    parser = argparse.ArgumentParser(description='''
+    This script computes the relative pose error from the ground truth trajectory and the estimated trajectory. 
+    ''')
+    parser.add_argument('groundtruth_file', help='ground-truth trajectory file (format: "timestamp tx ty tz qx qy qz qw")')
+    parser.add_argument('estimated_file', help='estimated trajectory file (format: "timestamp tx ty tz qx qy qz qw")')
+    parser.add_argument('--max_pairs', help='maximum number of pose comparisons (default: 10000, set to zero to disable downsampling)', default=10000)
+    parser.add_argument('--fixed_delta', help='only consider pose pairs that have a distance of delta delta_unit (e.g., for evaluating the drift per second/meter/radian)', action='store_true')
+    parser.add_argument('--delta', help='delta for evaluation (default: 1.0)',default=1.0)
+    parser.add_argument('--delta_unit', help='unit of delta (options: \'s\' for seconds, \'m\' for meters, \'rad\' for radians, \'f\' for frames; default: \'s\')',default='s')
+    parser.add_argument('--offset', help='time offset between ground-truth and estimated trajectory (default: 0.0)',default=0.0)
+    parser.add_argument('--scale', help='scaling factor for the estimated trajectory (default: 1.0)',default=1.0)
+    parser.add_argument('--save', help='text file to which the evaluation will be saved (format: stamp_est0 stamp_est1 stamp_gt0 stamp_gt1 trans_error rot_error)')
+    parser.add_argument('--plot', help='plot the result to a file (requires --fixed_delta, output format: png)')
+    parser.add_argument('--verbose', help='print all evaluation data (otherwise, only the mean translational error measured in meters will be printed)', action='store_true')
+    args = parser.parse_args()
+    
+    if args.plot and not args.fixed_delta:
+        sys.exit("The '--plot' option can only be used in combination with '--fixed_delta'")
+    
+    traj_gt = read_trajectory(args.groundtruth_file)
+    traj_est = read_trajectory(args.estimated_file)
+    
+    result = evaluate_trajectory(traj_gt,
+                                 traj_est,
+                                 int(args.max_pairs),
+                                 args.fixed_delta,
+                                 float(args.delta),
+                                 args.delta_unit,
+                                 float(args.offset),
+                                 float(args.scale))
+    
+    stamps = numpy.array(result)[:,0]
+    trans_error = numpy.array(result)[:,4]
+    rot_error = numpy.array(result)[:,5]
+    
+    if args.save:
+        f = open(args.save,"w")
+        f.write("\n".join([" ".join(["%f"%v for v in line]) for line in result]))
+        f.close()
+    
+    if args.verbose:
+        print( "compared_pose_pairs %d pairs"%(len(trans_error)))
+
+        print( "translational_error.rmse %f m"%numpy.sqrt(numpy.dot(trans_error,trans_error) / len(trans_error)))
+        print( "translational_error.mean %f m"%numpy.mean(trans_error))
+        print( "translational_error.median %f m"%numpy.median(trans_error))
+        print( "translational_error.std %f m"%numpy.std(trans_error))
+        print( "translational_error.min %f m"%numpy.min(trans_error))
+        print( "translational_error.max %f m"%numpy.max(trans_error))
+
+        print( "rotational_error.rmse %f deg"%(numpy.sqrt(numpy.dot(rot_error,rot_error) / len(rot_error)) * 180.0 / numpy.pi))
+        print( "rotational_error.mean %f deg"%(numpy.mean(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.median %f deg"%(numpy.median(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.std %f deg"%(numpy.std(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.min %f deg"%(numpy.min(rot_error) * 180.0 / numpy.pi))
+        print( "rotational_error.max %f deg"%(numpy.max(rot_error) * 180.0 / numpy.pi))
+    else:
+        print( numpy.mean(trans_error))
+
+    if args.plot:    
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        import matplotlib.pylab as pylab
+        fig = plt.figure()
+        ax = fig.add_subplot(111)        
+        ax.plot(stamps - stamps[0],trans_error,'-',color="blue")
+        #ax.plot([t for t,e in err_rot],[e for t,e in err_rot],'-',color="red")
+        ax.set_xlabel('time [s]')
+        ax.set_ylabel('translational error [m]')
+        plt.savefig(args.plot,dpi=300)
+        
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/train.py

@@ -0,0 +1,125 @@
+import sys
+sys.path.append('../core')
+
+import cv2
+import numpy as np
+from collections import OrderedDict
+
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from data_readers.factory import dataset_factory
+
+from lietorch import SO3, SE3, Sim3
+from geom.losses import geodesic_loss, residual_loss
+
+# network
+from networks.rslam import RaftSLAM
+from logger import Logger
+from evaluate import run_evaluation
+
+def show_image(image):
+    image = image.permute(1, 2, 0).cpu().numpy()
+    cv2.imshow('image', image / 255.0)
+    cv2.waitKey()
+
+def normalize_images(images):
+    images = images[:, :, [2,1,0]]
+    mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device)
+    std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device)
+    return (images/255.0).sub_(mean[:, None, None]).div_(std[:, None, None])
+
+def train(args):
+    """ Test to make sure project transform correctly maps points """
+
+    N = args.n_frames
+    model = RaftSLAM(args)
+    model.cuda()
+    model.train()
+
+    if args.ckpt is not None:
+        model.load_state_dict(torch.load(args.ckpt))
+
+    db = dataset_factory(args.datasets, n_frames=N, fmin=16.0, fmax=96.0)
+    train_loader = DataLoader(db, batch_size=args.batch, shuffle=True, num_workers=4)
+
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
+    scheduler = optim.lr_scheduler.OneCycleLR(optimizer, 
+        args.lr, args.steps, pct_start=0.01, cycle_momentum=False)
+
+    logger = Logger(args.name, scheduler)
+    should_keep_training = True
+    total_steps = 0
+
+    while should_keep_training:
+        for i_batch, item in enumerate(train_loader):
+            optimizer.zero_grad()
+
+            graph = OrderedDict()
+            for i in range(N):
+                graph[i] = [j for j in range(N) if i!=j and abs(i-j) <= 2]
+            
+            images, poses, depths, intrinsics = [x.to('cuda') for x in item]
+            
+            # convert poses w2c -> c2w
+            Ps = SE3(poses).inv()
+            Gs = SE3.Identity(Ps.shape, device='cuda')
+
+            images = normalize_images(images)
+            Gs, residuals = model(Gs, images, depths, intrinsics, graph, num_steps=args.iters)
+
+            geo_loss, geo_metrics = geodesic_loss(Ps, Gs, graph)
+            res_loss, res_metrics = residual_loss(residuals)
+
+            metrics = {}
+            metrics.update(geo_metrics)
+            metrics.update(res_metrics)
+
+            loss = args.w1 * geo_loss + args.w2 * res_loss
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
+            optimizer.step()
+            scheduler.step()
+            
+            logger.push(metrics)
+            total_steps += 1
+
+            if total_steps % 10000 == 0:
+                PATH = 'checkpoints/%s_%06d.pth' % (args.name, total_steps)
+                torch.save(model.state_dict(), PATH)
+
+                run_evaluation(PATH)
+
+            if total_steps >= args.steps:
+                should_keep_training = False
+                break
+
+    return model
+                
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--name', default='bla', help='name your experiment')
+    parser.add_argument('--ckpt', help='checkpoint to restore')
+    parser.add_argument('--datasets', nargs='+', help='lists of datasets for training')
+
+    parser.add_argument('--batch', type=int, default=2)
+    parser.add_argument('--iters', type=int, default=8)
+    parser.add_argument('--steps', type=int, default=100000)
+    parser.add_argument('--lr', type=float, default=0.0001)
+    parser.add_argument('--clip', type=float, default=2.5)
+    parser.add_argument('--n_frames', type=int, default=4)
+
+    parser.add_argument('--w1', type=float, default=10.0)
+    parser.add_argument('--w2', type=float, default=0.1)
+
+    args = parser.parse_args()
+
+    import os
+    if not os.path.isdir('checkpoints'):
+        os.mkdir('checkpoints')
+
+    model = train(args)
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/thirdparty/lietorch/examples/rgbdslam/viz.py

@@ -0,0 +1,160 @@
+import time
+import argparse
+import torch
+import scipy
+import numpy as np
+import open3d as o3d
+
+from queue import Empty
+from multiprocessing import Queue, Process
+from scipy.spatial.transform import Rotation
+
+def pose_matrix_from_quaternion(pvec):
+    """ convert 4x4 pose matrix to (t, q) """
+    pose = np.eye(4)
+    pose[:3,:3] = Rotation.from_quat(pvec[3:]).as_matrix()
+    pose[:3, 3] = pvec[:3]
+    return pose
+
+def create_camera_actor(is_gt=False, scale=0.05):
+    """ build open3d camera polydata """
+
+    cam_points = scale * np.array([
+        [ 0,   0,   0],
+        [-1,  -1, 1.5],
+        [ 1,  -1, 1.5],
+        [ 1,   1, 1.5],
+        [-1,   1, 1.5],
+        [-0.5, 1, 1.5],
+        [ 0.5, 1, 1.5],
+        [ 0, 1.2, 1.5]])
+
+    cam_lines = np.array([[1, 2], [2, 3], [3, 4], [4, 1],
+        [1, 0], [0, 2], [3, 0], [0, 4], [5, 7], [7, 6]])
+
+    camera_actor = o3d.geometry.LineSet(
+        points=o3d.utility.Vector3dVector(cam_points),
+        lines=o3d.utility.Vector2iVector(cam_lines))
+
+    color = (0.0, 0.0, 0.0) if is_gt else (0.0, 0.8, 0.8)
+    camera_actor.paint_uniform_color(color)
+
+    return camera_actor
+
+def create_point_cloud_actor(points, colors):
+    """ open3d point cloud from numpy array """
+
+    point_cloud = o3d.geometry.PointCloud()
+    point_cloud.points = o3d.utility.Vector3dVector(points)
+    point_cloud.colors = o3d.utility.Vector3dVector(colors)
+
+    return point_cloud
+
+def draw_trajectory(queue):
+
+    draw_trajectory.queue = queue
+    draw_trajectory.cameras = {}
+    draw_trajectory.points = {}
+    draw_trajectory.ix = 0
+    draw_trajectory.warmup = 8
+
+    def animation_callback(vis):
+        cam = vis.get_view_control().convert_to_pinhole_camera_parameters()
+        while True:
+            try:
+                data = draw_trajectory.queue.get_nowait()
+                if data[0] == 'pose':
+                    i, pose, is_gt = data[1:]
+                    
+                    # convert to 4x4 matrix
+                    pose = pose_matrix_from_quaternion(pose)
+
+                    if i in draw_trajectory.cameras:
+                        cam_actor, pose_prev = draw_trajectory.cameras[i]
+                        pose_change = pose @ np.linalg.inv(pose_prev)
+                        
+                        cam_actor.transform(pose_change)
+                        vis.update_geometry(cam_actor)
+
+                        if i in draw_trajectory.points:
+                            pc = draw_trajectory.points[i]
+                            pc.transform(pose_change)
+                            vis.update_geometry(pc)
+
+                    else:
+                        cam_actor = create_camera_actor(is_gt)
+                        cam_actor.transform(pose)
+                        vis.add_geometry(cam_actor)
+
+                    if not is_gt:
+                        draw_trajectory.cameras[i] = (cam_actor, pose)
+
+                elif data[0] == 'points':
+                    i, points, colors = data[1:]
+                    point_actor = create_point_cloud_actor(points, colors)
+
+                    pose = draw_trajectory.cameras[i][1]
+                    point_actor.transform(pose)
+                    vis.add_geometry(point_actor)
+
+                    draw_trajectory.points[i] = point_actor
+
+                elif data[0] == 'reset':
+                    draw_trajectory.warmup = -1
+                    
+                    for i in draw_trajectory.points:
+                        vis.remove_geometry(draw_trajectory.points[i])
+
+                    for i in draw_trajectory.cameras:
+                        vis.remove_geometry(draw_trajectory.cameras[i][0])
+
+                    draw_trajectory.cameras = {}
+                    draw_trajectory.points = {}
+
+            except Empty:
+                break
+
+        # hack to allow interacting with vizualization during inference
+        if len(draw_trajectory.cameras) >= draw_trajectory.warmup:
+            cam = vis.get_view_control().convert_from_pinhole_camera_parameters(cam)
+
+        vis.poll_events()
+        vis.update_renderer()
+
+    vis = o3d.visualization.Visualizer()
+
+    vis.register_animation_callback(animation_callback)
+    vis.create_window(height=540, width=960)
+    vis.get_render_option().load_from_json("assets/renderoption.json")
+
+    vis.run()
+    vis.destroy_window()
+
+
+class SLAMFrontend:
+    def __init__(self):
+        self.queue = Queue()
+        self.p = Process(target=draw_trajectory, args=(self.queue, ))
+
+    def update_pose(self, index, pose, gt=False):
+        if isinstance(pose, torch.Tensor):
+            pose = pose.cpu().numpy()
+        self.queue.put_nowait(('pose', index, pose, gt))
+
+    def update_points(self, index, points, colors):
+        if isinstance(points, torch.Tensor):
+            points = points.cpu().numpy()
+        self.queue.put_nowait(('points', index, points, colors))
+    
+    def reset(self):
+        self.queue.put_nowait(('reset', ))
+
+    def start(self):
+        self.p.start()
+        return self
+
+    def join(self):
+        self.p.join()
+
+
+