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

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

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

@@ -0,0 +1 @@
+

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,193 @@
+import numpy as np
+import torch
+import lietorch
+import droid_backends
+
+from torch.multiprocessing import Process, Queue, Lock, Value
+from collections import OrderedDict
+
+from droid_net import cvx_upsample
+import geom.projective_ops as pops
+
+class DepthVideo:
+    def __init__(self, image_size=[480, 640], buffer=1024, stereo=False, device="cuda:0"):
+                
+        # current keyframe count
+        self.counter = Value('i', 0)
+        self.ready = Value('i', 0)
+        self.ht = ht = image_size[0]
+        self.wd = wd = image_size[1]
+
+        ### state attributes ###
+        self.tstamp = torch.zeros(buffer, device="cuda", dtype=torch.float).share_memory_()
+        self.images = torch.zeros(buffer, 3, ht, wd, device="cuda", dtype=torch.uint8)
+        self.dirty = torch.zeros(buffer, device="cuda", dtype=torch.bool).share_memory_()
+        self.red = torch.zeros(buffer, device="cuda", dtype=torch.bool).share_memory_()
+        self.poses = torch.zeros(buffer, 7, device="cuda", dtype=torch.float).share_memory_()
+        self.disps = torch.ones(buffer, ht//8, wd//8, device="cuda", dtype=torch.float).share_memory_()
+        self.disps_sens = torch.zeros(buffer, ht//8, wd//8, device="cuda", dtype=torch.float).share_memory_()
+        self.disps_up = torch.zeros(buffer, ht, wd, device="cuda", dtype=torch.float).share_memory_()
+        self.intrinsics = torch.zeros(buffer, 4, device="cuda", dtype=torch.float).share_memory_()
+
+        self.stereo = stereo
+        c = 1 if not self.stereo else 2
+
+        ### feature attributes ###
+        self.fmaps = torch.zeros(buffer, c, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
+        self.nets = torch.zeros(buffer, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
+        self.inps = torch.zeros(buffer, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
+
+        # initialize poses to identity transformation
+        self.poses[:] = torch.as_tensor([0, 0, 0, 0, 0, 0, 1], dtype=torch.float, device="cuda")
+        
+    def get_lock(self):
+        return self.counter.get_lock()
+
+    def __item_setter(self, index, item):
+        if isinstance(index, int) and index >= self.counter.value:
+            self.counter.value = index + 1
+        
+        elif isinstance(index, torch.Tensor) and index.max().item() > self.counter.value:
+            self.counter.value = index.max().item() + 1
+
+        # self.dirty[index] = True
+        self.tstamp[index] = item[0]
+        self.images[index] = item[1]
+
+        if item[2] is not None:
+            self.poses[index] = item[2]
+
+        if item[3] is not None:
+            self.disps[index] = item[3]
+
+        if item[4] is not None:
+            depth = item[4][3::8,3::8]
+            self.disps_sens[index] = torch.where(depth>0, 1.0/depth, depth)
+
+        if item[5] is not None:
+            self.intrinsics[index] = item[5]
+
+        if len(item) > 6:
+            self.fmaps[index] = item[6]
+
+        if len(item) > 7:
+            self.nets[index] = item[7]
+
+        if len(item) > 8:
+            self.inps[index] = item[8]
+
+    def __setitem__(self, index, item):
+        with self.get_lock():
+            self.__item_setter(index, item)
+
+    def __getitem__(self, index):
+        """ index the depth video """
+
+        with self.get_lock():
+            # support negative indexing
+            if isinstance(index, int) and index < 0:
+                index = self.counter.value + index
+
+            item = (
+                self.poses[index],
+                self.disps[index],
+                self.intrinsics[index],
+                self.fmaps[index],
+                self.nets[index],
+                self.inps[index])
+
+        return item
+
+    def append(self, *item):
+        with self.get_lock():
+            self.__item_setter(self.counter.value, item)
+
+
+    ### geometric operations ###
+
+    @staticmethod
+    def format_indicies(ii, jj):
+        """ to device, long, {-1} """
+
+        if not isinstance(ii, torch.Tensor):
+            ii = torch.as_tensor(ii)
+
+        if not isinstance(jj, torch.Tensor):
+            jj = torch.as_tensor(jj)
+
+        ii = ii.to(device="cuda", dtype=torch.long).reshape(-1)
+        jj = jj.to(device="cuda", dtype=torch.long).reshape(-1)
+
+        return ii, jj
+
+    def upsample(self, ix, mask):
+        """ upsample disparity """
+
+        disps_up = cvx_upsample(self.disps[ix].unsqueeze(-1), mask)
+        self.disps_up[ix] = disps_up.squeeze()
+
+    def normalize(self):
+        """ normalize depth and poses """
+
+        with self.get_lock():
+            s = self.disps[:self.counter.value].mean()
+            self.disps[:self.counter.value] /= s
+            self.poses[:self.counter.value,:3] *= s
+            self.dirty[:self.counter.value] = True
+
+
+    def reproject(self, ii, jj):
+        """ project points from ii -> jj """
+        ii, jj = DepthVideo.format_indicies(ii, jj)
+        Gs = lietorch.SE3(self.poses[None])
+
+        coords, valid_mask = \
+            pops.projective_transform(Gs, self.disps[None], self.intrinsics[None], ii, jj)
+
+        return coords, valid_mask
+
+    def distance(self, ii=None, jj=None, beta=0.3, bidirectional=True):
+        """ frame distance metric """
+
+        return_matrix = False
+        if ii is None:
+            return_matrix = True
+            N = self.counter.value
+            ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
+        
+        ii, jj = DepthVideo.format_indicies(ii, jj)
+
+        if bidirectional:
+
+            poses = self.poses[:self.counter.value].clone()
+
+            d1 = droid_backends.frame_distance(
+                poses, self.disps, self.intrinsics[0], ii, jj, beta)
+
+            d2 = droid_backends.frame_distance(
+                poses, self.disps, self.intrinsics[0], jj, ii, beta)
+
+            d = .5 * (d1 + d2)
+
+        else:
+            d = droid_backends.frame_distance(
+                self.poses, self.disps, self.intrinsics[0], ii, jj, beta)
+
+        if return_matrix:
+            return d.reshape(N, N)
+
+        return d
+
+    def ba(self, target, weight, eta, ii, jj, t0=1, t1=None, itrs=2, lm=1e-4, ep=0.1, motion_only=False):
+        """ dense bundle adjustment (DBA) """
+
+        with self.get_lock():
+
+            # [t0, t1] window of bundle adjustment optimization
+            if t1 is None:
+                t1 = max(ii.max().item(), jj.max().item()) + 1
+
+            droid_backends.ba(self.poses, self.disps, self.intrinsics[0], self.disps_sens,
+                target, weight, eta, ii, jj, t0, t1, itrs, lm, ep, motion_only)
+
+            self.disps.clamp_(min=0.001)

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

@@ -0,0 +1,89 @@
+import torch
+import lietorch
+import numpy as np
+
+from droid_net import DroidNet
+from depth_video import DepthVideo
+from motion_filter import MotionFilter
+from droid_frontend import DroidFrontend
+from droid_backend import DroidBackend
+from trajectory_filler import PoseTrajectoryFiller
+
+from collections import OrderedDict
+from torch.multiprocessing import Process
+
+
+class Droid:
+    def __init__(self, args):
+        super(Droid, self).__init__()
+        self.load_weights(args.weights)
+        self.args = args
+        self.disable_vis = args.disable_vis
+
+        # store images, depth, poses, intrinsics (shared between processes)
+        self.video = DepthVideo(args.image_size, args.buffer, stereo=args.stereo)
+
+        # filter incoming frames so that there is enough motion
+        self.filterx = MotionFilter(self.net, self.video, thresh=args.filter_thresh)
+
+        # frontend process
+        self.frontend = DroidFrontend(self.net, self.video, self.args)
+        
+        # backend process
+        self.backend = DroidBackend(self.net, self.video, self.args)
+
+        # visualizer
+        if not self.disable_vis:
+            from visualization import droid_visualization
+            self.visualizer = Process(target=droid_visualization, args=(self.video,))
+            self.visualizer.start()
+
+        # post processor - fill in poses for non-keyframes
+        self.traj_filler = PoseTrajectoryFiller(self.net, self.video)
+
+
+    def load_weights(self, weights):
+        """ load trained model weights """
+
+        print(weights)
+        self.net = DroidNet()
+        state_dict = OrderedDict([
+            (k.replace("module.", ""), v) for (k, v) in torch.load(weights).items()])
+
+        state_dict["update.weight.2.weight"] = state_dict["update.weight.2.weight"][:2]
+        state_dict["update.weight.2.bias"] = state_dict["update.weight.2.bias"][:2]
+        state_dict["update.delta.2.weight"] = state_dict["update.delta.2.weight"][:2]
+        state_dict["update.delta.2.bias"] = state_dict["update.delta.2.bias"][:2]
+
+        self.net.load_state_dict(state_dict)
+        self.net.to("cuda:0").eval()
+
+    def track(self, tstamp, image, depth=None, intrinsics=None):
+        """ main thread - update map """
+
+        with torch.no_grad():
+            # check there is enough motion
+            self.filterx.track(tstamp, image, depth, intrinsics)
+
+            # local bundle adjustment
+            self.frontend()
+
+            # global bundle adjustment
+            # self.backend()
+
+    def terminate(self, stream=None):
+        """ terminate the visualization process, return poses [t, q] """
+
+        del self.frontend
+
+        torch.cuda.empty_cache()
+        print("#" * 32)
+        self.backend(7)
+
+        torch.cuda.empty_cache()
+        print("#" * 32)
+        self.backend(12)
+
+        camera_trajectory = self.traj_filler(stream)
+        return camera_trajectory.inv().data.cpu().numpy()
+

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

@@ -0,0 +1,41 @@
+import torch
+import lietorch
+import numpy as np
+
+from lietorch import SE3
+from factor_graph import FactorGraph
+
+
+class DroidBackend:
+    def __init__(self, net, video, args):
+        self.video = video
+        self.update_op = net.update
+
+        # global optimization window
+        self.t0 = 0
+        self.t1 = 0
+
+        self.upsample = args.upsample
+        self.beta = args.beta
+        self.backend_thresh = args.backend_thresh
+        self.backend_radius = args.backend_radius
+        self.backend_nms = args.backend_nms
+        
+    @torch.no_grad()
+    def __call__(self, steps=12):
+        """ main update """
+
+        t = self.video.counter.value
+        if not self.video.stereo and not torch.any(self.video.disps_sens):
+             self.video.normalize()
+
+        graph = FactorGraph(self.video, self.update_op, corr_impl="alt", max_factors=16*t, upsample=self.upsample)
+
+        graph.add_proximity_factors(rad=self.backend_radius, 
+                                    nms=self.backend_nms, 
+                                    thresh=self.backend_thresh, 
+                                    beta=self.beta)
+
+        graph.update_lowmem(steps=steps)
+        graph.clear_edges()
+        self.video.dirty[:t] = True

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

@@ -0,0 +1,119 @@
+import torch
+import lietorch
+import numpy as np
+
+from lietorch import SE3
+from factor_graph import FactorGraph
+
+
+class DroidFrontend:
+    def __init__(self, net, video, args):
+        self.video = video
+        self.update_op = net.update
+        self.graph = FactorGraph(video, net.update, max_factors=48, upsample=args.upsample)
+
+        # local optimization window
+        self.t0 = 0
+        self.t1 = 0
+
+        # frontent variables
+        self.is_initialized = False
+        self.count = 0
+
+        self.max_age = 25
+        self.iters1 = 4
+        self.iters2 = 2
+
+        self.warmup = args.warmup
+        self.beta = args.beta
+        self.frontend_nms = args.frontend_nms
+        self.keyframe_thresh = args.keyframe_thresh
+        self.frontend_window = args.frontend_window
+        self.frontend_thresh = args.frontend_thresh
+        self.frontend_radius = args.frontend_radius
+
+    def __update(self):
+        """ add edges, perform update """
+
+        self.count += 1
+        self.t1 += 1
+
+        if self.graph.corr is not None:
+            self.graph.rm_factors(self.graph.age > self.max_age, store=True)
+
+        self.graph.add_proximity_factors(self.t1-5, max(self.t1-self.frontend_window, 0), 
+            rad=self.frontend_radius, nms=self.frontend_nms, thresh=self.frontend_thresh, beta=self.beta, remove=True)
+
+        self.video.disps[self.t1-1] = torch.where(self.video.disps_sens[self.t1-1] > 0, 
+           self.video.disps_sens[self.t1-1], self.video.disps[self.t1-1])
+
+        for itr in range(self.iters1):
+            self.graph.update(None, None, use_inactive=True)
+
+        # set initial pose for next frame
+        poses = SE3(self.video.poses)
+        d = self.video.distance([self.t1-3], [self.t1-2], beta=self.beta, bidirectional=True)
+
+        if d.item() < self.keyframe_thresh:
+            self.graph.rm_keyframe(self.t1 - 2)
+            
+            with self.video.get_lock():
+                self.video.counter.value -= 1
+                self.t1 -= 1
+
+        else:
+            for itr in range(self.iters2):
+                self.graph.update(None, None, use_inactive=True)
+
+        # set pose for next itration
+        self.video.poses[self.t1] = self.video.poses[self.t1-1]
+        self.video.disps[self.t1] = self.video.disps[self.t1-1].mean()
+
+        # update visualization
+        self.video.dirty[self.graph.ii.min():self.t1] = True
+
+    def __initialize(self):
+        """ initialize the SLAM system """
+
+        self.t0 = 0
+        self.t1 = self.video.counter.value
+
+        self.graph.add_neighborhood_factors(self.t0, self.t1, r=3)
+
+        for itr in range(8):
+            self.graph.update(1, use_inactive=True)
+
+        self.graph.add_proximity_factors(0, 0, rad=2, nms=2, thresh=self.frontend_thresh, remove=False)
+
+        for itr in range(8):
+            self.graph.update(1, use_inactive=True)
+
+
+        # self.video.normalize()
+        self.video.poses[self.t1] = self.video.poses[self.t1-1].clone()
+        self.video.disps[self.t1] = self.video.disps[self.t1-4:self.t1].mean()
+
+        # initialization complete
+        self.is_initialized = True
+        self.last_pose = self.video.poses[self.t1-1].clone()
+        self.last_disp = self.video.disps[self.t1-1].clone()
+        self.last_time = self.video.tstamp[self.t1-1].clone()
+
+        with self.video.get_lock():
+            self.video.ready.value = 1
+            self.video.dirty[:self.t1] = True
+
+        self.graph.rm_factors(self.graph.ii < self.warmup-4, store=True)
+
+    def __call__(self):
+        """ main update """
+
+        # do initialization
+        if not self.is_initialized and self.video.counter.value == self.warmup:
+            self.__initialize()
+            
+        # do update
+        elif self.is_initialized and self.t1 < self.video.counter.value:
+            self.__update()
+
+        

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

@@ -0,0 +1,222 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+from modules.extractor import BasicEncoder
+from modules.corr import CorrBlock
+from modules.gru import ConvGRU
+from modules.clipping import GradientClip
+
+from lietorch import SE3
+from geom.ba import BA
+
+import geom.projective_ops as pops
+from geom.graph_utils import graph_to_edge_list, keyframe_indicies
+
+from torch_scatter import scatter_mean
+
+
+def cvx_upsample(data, mask):
+    """ upsample pixel-wise transformation field """
+    batch, ht, wd, dim = data.shape
+    data = data.permute(0, 3, 1, 2)
+    mask = mask.view(batch, 1, 9, 8, 8, ht, wd)
+    mask = torch.softmax(mask, dim=2)
+
+    up_data = F.unfold(data, [3,3], padding=1)
+    up_data = up_data.view(batch, dim, 9, 1, 1, ht, wd)
+
+    up_data = torch.sum(mask * up_data, dim=2)
+    up_data = up_data.permute(0, 4, 2, 5, 3, 1)
+    up_data = up_data.reshape(batch, 8*ht, 8*wd, dim)
+
+    return up_data
+
+def upsample_disp(disp, mask):
+    batch, num, ht, wd = disp.shape
+    disp = disp.view(batch*num, ht, wd, 1)
+    mask = mask.view(batch*num, -1, ht, wd)
+    return cvx_upsample(disp, mask).view(batch, num, 8*ht, 8*wd)
+
+
+class GraphAgg(nn.Module):
+    def __init__(self):
+        super(GraphAgg, self).__init__()
+        self.conv1 = nn.Conv2d(128, 128, 3, padding=1)
+        self.conv2 = nn.Conv2d(128, 128, 3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.eta = nn.Sequential(
+            nn.Conv2d(128, 1, 3, padding=1),
+            GradientClip(),
+            nn.Softplus())
+
+        self.upmask = nn.Sequential(
+            nn.Conv2d(128, 8*8*9, 1, padding=0))
+
+    def forward(self, net, ii):
+        batch, num, ch, ht, wd = net.shape
+        net = net.view(batch*num, ch, ht, wd)
+
+        _, ix = torch.unique(ii, return_inverse=True)
+        net = self.relu(self.conv1(net))
+
+        net = net.view(batch, num, 128, ht, wd)
+        net = scatter_mean(net, ix, dim=1)
+        net = net.view(-1, 128, ht, wd)
+
+        net = self.relu(self.conv2(net))
+
+        eta = self.eta(net).view(batch, -1, ht, wd)
+        upmask = self.upmask(net).view(batch, -1, 8*8*9, ht, wd)
+
+        return .01 * eta, upmask
+
+
+class UpdateModule(nn.Module):
+    def __init__(self):
+        super(UpdateModule, self).__init__()
+        cor_planes = 4 * (2*3 + 1)**2
+
+        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(4, 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, 2, 3, padding=1),
+            GradientClip(),
+            nn.Sigmoid())
+
+        self.delta = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 2, 3, padding=1),
+            GradientClip())
+
+        self.gru = ConvGRU(128, 128+128+64)
+        self.agg = GraphAgg()
+
+    def forward(self, net, inp, corr, flow=None, ii=None, jj=None):
+        """ RaftSLAM update operator """
+
+        batch, num, ch, ht, wd = net.shape
+
+        if flow is None:
+            flow = torch.zeros(batch, num, 4, ht, wd, device=net.device)
+
+        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)[...,:2].contiguous()
+        weight = weight.permute(0,1,3,4,2)[...,:2].contiguous()
+
+        net = net.view(*output_dim)
+
+        if ii is not None:
+            eta, upmask = self.agg(net, ii.to(net.device))
+            return net, delta, weight, eta, upmask
+
+        else:
+            return net, delta, weight
+
+
+class DroidNet(nn.Module):
+    def __init__(self):
+        super(DroidNet, self).__init__()
+        self.fnet = BasicEncoder(output_dim=128, norm_fn='instance')
+        self.cnet = BasicEncoder(output_dim=256, norm_fn='none')
+        self.update = UpdateModule()
+
+
+    def extract_features(self, images):
+        """ run feeature extraction networks """
+
+        # normalize images
+        images = images[:, :, [2,1,0]] / 255.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)
+        images = images.sub_(mean[:, None, None]).div_(std[:, None, None])
+
+        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, disps, intrinsics, graph=None, num_steps=12, fixedp=2):
+        """ Estimates SE3 or Sim3 between pair of frames """
+
+        u = keyframe_indicies(graph)
+        ii, jj, kk = graph_to_edge_list(graph)
+
+        ii = ii.to(device=images.device, dtype=torch.long)
+        jj = jj.to(device=images.device, dtype=torch.long)
+
+        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)
+
+        ht, wd = images.shape[-2:]
+        coords0 = pops.coords_grid(ht//8, wd//8, device=images.device)
+        
+        coords1, _ = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+        target = coords1.clone()
+
+        Gs_list, disp_list, residual_list = [], [], []
+        for step in range(num_steps):
+            Gs = Gs.detach()
+            disps = disps.detach()
+            coords1 = coords1.detach()
+            target = target.detach()
+
+            # extract motion features
+            corr = corr_fn(coords1)
+            resd = target - coords1
+            flow = coords1 - coords0
+
+            motion = torch.cat([flow, resd], dim=-1)
+            motion = motion.permute(0,1,4,2,3).clamp(-64.0, 64.0)
+
+            net, delta, weight, eta, upmask = \
+                self.update(net, inp, corr, motion, ii, jj)
+
+            target = coords1 + delta
+
+            for i in range(2):
+                Gs, disps = BA(target, weight, eta, Gs, disps, intrinsics, ii, jj, fixedp=2)
+
+            coords1, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+            residual = (target - coords1)
+
+            Gs_list.append(Gs)
+            disp_list.append(upsample_disp(disps, upmask))
+            residual_list.append(valid_mask * residual)
+
+
+        return Gs_list, disp_list, residual_list

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

@@ -0,0 +1,375 @@
+import torch
+import lietorch
+import numpy as np
+
+import matplotlib.pyplot as plt
+from lietorch import SE3
+from modules.corr import CorrBlock, AltCorrBlock
+import geom.projective_ops as pops
+
+
+class FactorGraph:
+    def __init__(self, video, update_op, device="cuda:0", corr_impl="volume", max_factors=-1, upsample=False):
+        self.video = video
+        self.update_op = update_op
+        self.device = device
+        self.max_factors = max_factors
+        self.corr_impl = corr_impl
+        self.upsample = upsample
+
+        # operator at 1/8 resolution
+        self.ht = ht = video.ht // 8
+        self.wd = wd = video.wd // 8
+
+        self.coords0 = pops.coords_grid(ht, wd, device=device)
+        self.ii = torch.as_tensor([], dtype=torch.long, device=device)
+        self.jj = torch.as_tensor([], dtype=torch.long, device=device)
+        self.age = torch.as_tensor([], dtype=torch.long, device=device)
+
+        self.corr, self.net, self.inp = None, None, None
+        self.damping = 1e-6 * torch.ones_like(self.video.disps)
+
+        self.target = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
+        self.weight = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
+
+        # inactive factors
+        self.ii_inac = torch.as_tensor([], dtype=torch.long, device=device)
+        self.jj_inac = torch.as_tensor([], dtype=torch.long, device=device)
+        self.ii_bad = torch.as_tensor([], dtype=torch.long, device=device)
+        self.jj_bad = torch.as_tensor([], dtype=torch.long, device=device)
+
+        self.target_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
+        self.weight_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
+
+    def __filter_repeated_edges(self, ii, jj):
+        """ remove duplicate edges """
+
+        keep = torch.zeros(ii.shape[0], dtype=torch.bool, device=ii.device)
+        eset = set(
+            [(i.item(), j.item()) for i, j in zip(self.ii, self.jj)] +
+            [(i.item(), j.item()) for i, j in zip(self.ii_inac, self.jj_inac)])
+
+        for k, (i, j) in enumerate(zip(ii, jj)):
+            keep[k] = (i.item(), j.item()) not in eset
+
+        return ii[keep], jj[keep]
+
+    def print_edges(self):
+        ii = self.ii.cpu().numpy()
+        jj = self.jj.cpu().numpy()
+
+        ix = np.argsort(ii)
+        ii = ii[ix]
+        jj = jj[ix]
+
+        w = torch.mean(self.weight, dim=[0,2,3,4]).cpu().numpy()
+        w = w[ix]
+        for e in zip(ii, jj, w):
+            print(e)
+        print()
+
+    def filter_edges(self):
+        """ remove bad edges """
+        conf = torch.mean(self.weight, dim=[0,2,3,4])
+        mask = (torch.abs(self.ii-self.jj) > 2) & (conf < 0.001)
+
+        self.ii_bad = torch.cat([self.ii_bad, self.ii[mask]])
+        self.jj_bad = torch.cat([self.jj_bad, self.jj[mask]])
+        self.rm_factors(mask, store=False)
+
+    def clear_edges(self):
+        self.rm_factors(self.ii >= 0)
+        self.net = None
+        self.inp = None
+
+    @torch.cuda.amp.autocast(enabled=True)
+    def add_factors(self, ii, jj, remove=False):
+        """ add edges to factor graph """
+
+        if not isinstance(ii, torch.Tensor):
+            ii = torch.as_tensor(ii, dtype=torch.long, device=self.device)
+
+        if not isinstance(jj, torch.Tensor):
+            jj = torch.as_tensor(jj, dtype=torch.long, device=self.device)
+
+        # remove duplicate edges
+        ii, jj = self.__filter_repeated_edges(ii, jj)
+
+
+        if ii.shape[0] == 0:
+            return
+
+        # place limit on number of factors
+        if self.max_factors > 0 and self.ii.shape[0] + ii.shape[0] > self.max_factors \
+                and self.corr is not None and remove:
+            
+            ix = torch.arange(len(self.age))[torch.argsort(self.age).cpu()]
+            self.rm_factors(ix >= self.max_factors - ii.shape[0], store=True)
+
+        net = self.video.nets[ii].to(self.device).unsqueeze(0)
+
+        # correlation volume for new edges
+        if self.corr_impl == "volume":
+            c = (ii == jj).long()
+            fmap1 = self.video.fmaps[ii,0].to(self.device).unsqueeze(0)
+            fmap2 = self.video.fmaps[jj,c].to(self.device).unsqueeze(0)
+            corr = CorrBlock(fmap1, fmap2)
+            self.corr = corr if self.corr is None else self.corr.cat(corr)
+
+            inp = self.video.inps[ii].to(self.device).unsqueeze(0)
+            self.inp = inp if self.inp is None else torch.cat([self.inp, inp], 1)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            target, _ = self.video.reproject(ii, jj)
+            weight = torch.zeros_like(target)
+
+        self.ii = torch.cat([self.ii, ii], 0)
+        self.jj = torch.cat([self.jj, jj], 0)
+        self.age = torch.cat([self.age, torch.zeros_like(ii)], 0)
+
+        # reprojection factors
+        self.net = net if self.net is None else torch.cat([self.net, net], 1)
+
+        self.target = torch.cat([self.target, target], 1)
+        self.weight = torch.cat([self.weight, weight], 1)
+
+    @torch.cuda.amp.autocast(enabled=True)
+    def rm_factors(self, mask, store=False):
+        """ drop edges from factor graph """
+
+        # store estimated factors
+        if store:
+            self.ii_inac = torch.cat([self.ii_inac, self.ii[mask]], 0)
+            self.jj_inac = torch.cat([self.jj_inac, self.jj[mask]], 0)
+            self.target_inac = torch.cat([self.target_inac, self.target[:,mask]], 1)
+            self.weight_inac = torch.cat([self.weight_inac, self.weight[:,mask]], 1)
+
+        self.ii = self.ii[~mask]
+        self.jj = self.jj[~mask]
+        self.age = self.age[~mask]
+        
+        if self.corr_impl == "volume":
+            self.corr = self.corr[~mask]
+
+        if self.net is not None:
+            self.net = self.net[:,~mask]
+
+        if self.inp is not None:
+            self.inp = self.inp[:,~mask]
+
+        self.target = self.target[:,~mask]
+        self.weight = self.weight[:,~mask]
+
+
+    @torch.cuda.amp.autocast(enabled=True)
+    def rm_keyframe(self, ix):
+        """ drop edges from factor graph """
+
+
+        with self.video.get_lock():
+            self.video.images[ix] = self.video.images[ix+1]
+            self.video.poses[ix] = self.video.poses[ix+1]
+            self.video.disps[ix] = self.video.disps[ix+1]
+            self.video.disps_sens[ix] = self.video.disps_sens[ix+1]
+            self.video.intrinsics[ix] = self.video.intrinsics[ix+1]
+
+            self.video.nets[ix] = self.video.nets[ix+1]
+            self.video.inps[ix] = self.video.inps[ix+1]
+            self.video.fmaps[ix] = self.video.fmaps[ix+1]
+
+        m = (self.ii_inac == ix) | (self.jj_inac == ix)
+        self.ii_inac[self.ii_inac >= ix] -= 1
+        self.jj_inac[self.jj_inac >= ix] -= 1
+
+        if torch.any(m):
+            self.ii_inac = self.ii_inac[~m]
+            self.jj_inac = self.jj_inac[~m]
+            self.target_inac = self.target_inac[:,~m]
+            self.weight_inac = self.weight_inac[:,~m]
+
+        m = (self.ii == ix) | (self.jj == ix)
+
+        self.ii[self.ii >= ix] -= 1
+        self.jj[self.jj >= ix] -= 1
+        self.rm_factors(m, store=False)
+
+
+    @torch.cuda.amp.autocast(enabled=True)
+    def update(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, motion_only=False):
+        """ run update operator on factor graph """
+
+        # motion features
+        with torch.cuda.amp.autocast(enabled=False):
+            coords1, mask = self.video.reproject(self.ii, self.jj)
+            motn = torch.cat([coords1 - self.coords0, self.target - coords1], dim=-1)
+            motn = motn.permute(0,1,4,2,3).clamp(-64.0, 64.0)
+        
+        # correlation features
+        corr = self.corr(coords1)
+
+        self.net, delta, weight, damping, upmask = \
+            self.update_op(self.net, self.inp, corr, motn, self.ii, self.jj)
+
+        if t0 is None:
+            t0 = max(1, self.ii.min().item()+1)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            self.target = coords1 + delta.to(dtype=torch.float)
+            self.weight = weight.to(dtype=torch.float)
+
+            ht, wd = self.coords0.shape[0:2]
+            self.damping[torch.unique(self.ii)] = damping
+
+            if use_inactive:
+                m = (self.ii_inac >= t0 - 3) & (self.jj_inac >= t0 - 3)
+                ii = torch.cat([self.ii_inac[m], self.ii], 0)
+                jj = torch.cat([self.jj_inac[m], self.jj], 0)
+                target = torch.cat([self.target_inac[:,m], self.target], 1)
+                weight = torch.cat([self.weight_inac[:,m], self.weight], 1)
+
+            else:
+                ii, jj, target, weight = self.ii, self.jj, self.target, self.weight
+
+
+            damping = .2 * self.damping[torch.unique(ii)].contiguous() + EP
+
+            target = target.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
+            weight = weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
+
+            # dense bundle adjustment
+            self.video.ba(target, weight, damping, ii, jj, t0, t1, 
+                itrs=itrs, lm=1e-4, ep=0.1, motion_only=motion_only)
+        
+            if self.upsample:
+                self.video.upsample(torch.unique(self.ii), upmask)
+
+        self.age += 1
+
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def update_lowmem(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, steps=8):
+        """ run update operator on factor graph - reduced memory implementation """
+
+        # alternate corr implementation
+        t = self.video.counter.value
+
+        num, rig, ch, ht, wd = self.video.fmaps.shape
+        corr_op = AltCorrBlock(self.video.fmaps.view(1, num*rig, ch, ht, wd))
+
+        for step in range(steps):
+            print("Global BA Iteration #{}".format(step+1))
+            with torch.cuda.amp.autocast(enabled=False):
+                coords1, mask = self.video.reproject(self.ii, self.jj)
+                motn = torch.cat([coords1 - self.coords0, self.target - coords1], dim=-1)
+                motn = motn.permute(0,1,4,2,3).clamp(-64.0, 64.0)
+
+            s = 8
+            for i in range(0, self.jj.max()+1, s):
+                v = (self.ii >= i) & (self.ii < i + s)
+                iis = self.ii[v]
+                jjs = self.jj[v]
+
+                ht, wd = self.coords0.shape[0:2]
+                corr1 = corr_op(coords1[:,v], rig * iis, rig * jjs + (iis == jjs).long())
+
+                with torch.cuda.amp.autocast(enabled=True):
+                 
+                    net, delta, weight, damping, upmask = \
+                        self.update_op(self.net[:,v], self.video.inps[None,iis], corr1, motn[:,v], iis, jjs)
+
+                    if self.upsample:
+                        self.video.upsample(torch.unique(iis), upmask)
+
+                self.net[:,v] = net
+                self.target[:,v] = coords1[:,v] + delta.float()
+                self.weight[:,v] = weight.float()
+                self.damping[torch.unique(iis)] = damping
+
+            damping = .2 * self.damping[torch.unique(self.ii)].contiguous() + EP
+            target = self.target.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
+            weight = self.weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
+
+            # dense bundle adjustment
+            self.video.ba(target, weight, damping, self.ii, self.jj, 1, t, 
+                itrs=itrs, lm=1e-5, ep=1e-2, motion_only=False)
+
+            self.video.dirty[:t] = True
+
+    def add_neighborhood_factors(self, t0, t1, r=3):
+        """ add edges between neighboring frames within radius r """
+
+        ii, jj = torch.meshgrid(torch.arange(t0,t1), torch.arange(t0,t1))
+        ii = ii.reshape(-1).to(dtype=torch.long, device=self.device)
+        jj = jj.reshape(-1).to(dtype=torch.long, device=self.device)
+
+        c = 1 if self.video.stereo else 0
+
+        keep = ((ii - jj).abs() > c) & ((ii - jj).abs() <= r)
+        self.add_factors(ii[keep], jj[keep])
+
+    
+    def add_proximity_factors(self, t0=0, t1=0, rad=2, nms=2, beta=0.25, thresh=16.0, remove=False):
+        """ add edges to the factor graph based on distance """
+
+        t = self.video.counter.value
+        ix = torch.arange(t0, t)
+        jx = torch.arange(t1, t)
+
+        ii, jj = torch.meshgrid(ix, jx)
+        ii = ii.reshape(-1)
+        jj = jj.reshape(-1)
+
+        d = self.video.distance(ii, jj, beta=beta)
+        d[ii - rad < jj] = np.inf
+        d[d > 100] = np.inf
+
+        ii1 = torch.cat([self.ii, self.ii_bad, self.ii_inac], 0)
+        jj1 = torch.cat([self.jj, self.jj_bad, self.jj_inac], 0)
+        for i, j in zip(ii1.cpu().numpy(), jj1.cpu().numpy()):
+            for di in range(-nms, nms+1):
+                for dj in range(-nms, nms+1):
+                    if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0):
+                        i1 = i + di
+                        j1 = j + dj
+
+                        if (t0 <= i1 < t) and (t1 <= j1 < t):
+                            d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf
+
+
+        es = []
+        for i in range(t0, t):
+            if self.video.stereo:
+                es.append((i, i))
+                d[(i-t0)*(t-t1) + (i-t1)] = np.inf
+
+            for j in range(max(i-rad-1,0), i):
+                es.append((i,j))
+                es.append((j,i))
+                d[(i-t0)*(t-t1) + (j-t1)] = np.inf
+
+        ix = torch.argsort(d)
+        for k in ix:
+            if d[k].item() > thresh:
+                continue
+
+            if len(es) > self.max_factors:
+                break
+
+            i = ii[k]
+            j = jj[k]
+            
+            # bidirectional
+            es.append((i, j))
+            es.append((j, i))
+
+            for di in range(-nms, nms+1):
+                for dj in range(-nms, nms+1):
+                    if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0):
+                        i1 = i + di
+                        j1 = j + dj
+
+                        if (t0 <= i1 < t) and (t1 <= j1 < t):
+                            d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf
+
+        ii, jj = torch.as_tensor(es, device=self.device).unbind(dim=-1)
+        self.add_factors(ii, jj, remove)

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/geom/ba.py

