-
Notifications
You must be signed in to change notification settings - Fork 302
/
demo.py
134 lines (101 loc) · 4.84 KB
/
demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import sys
sys.path.append('droid_slam')
from tqdm import tqdm
import numpy as np
import torch
import lietorch
import cv2
import os
import glob
import time
import argparse
from torch.multiprocessing import Process
from droid import Droid
import torch.nn.functional as F
def show_image(image):
image = image.permute(1, 2, 0).cpu().numpy()
cv2.imshow('image', image / 255.0)
cv2.waitKey(1)
def image_stream(imagedir, calib, stride):
""" image generator """
calib = np.loadtxt(calib, delimiter=" ")
fx, fy, cx, cy = calib[:4]
K = np.eye(3)
K[0,0] = fx
K[0,2] = cx
K[1,1] = fy
K[1,2] = cy
image_list = sorted(os.listdir(imagedir))[::stride]
for t, imfile in enumerate(image_list):
image = cv2.imread(os.path.join(imagedir, imfile))
if len(calib) > 4:
image = cv2.undistort(image, K, calib[4:])
h0, w0, _ = image.shape
h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))
image = cv2.resize(image, (w1, h1))
image = image[:h1-h1%8, :w1-w1%8]
image = torch.as_tensor(image).permute(2, 0, 1)
intrinsics = torch.as_tensor([fx, fy, cx, cy])
intrinsics[0::2] *= (w1 / w0)
intrinsics[1::2] *= (h1 / h0)
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")
parser.add_argument("--calib", type=str, help="path to calibration file")
parser.add_argument("--t0", default=0, type=int, help="starting frame")
parser.add_argument("--stride", default=3, type=int, help="frame stride")
parser.add_argument("--weights", default="droid.pth")
parser.add_argument("--buffer", type=int, default=512)
parser.add_argument("--image_size", default=[240, 320])
parser.add_argument("--disable_vis", action="store_true")
parser.add_argument("--beta", type=float, default=0.3, help="weight for translation / rotation components of flow")
parser.add_argument("--filter_thresh", type=float, default=2.4, help="how much motion before considering new keyframe")
parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
parser.add_argument("--keyframe_thresh", type=float, default=4.0, help="threshold to create a new keyframe")
parser.add_argument("--frontend_thresh", type=float, default=16.0, help="add edges between frames whithin this distance")
parser.add_argument("--frontend_window", type=int, default=25, help="frontend optimization window")
parser.add_argument("--frontend_radius", type=int, default=2, help="force edges between frames within radius")
parser.add_argument("--frontend_nms", type=int, default=1, help="non-maximal supression of edges")
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:
continue
if not args.disable_vis:
show_image(image[0])
if droid is None:
args.image_size = [image.shape[2], image.shape[3]]
droid = Droid(args)
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))