diff --git a/README.md b/README.md index 6d0622bb..dc73a219 100644 --- a/README.md +++ b/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 diff --git a/demo.py b/demo.py index c6ccabfa..382f080a 100644 --- a/demo.py +++ b/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,6 +101,8 @@ 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 @@ -87,6 +110,10 @@ def image_stream(imagedir, calib, stride): 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)) diff --git a/droid_slam/droid_backend.py b/droid_slam/droid_backend.py index c452bdf0..665eee0e 100644 --- a/droid_slam/droid_backend.py +++ b/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, diff --git a/droid_slam/droid_frontend.py b/droid_slam/droid_frontend.py index 9e55de6c..5d2c56d6 100644 --- a/droid_slam/droid_frontend.py +++ b/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 diff --git a/droid_slam/factor_graph.py b/droid_slam/factor_graph.py index 8f288370..eb605d63 100644 --- a/droid_slam/factor_graph.py +++ b/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() diff --git a/environment.yaml b/environment.yaml index 9c78541b..f92153ac 100644 --- a/environment.yaml +++ b/environment.yaml @@ -1,4 +1,4 @@ -name: droidenv5 +name: droidenv channels: - rusty1s - pytorch