@@ -0,0 +1,158 @@
+import lietorch
+import torch
+import torch.nn.functional as F
+
+from .chol import block_solve, schur_solve
+import geom.projective_ops as pops
+
+from torch_scatter import scatter_sum
+
+
+# utility functions for scattering ops
+def safe_scatter_add_mat(A, ii, jj, n, m):
+    v = (ii >= 0) & (jj >= 0) & (ii < n) & (jj < m)
+    return scatter_sum(A[:,v], ii[v]*m + jj[v], dim=1, dim_size=n*m)
+
+def safe_scatter_add_vec(b, ii, n):
+    v = (ii >= 0) & (ii < n)
+    return scatter_sum(b[:,v], ii[v], dim=1, dim_size=n)
+
+# apply retraction operator to inv-depth maps
+def disp_retr(disps, dz, ii):
+    ii = ii.to(device=dz.device)
+    return disps + scatter_sum(dz, ii, dim=1, dim_size=disps.shape[1])
+
+# apply retraction operator to poses
+def pose_retr(poses, dx, ii):
+    ii = ii.to(device=dx.device)
+    return poses.retr(scatter_sum(dx, ii, dim=1, dim_size=poses.shape[1]))
+
+
+def BA(target, weight, eta, poses, disps, intrinsics, ii, jj, fixedp=1, rig=1):
+    """ Full Bundle Adjustment """
+
+    B, P, ht, wd = disps.shape
+    N = ii.shape[0]
+    D = poses.manifold_dim
+
+    ### 1: commpute jacobians and residuals ###
+    coords, valid, (Ji, Jj, Jz) = pops.projective_transform(
+        poses, disps, intrinsics, ii, jj, jacobian=True)
+
+    r = (target - coords).view(B, N, -1, 1)
+    w = .001 * (valid * weight).view(B, N, -1, 1)
+
+    ### 2: construct linear system ###
+    Ji = Ji.reshape(B, N, -1, D)
+    Jj = Jj.reshape(B, N, -1, D)
+    wJiT = (w * Ji).transpose(2,3)
+    wJjT = (w * Jj).transpose(2,3)
+
+    Jz = Jz.reshape(B, N, ht*wd, -1)
+
+    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)
+
+    Ei = (wJiT.view(B,N,D,ht*wd,-1) * Jz[:,:,None]).sum(dim=-1)
+    Ej = (wJjT.view(B,N,D,ht*wd,-1) * Jz[:,:,None]).sum(dim=-1)
+
+    w = w.view(B, N, ht*wd, -1)
+    r = r.view(B, N, ht*wd, -1)
+    wk = torch.sum(w*r*Jz, dim=-1)
+    Ck = torch.sum(w*Jz*Jz, dim=-1)
+
+    kx, kk = torch.unique(ii, return_inverse=True)
+    M = kx.shape[0]
+
+    # only optimize keyframe poses
+    P = P // rig - fixedp
+    ii = ii // rig - fixedp
+    jj = jj // rig - fixedp
+
+    H = safe_scatter_add_mat(Hii, ii, ii, P, P) + \
+        safe_scatter_add_mat(Hij, ii, jj, P, P) + \
+        safe_scatter_add_mat(Hji, jj, ii, P, P) + \
+        safe_scatter_add_mat(Hjj, jj, jj, P, P)
+
+    E = safe_scatter_add_mat(Ei, ii, kk, P, M) + \
+        safe_scatter_add_mat(Ej, jj, kk, P, M)
+
+    v = safe_scatter_add_vec(vi, ii, P) + \
+        safe_scatter_add_vec(vj, jj, P)
+
+    C = safe_scatter_add_vec(Ck, kk, M)
+    w = safe_scatter_add_vec(wk, kk, M)
+
+    C = C + eta.view(*C.shape) + 1e-7
+
+    H = H.view(B, P, P, D, D)
+    E = E.view(B, P, M, D, ht*wd)
+
+    ### 3: solve the system ###
+    dx, dz = schur_solve(H, E, C, v, w)
+    
+    ### 4: apply retraction ###
+    poses = pose_retr(poses, dx, torch.arange(P) + fixedp)
+    disps = disp_retr(disps, dz.view(B,-1,ht,wd), kx)
+
+    disps = torch.where(disps > 10, torch.zeros_like(disps), disps)
+    disps = disps.clamp(min=0.0)
+
+    return poses, disps
+
+
+def MoBA(target, weight, eta, poses, disps, intrinsics, ii, jj, fixedp=1, rig=1):
+    """ Motion only bundle adjustment """
+
+    B, P, ht, wd = disps.shape
+    N = ii.shape[0]
+    D = poses.manifold_dim
+
+    ### 1: commpute jacobians and residuals ###
+    coords, valid, (Ji, Jj, Jz) = pops.projective_transform(
+        poses, disps, intrinsics, ii, jj, jacobian=True)
+
+    r = (target - coords).view(B, N, -1, 1)
+    w = .001 * (valid * weight).view(B, N, -1, 1)
+
+    ### 2: construct linear system ###
+    Ji = Ji.reshape(B, N, -1, D)
+    Jj = Jj.reshape(B, N, -1, D)
+    wJiT = (w * Ji).transpose(2,3)
+    wJjT = (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
+    P = P // rig - fixedp
+    ii = ii // rig - fixedp
+    jj = jj // rig - fixedp
+
+    H = safe_scatter_add_mat(Hii, ii, ii, P, P) + \
+        safe_scatter_add_mat(Hij, ii, jj, P, P) + \
+        safe_scatter_add_mat(Hji, jj, ii, P, P) + \
+        safe_scatter_add_mat(Hjj, jj, jj, P, P)
+
+    v = safe_scatter_add_vec(vi, ii, P) + \
+        safe_scatter_add_vec(vj, jj, P)
+    
+    H = H.view(B, P, P, D, D)
+
+    ### 3: solve the system ###
+    dx = block_solve(H, v)
+
+    ### 4: apply retraction ###
+    poses = pose_retr(poses, dx, torch.arange(P) + fixedp)
+    return poses
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/geom/chol.py

@@ -0,0 +1,73 @@
+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.linalg.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)
+
+
+def schur_solve(H, E, C, v, w, ep=0.1, lm=0.0001, sless=False):
+    """ solve using shur complement """
+    
+    B, P, M, D, HW = E.shape
+    H = H.permute(0,1,3,2,4).reshape(B, P*D, P*D)
+    E = E.permute(0,1,3,2,4).reshape(B, P*D, M*HW)
+    Q = (1.0 / C).view(B, M*HW, 1)
+
+    # damping
+    I = torch.eye(P*D).to(H.device)
+    H = H + (ep + lm*H) * I
+    
+    v = v.reshape(B, P*D, 1)
+    w = w.reshape(B, M*HW, 1)
+
+    Et = E.transpose(1,2)
+    S = H - torch.matmul(E, Q*Et)
+    v = v - torch.matmul(E, Q*w)
+
+    dx = CholeskySolver.apply(S, v)
+    if sless:
+        return dx.reshape(B, P, D)
+
+    dz = Q * (w - Et @ dx)    
+    dx = dx.reshape(B, P, D)
+    dz = dz.reshape(B, M, HW)
+
+    return dx, dz

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/geom/graph_utils.py

