option to save reconstruction w/ full res depths

Author: zachteed

README.md

@@ -58,7 +58,7 @@ python setup.py install
 ./tools/download_sample_data.sh
 ```
 
-Run the demo on any of the samples (all demos can be run on a GPU with 11G of memory). While running, press the "s" key to increase the filtering threshold (= more points) and "a" to decrease the filtering threshold (= fewer points).
+Run the demo on any of the samples (all demos can be run on a GPU with 11G of memory). While running, press the "s" key to increase the filtering threshold (= more points) and "a" to decrease the filtering threshold (= fewer points). To save the reconstruction with full resolution depth maps use the `--reconstruction_path` flag.
 
 
 ```Python

demo.py

@@ -56,6 +56,27 @@ def image_stream(imagedir, calib, stride):
         yield t, image[None], intrinsics
 
 
+def save_reconstruction(droid, reconstruction_path):
+
+    from pathlib import Path
+    import random
+    import string
+
+    t = droid.video.counter.value
+    tstamps = droid.video.tstamp[:t].cpu().numpy()
+    images = droid.video.images[:t].cpu().numpy()
+    disps = droid.video.disps_up[:t].cpu().numpy()
+    poses = droid.video.poses[:t].cpu().numpy()
+    intrinsics = droid.video.intrinsics[:t].cpu().numpy()
+
+    Path("reconstructions/{}".format(reconstruction_path)).mkdir(parents=True, exist_ok=True)
+    np.save("reconstructions/{}/tstamps.npy".format(reconstruction_path), tstamps)
+    np.save("reconstructions/{}/images.npy".format(reconstruction_path), images)
+    np.save("reconstructions/{}/disps.npy".format(reconstruction_path), disps)
+    np.save("reconstructions/{}/poses.npy".format(reconstruction_path), poses)
+    np.save("reconstructions/{}/intrinsics.npy".format(reconstruction_path), intrinsics)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("--imagedir", type=str, help="path to image directory")
@@ -80,13 +101,19 @@ def image_stream(imagedir, calib, stride):
     parser.add_argument("--backend_thresh", type=float, default=22.0)
     parser.add_argument("--backend_radius", type=int, default=2)
     parser.add_argument("--backend_nms", type=int, default=3)
+    parser.add_argument("--upsample", action="store_true")
+    parser.add_argument("--reconstruction_path", help="path to saved reconstruction")
     args = parser.parse_args()
 
     args.stereo = False
     torch.multiprocessing.set_start_method('spawn')
 
     droid = None
 
+    # need high resolution depths
+    if args.reconstruction_path is not None:
+        args.upsample = True
+
     tstamps = []
     for (t, image, intrinsics) in tqdm(image_stream(args.imagedir, args.calib, args.stride)):
         if t < args.t0:
@@ -101,4 +128,7 @@ def image_stream(imagedir, calib, stride):
 
         droid.track(t, image, intrinsics=intrinsics)
 
+    if args.reconstruction_path is not None:
+        save_reconstruction(droid, args.reconstruction_path)
+
     traj_est = droid.terminate(image_stream(args.imagedir, args.calib, args.stride))

droid_slam/droid_backend.py

@@ -15,6 +15,7 @@ def __init__(self, net, video, args):
         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
@@ -28,7 +29,7 @@ def __call__(self, steps=12):
         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)
+        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, 

droid_slam/droid_frontend.py

@@ -10,7 +10,7 @@ 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)
+        self.graph = FactorGraph(video, net.update, max_factors=48, upsample=args.upsample)
 
         # local optimization window
         self.t0 = 0

droid_slam/factor_graph.py

@@ -9,12 +9,13 @@
 
 
 class FactorGraph:
-    def __init__(self, video, update_op, device="cuda:0", corr_impl="volume", max_factors=-1):
+    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
@@ -239,6 +240,9 @@ def update(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, motion_o
             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
 
 
@@ -270,9 +274,11 @@ def update_lowmem(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, s
 
                 with torch.cuda.amp.autocast(enabled=True):
 
-                    net, delta, weight, damping, _ = \
+                    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()

environment.yaml

@@ -1,4 +1,4 @@
-name: droidenv5
+name: droidenv
 channels:
   - rusty1s
   - pytorch