@@ -0,0 +1,113 @@
+
+import torch
+import numpy as np
+from collections import OrderedDict
+
+import lietorch
+from data_readers.rgbd_utils import compute_distance_matrix_flow, compute_distance_matrix_flow2
+
+
+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)
+    jj = torch.as_tensor(jj)
+    kk = torch.as_tensor(kk)
+    return ii, jj, kk
+
+def keyframe_indicies(graph):
+    return torch.as_tensor([u for u in graph])
+
+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 neighbourhood_graph(n, r):
+    ii, jj = meshgrid(n, n)
+    d = (ii - jj).abs()
+    keep = (d >= 1) & (d <= r)
+    return ii[keep], jj[keep]
+
+
+def build_frame_graph(poses, disps, intrinsics, num=16, thresh=24.0, r=2):
+    """ construct a frame graph between co-visible frames """
+    N = poses.shape[1]
+    poses = poses[0].cpu().numpy()
+    disps = disps[0][:,3::8,3::8].cpu().numpy()
+    intrinsics = intrinsics[0].cpu().numpy() / 8.0
+    d = compute_distance_matrix_flow(poses, disps, intrinsics)
+
+    count = 0
+    graph = OrderedDict()
+    
+    for i in range(N):
+        graph[i] = []
+        d[i,i] = np.inf
+        for j in range(i-r, i+r+1):
+            if 0 <= j < N and i != j:
+                graph[i].append(j)
+                d[i,j] = np.inf
+                count += 1
+
+    while count < num:
+        ix = np.argmin(d)
+        i, j = ix // N, ix % N
+
+        if d[i,j] < thresh:
+            graph[i].append(j)
+            d[i,j] = np.inf
+            count += 1
+        else:
+            break
+    
+    return graph
+
+
+
+def build_frame_graph_v2(poses, disps, intrinsics, num=16, thresh=24.0, r=2):
+    """ construct a frame graph between co-visible frames """
+    N = poses.shape[1]
+    # poses = poses[0].cpu().numpy()
+    # disps = disps[0].cpu().numpy()
+    # intrinsics = intrinsics[0].cpu().numpy()
+    d = compute_distance_matrix_flow2(poses, disps, intrinsics)
+
+    # import matplotlib.pyplot as plt
+    # plt.imshow(d)
+    # plt.show()
+
+    count = 0
+    graph = OrderedDict()
+    
+    for i in range(N):
+        graph[i] = []
+        d[i,i] = np.inf
+        for j in range(i-r, i+r+1):
+            if 0 <= j < N and i != j:
+                graph[i].append(j)
+                d[i,j] = np.inf
+                count += 1
+
+    while 1:
+        ix = np.argmin(d)
+        i, j = ix // N, ix % N
+
+        if d[i,j] < thresh:
+            graph[i].append(j)
+
+            for i1 in range(i-1, i+2):
+                for j1 in range(j-1, j+2):
+                    if 0 <= i1 < N and 0 <= j1 < N:
+                        d[i1, j1] = np.inf
+
+            count += 1
+        else:
+            break
+    
+    return graph
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/geom/losses.py

@@ -0,0 +1,118 @@
+from collections import OrderedDict
+import numpy as np
+import torch
+from lietorch import SO3, SE3, Sim3
+from .graph_utils import graph_to_edge_list
+from .projective_ops import projective_transform
+
+
+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 fit_scale(Ps, Gs):
+    b = Ps.shape[0]
+    t1 = Ps.data[...,:3].detach().reshape(b, -1)
+    t2 = Gs.data[...,:3].detach().reshape(b, -1)
+
+    s = (t1*t2).sum(-1) / ((t2*t2).sum(-1) + 1e-8)
+    return s
+
+
+def geodesic_loss(Ps, Gs, graph, gamma=0.9, do_scale=True):
+    """ 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()
+
+        if do_scale:
+            s = fit_scale(dP, dG)
+            dG = dG.scale(s[:,None])
+        
+        # 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 = {
+        'rot_error': r_err.mean().item(),
+        'tr_error': t_err.mean().item(),
+        'bad_rot': (r_err < .1).float().mean().item(),
+        'bad_tr': (t_err < .01).float().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()}
+
+
+def flow_loss(Ps, disps, poses_est, disps_est, intrinsics, graph, gamma=0.9):
+    """ optical flow loss """
+
+    N = Ps.shape[1]
+    graph = OrderedDict()
+    for i in range(N):
+        graph[i] = [j for j in range(N) if abs(i-j)==1]
+
+    ii, jj, kk = graph_to_edge_list(graph)
+    coords0, val0 = projective_transform(Ps, disps, intrinsics, ii, jj)
+    val0 = val0 * (disps[:,ii] > 0).float().unsqueeze(dim=-1)
+
+    n = len(poses_est)
+    flow_loss = 0.0
+
+    for i in range(n):
+        w = gamma ** (n - i - 1)
+        coords1, val1 = projective_transform(poses_est[i], disps_est[i], intrinsics, ii, jj)
+
+        v = (val0 * val1).squeeze(dim=-1)
+        epe = v * (coords1 - coords0).norm(dim=-1)
+        flow_loss += w * epe.mean()
+
+    epe = epe.reshape(-1)[v.reshape(-1) > 0.5]
+    metrics = {
+        'f_error': epe.mean().item(),
+        '1px': (epe<1.0).float().mean().item(),
+    }
+
+    return flow_loss, metrics

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/geom/projective_ops.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn.functional as F
+
+from lietorch import SE3, Sim3
+
+MIN_DEPTH = 0.2
+
+def extract_intrinsics(intrinsics):
+    return intrinsics[...,None,None,:].unbind(dim=-1)
+
+def coords_grid(ht, wd, **kwargs):
+    y, x = torch.meshgrid(
+        torch.arange(ht).to(**kwargs).float(),
+        torch.arange(wd).to(**kwargs).float())
+
+    return torch.stack([x, y], dim=-1)
+
+def iproj(disps, intrinsics, jacobian=False):
+    """ 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
+    pts = torch.stack([X, Y, i, disps], dim=-1)
+
+    if jacobian:
+        J = torch.zeros_like(pts)
+        J[...,-1] = 1.0
+        return pts, J
+
+    return pts, None
+
+def proj(Xs, intrinsics, jacobian=False, return_depth=False):
+    """ pinhole camera projection """
+    fx, fy, cx, cy = extract_intrinsics(intrinsics)
+    X, Y, Z, D = Xs.unbind(dim=-1)
+
+    Z = torch.where(Z < 0.5*MIN_DEPTH, torch.ones_like(Z), Z)
+    d = 1.0 / Z
+
+    x = fx * (X * d) + cx
+    y = fy * (Y * d) + cy
+    if return_depth:
+        coords = torch.stack([x, y, D*d], dim=-1)
+    else:
+        coords = torch.stack([x, y], 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, 2, 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, return_depth=False):
+    """ map points from ii->jj """
+
+    # inverse project (pinhole)
+    X0, Jz = iproj(depths[:,ii], intrinsics[:,ii], jacobian=jacobian)
+    
+    # transform
+    Gij = poses[:,jj] * poses[:,ii].inv()
+
+    Gij.data[:,ii==jj] = torch.as_tensor([-0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], device="cuda")
+    X1, Ja = actp(Gij, X0, jacobian=jacobian)
+    
+    # project (pinhole)
+    x1, Jp = proj(X1, intrinsics[:,jj], jacobian=jacobian, return_depth=return_depth)
+
+    # exclude points too close to camera
+    valid = ((X1[...,2] > MIN_DEPTH) & (X0[...,2] > MIN_DEPTH)).float()
+    valid = valid.unsqueeze(-1)
+
+    if jacobian:
+        # Ji transforms according to dual adjoint
+        Jj = torch.matmul(Jp, Ja)
+        Ji = -Gij[:,:,None,None,None].adjT(Jj)
+
+        Jz = Gij[:,:,None,None] * Jz
+        Jz = torch.matmul(Jp, Jz.unsqueeze(-1))
+
+        return x1, valid, (Ji, Jj, Jz)
+
+    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, False)
+
+    return coords1[...,:2] - coords0, valid
+

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

@@ -0,0 +1,54 @@
+
+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
+
+    def write_dict(self, results):
+        for key in results:
+            self.writer.add_scalar(key, results[key], self.total_steps)
+
+    def close(self):
+        self.writer.close()
+

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/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/droid_slam/modules/corr.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn.functional as F
+
+import droid_backends
+
+class CorrSampler(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, volume, coords, radius):
+        ctx.save_for_backward(volume,coords)
+        ctx.radius = radius
+        corr, = droid_backends.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, = droid_backends.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=3):
+        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 cat(self, other):
+        for i in range(self.num_levels):
+            self.corr_pyramid[i] = torch.cat([self.corr_pyramid[i], other.corr_pyramid[i]], 0)
+        return self
+
+    def __getitem__(self, index):
+        for i in range(self.num_levels):
+            self.corr_pyramid[i] = self.corr_pyramid[i][index]
+        return self
+
+
+    @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
+        ctx.save_for_backward(fmap1, fmap2, coords)
+        corr, = droid_backends.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 = \
+            droid_backends.altcorr_backward(fmap1, fmap2, coords, grad_corr, ctx.r)
+        return fmap1_grad, fmap2_grad, coords_grad, None
+
+
+class AltCorrBlock:
+    def __init__(self, fmaps, num_levels=4, radius=3):
+        self.num_levels = num_levels
+        self.radius = radius
+
+        B, N, C, H, W = fmaps.shape
+        fmaps = fmaps.view(B*N, C, H, W) / 4.0
+        
+        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).contiguous()
+            self.pyramid.append(fmap_lvl.view(*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.float(), fmap2_i.float(), 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
+
+
+    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/droid_slam/modules/extractor.py

@@ -0,0 +1,198 @@
+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)

third_party/Dyn-HaMR/third-party/DROID-SLAM/droid_slam/modules/gru.py

@@ -0,0 +1,34 @@
+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)
+
+        self.w = nn.Conv2d(h_planes, h_planes, 1, padding=0)
+
+        self.convz_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
+        self.convr_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
+        self.convq_glo = nn.Conv2d(h_planes, h_planes, 1, padding=0)
+
+    def forward(self, net, *inputs):
+        inp = torch.cat(inputs, dim=1)
+        net_inp = torch.cat([net, inp], dim=1)
+
+        b, c, h, w = net.shape
+        glo = torch.sigmoid(self.w(net)) * net
+        glo = glo.view(b, c, h*w).mean(-1).view(b, c, 1, 1)
+
+        z = torch.sigmoid(self.convz(net_inp) + self.convz_glo(glo))
+        r = torch.sigmoid(self.convr(net_inp) + self.convr_glo(glo))
+        q = torch.tanh(self.convq(torch.cat([r*net, inp], dim=1)) + self.convq_glo(glo))
+
+        net = (1-z) * net + z * q
+        return net
+
+

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

@@ -0,0 +1,159 @@
+import cv2
+import torch
+import lietorch
+
+from collections import OrderedDict
+from droid_net import DroidNet
+
+import geom.projective_ops as pops
+from modules.corr import CorrBlock
+
+
+class MotionFilter:
+    """ This class is used to filter incoming frames and extract features """
+
+    def __init__(self, net, video, thresh=2.5, device="cuda:0"):
+        
+        # split net modules
+        self.cnet = net.cnet
+        self.fnet = net.fnet
+        self.update = net.update
+
+        self.video = video
+        self.thresh = thresh
+        self.device = device
+
+        self.count = 0
+
+        # mean, std for image normalization
+        self.MEAN = torch.as_tensor([0.485, 0.456, 0.406], device=self.device)[:, None, None]
+        self.STDV = torch.as_tensor([0.229, 0.224, 0.225], device=self.device)[:, None, None]
+        
+    @torch.cuda.amp.autocast(enabled=True)
+    def __context_encoder(self, image):
+        """ context features """
+        net, inp = self.cnet(image).split([128,128], dim=2)
+        return net.tanh().squeeze(0), inp.relu().squeeze(0)
+
+    @torch.cuda.amp.autocast(enabled=True)
+    def __feature_encoder(self, image):
+        """ features for correlation volume """
+        return self.fnet(image).squeeze(0)
+
+    @torch.cuda.amp.autocast(enabled=True)
+    @torch.no_grad()
+    def track(self, tstamp, image, depth=None, intrinsics=None):
+        """ main update operation - run on every frame in video """
+
+        Id = lietorch.SE3.Identity(1,).data.squeeze()
+        ht = image.shape[-2] // 8
+        wd = image.shape[-1] // 8
+
+        # normalize images
+        inputs = image[None, :, [2,1,0]].to(self.device) / 255.0
+        inputs = inputs.sub_(self.MEAN).div_(self.STDV)
+
+        # extract features
+        gmap = self.__feature_encoder(inputs)
+
+        ### always add first frame to the depth video ###
+        if self.video.counter.value == 0:
+            net, inp = self.__context_encoder(inputs[:,[0]])
+            self.net, self.inp, self.fmap = net, inp, gmap
+            self.video.append(tstamp, image[0], Id, 1.0, depth, intrinsics / 8.0, gmap, net[0,0], inp[0,0])
+
+        ### only add new frame if there is enough motion ###
+        else:                
+            # index correlation volume
+            coords0 = pops.coords_grid(ht, wd, device=self.device)[None,None]
+            corr = CorrBlock(self.fmap[None,[0]], gmap[None,[0]])(coords0)
+
+            # approximate flow magnitude using 1 update iteration
+            _, delta, weight = self.update(self.net[None], self.inp[None], corr)
+
+            # check motion magnitue / add new frame to video
+            if delta.norm(dim=-1).mean().item() > self.thresh:
+                self.count = 0
+                net, inp = self.__context_encoder(inputs[:,[0]])
+                self.net, self.inp, self.fmap = net, inp, gmap
+                self.video.append(tstamp, image[0], None, None, depth, intrinsics / 8.0, gmap, net[0], inp[0])
+
+            else:
+                self.count += 1
+
+
+
+
+# class MotionFilter:
+#     """ This class is used to filter incoming frames and extract features """
+
+#     def __init__(self, net, video, thresh=2.5, device="cuda:0"):
+        
+#         # split net modules
+#         self.cnet = net.cnet
+#         self.fnet = net.fnet
+#         self.update = net.update
+
+#         self.video = video
+#         self.thresh = thresh
+#         self.device = device
+
+#         self.count = 0
+
+#         # mean, std for image normalization
+#         self.MEAN = torch.as_tensor([0.485, 0.456, 0.406], device=self.device)[:, None, None]
+#         self.STDV = torch.as_tensor([0.229, 0.224, 0.225], device=self.device)[:, None, None]
+        
+#     @torch.cuda.amp.autocast(enabled=True)
+#     def __context_encoder(self, image):
+#         """ context features """
+#         x = self.cnet(image)
+#         net, inp = self.cnet(image).split([128,128], dim=2)
+#         return net.tanh().squeeze(0), inp.relu().squeeze(0)
+
+#     @torch.cuda.amp.autocast(enabled=True)
+#     def __feature_encoder(self, image):
+#         """ features for correlation volume """
+#         return self.fnet(image).squeeze(0)
+
+#     @torch.cuda.amp.autocast(enabled=True)
+#     @torch.no_grad()
+#     def track(self, tstamp, image, depth=None, intrinsics=None):
+#         """ main update operation - run on every frame in video """
+
+#         Id = lietorch.SE3.Identity(1,).data.squeeze()
+#         ht = image.shape[-2] // 8
+#         wd = image.shape[-1] // 8
+
+#         # normalize images
+#         inputs = image[None, None, [2,1,0]].to(self.device) / 255.0
+#         inputs = inputs.sub_(self.MEAN).div_(self.STDV)
+
+#         # extract features
+#         gmap = self.__feature_encoder(inputs)
+
+#         ### always add first frame to the depth video ###
+#         if self.video.counter.value == 0:
+#             net, inp = self.__context_encoder(inputs)
+#             self.net, self.inp, self.fmap = net, inp, gmap
+#             self.video.append(tstamp, image, Id, 1.0, intrinsics / 8.0, gmap[0], net[0], inp[0])
+
+#         ### only add new frame if there is enough motion ###
+#         else:                
+#             # index correlation volume
+#             coords0 = pops.coords_grid(ht, wd, device=self.device)[None,None]
+#             corr = CorrBlock(self.fmap[None], gmap[None])(coords0)
+
+#             # approximate flow magnitude using 1 update iteration
+#             _, delta, weight = self.update(self.net[None], self.inp[None], corr)
+
+#             # check motion magnitue / add new frame to video
+#             if delta.norm(dim=-1).mean().item() > self.thresh:
+#                 self.count = 0
+#                 net, inp = self.__context_encoder(inputs)
+#                 self.net, self.inp, self.fmap = net, inp, gmap
+#                 self.video.append(tstamp, image, None, None, intrinsics / 8.0, gmap[0], net[0], inp[0])
+
+#             else:
+#                 self.count += 1
+

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

@@ -0,0 +1,104 @@
+import cv2
+import torch
+import lietorch
+
+from lietorch import SE3
+from collections import OrderedDict
+from factor_graph import FactorGraph
+from droid_net import DroidNet
+import geom.projective_ops as pops
+
+
+class PoseTrajectoryFiller:
+    """ This class is used to fill in non-keyframe poses """
+
+    def __init__(self, net, video, device="cuda:0"):
+        
+        # split net modules
+        self.cnet = net.cnet
+        self.fnet = net.fnet
+        self.update = net.update
+
+        self.count = 0
+        self.video = video
+        self.device = device
+
+        # mean, std for image normalization
+        self.MEAN = torch.as_tensor([0.485, 0.456, 0.406], device=self.device)[:, None, None]
+        self.STDV = torch.as_tensor([0.229, 0.224, 0.225], device=self.device)[:, None, None]
+        
+    @torch.cuda.amp.autocast(enabled=True)
+    def __feature_encoder(self, image):
+        """ features for correlation volume """
+        return self.fnet(image)
+
+    def __fill(self, tstamps, images, intrinsics):
+        """ fill operator """
+
+        tt = torch.as_tensor(tstamps, device="cuda")
+        images = torch.stack(images, 0)
+        intrinsics = torch.stack(intrinsics, 0)
+        inputs = images[:,:,[2,1,0]].to(self.device) / 255.0
+        
+        ### linear pose interpolation ###
+        N = self.video.counter.value
+        M = len(tstamps)
+
+        ts = self.video.tstamp[:N]
+        Ps = SE3(self.video.poses[:N])
+
+        t0 = torch.as_tensor([ts[ts<=t].shape[0] - 1 for t in tstamps])
+        t1 = torch.where(t0<N-1, t0+1, t0)
+
+        dt = ts[t1] - ts[t0] + 1e-3
+        dP = Ps[t1] * Ps[t0].inv()
+
+        v = dP.log() / dt.unsqueeze(-1)
+        w = v * (tt - ts[t0]).unsqueeze(-1)
+        Gs = SE3.exp(w) * Ps[t0]
+
+        # extract features (no need for context features)
+        inputs = inputs.sub_(self.MEAN).div_(self.STDV)
+        fmap = self.__feature_encoder(inputs)
+
+        self.video.counter.value += M
+        self.video[N:N+M] = (tt, images[:,0], Gs.data, 1, None, intrinsics / 8.0, fmap)
+
+        graph = FactorGraph(self.video, self.update)
+        graph.add_factors(t0.cuda(), torch.arange(N, N+M).cuda())
+        graph.add_factors(t1.cuda(), torch.arange(N, N+M).cuda())
+
+        for itr in range(6):
+            graph.update(N, N+M, motion_only=True)
+    
+        Gs = SE3(self.video.poses[N:N+M].clone())
+        self.video.counter.value -= M
+
+        return [ Gs ]
+
+    @torch.no_grad()
+    def __call__(self, image_stream):
+        """ fill in poses of non-keyframe images """
+
+        # store all camera poses
+        pose_list = []
+
+        tstamps = []
+        images = []
+        intrinsics = []
+        
+        for (tstamp, image, intrinsic) in image_stream:
+            tstamps.append(tstamp)
+            images.append(image)
+            intrinsics.append(intrinsic)
+
+            if len(tstamps) == 16:
+                pose_list += self.__fill(tstamps, images, intrinsics)
+                tstamps, images, intrinsics = [], [], []
+
+        if len(tstamps) > 0:
+            pose_list += self.__fill(tstamps, images, intrinsics)
+
+        # stitch pose segments together
+        return lietorch.cat(pose_list, 0)
+

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

@@ -0,0 +1,154 @@
+import torch
+import cv2
+import lietorch
+import droid_backends
+import time
+import argparse
+import numpy as np
+import open3d as o3d
+
+from lietorch import SE3
+import geom.projective_ops as pops
+
+CAM_POINTS = 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]])
+
+def white_balance(img):
+    # from https://stackoverflow.com/questions/46390779/automatic-white-balancing-with-grayworld-assumption
+    result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+    avg_a = np.average(result[:, :, 1])
+    avg_b = np.average(result[:, :, 2])
+    result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1)
+    result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1)
+    result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
+    return result
+
+def create_camera_actor(g, scale=0.05):
+    """ build open3d camera polydata """
+    camera_actor = o3d.geometry.LineSet(
+        points=o3d.utility.Vector3dVector(scale * CAM_POINTS),
+        lines=o3d.utility.Vector2iVector(CAM_LINES))
+
+    color = (g * 1.0, 0.5 * (1-g), 0.9 * (1-g))
+    camera_actor.paint_uniform_color(color)
+    return camera_actor
+
+def create_point_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 droid_visualization(video, device="cuda:0"):
+    """ DROID visualization frontend """
+
+    torch.cuda.set_device(device)
+    droid_visualization.video = video
+    droid_visualization.cameras = {}
+    droid_visualization.points = {}
+    droid_visualization.warmup = 8
+    droid_visualization.scale = 1.0
+    droid_visualization.ix = 0
+
+    droid_visualization.filter_thresh = 0.005
+
+    def increase_filter(vis):
+        droid_visualization.filter_thresh *= 2
+        with droid_visualization.video.get_lock():
+            droid_visualization.video.dirty[:droid_visualization.video.counter.value] = True
+
+    def decrease_filter(vis):
+        droid_visualization.filter_thresh *= 0.5
+        with droid_visualization.video.get_lock():
+            droid_visualization.video.dirty[:droid_visualization.video.counter.value] = True
+
+    def animation_callback(vis):
+        cam = vis.get_view_control().convert_to_pinhole_camera_parameters()
+
+        with torch.no_grad():
+
+            with video.get_lock():
+                t = video.counter.value 
+                dirty_index, = torch.where(video.dirty.clone())
+                dirty_index = dirty_index
+
+            if len(dirty_index) == 0:
+                return
+
+            video.dirty[dirty_index] = False
+
+            # convert poses to 4x4 matrix
+            poses = torch.index_select(video.poses, 0, dirty_index)
+            disps = torch.index_select(video.disps, 0, dirty_index)
+            Ps = SE3(poses).inv().matrix().cpu().numpy()
+
+            images = torch.index_select(video.images, 0, dirty_index)
+            images = images.cpu()[:,[2,1,0],3::8,3::8].permute(0,2,3,1) / 255.0
+            points = droid_backends.iproj(SE3(poses).inv().data, disps, video.intrinsics[0]).cpu()
+
+            thresh = droid_visualization.filter_thresh * torch.ones_like(disps.mean(dim=[1,2]))
+            
+            count = droid_backends.depth_filter(
+                video.poses, video.disps, video.intrinsics[0], dirty_index, thresh)
+
+            count = count.cpu()
+            disps = disps.cpu()
+            masks = ((count >= 2) & (disps > .5*disps.mean(dim=[1,2], keepdim=True)))
+            
+            for i in range(len(dirty_index)):
+                pose = Ps[i]
+                ix = dirty_index[i].item()
+
+                if ix in droid_visualization.cameras:
+                    vis.remove_geometry(droid_visualization.cameras[ix])
+                    del droid_visualization.cameras[ix]
+
+                if ix in droid_visualization.points:
+                    vis.remove_geometry(droid_visualization.points[ix])
+                    del droid_visualization.points[ix]
+
+                ### add camera actor ###
+                cam_actor = create_camera_actor(True)
+                cam_actor.transform(pose)
+                vis.add_geometry(cam_actor)
+                droid_visualization.cameras[ix] = cam_actor
+
+                mask = masks[i].reshape(-1)
+                pts = points[i].reshape(-1, 3)[mask].cpu().numpy()
+                clr = images[i].reshape(-1, 3)[mask].cpu().numpy()
+                
+                ## add point actor ###
+                point_actor = create_point_actor(pts, clr)
+                vis.add_geometry(point_actor)
+                droid_visualization.points[ix] = point_actor
+
+            # hack to allow interacting with vizualization during inference
+            if len(droid_visualization.cameras) >= droid_visualization.warmup:
+                cam = vis.get_view_control().convert_from_pinhole_camera_parameters(cam)
+
+            droid_visualization.ix += 1
+            vis.poll_events()
+            vis.update_renderer()
+
+    ### create Open3D visualization ###
+    vis = o3d.visualization.VisualizerWithKeyCallback()
+    vis.register_animation_callback(animation_callback)
+    vis.register_key_callback(ord("S"), increase_filter)
+    vis.register_key_callback(ord("A"), decrease_filter)
+
+    vis.create_window(height=540, width=960)
+    vis.get_render_option().load_from_json("misc/renderoption.json")
+
+    vis.run()
+    vis.destroy_window()