File size: 3,755 Bytes
59630ba | 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 | import argparse
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), 'droid_slam'))
from typing import List
import cv2
import numpy as np
import torch
from droid import Droid
from tqdm import tqdm
def image_stream(image_list, stride, calib=None):
""" image generator """
# If calib is None, compute intrinsics from image size
image_list = image_list[::stride]
for t, imfile in enumerate(image_list):
image = cv2.imread(imfile)
h0, w0, _ = image.shape
# If calib is None, set fx, fy to a reasonable value, cx, cy to half image size
if calib is None:
fx = fy = 0.9 * max(h0, w0) # 0.9 * max dimension as a reasonable focal length
cx = w0 / 2.0
cy = h0 / 2.0
else:
fx, fy, cx, cy = calib
K = np.eye(3)
K[0,0] = fx
K[0,2] = cx
K[1,1] = fy
K[1,2] = cy
h1 = int(h0 * np.sqrt((512 * 512) / (h0 * w0)))
w1 = int(w0 * np.sqrt((512 * 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, save_path):
if hasattr(droid, "video2"):
video = droid.video2
else:
video = droid.video
t = video.counter.value
save_data = {
"tstamps": video.tstamp[:t].cpu(),
"images": video.images[:t].cpu(),
"disps": video.disps_up[:t].cpu(),
"poses": video.poses[:t].cpu(),
"intrinsics": video.intrinsics[:t].cpu(),
}
torch.save(save_data, save_path)
class ReprojectionErrorMetric():
"""
return: Reprojection error
"""
def __init__(self) -> None:
super().__init__()
args = {
"t0": 0,
"stride": 1,
"weights": os.path.join(os.path.dirname(__file__), 'checkpoints', 'droid.pth'),
"buffer": 512,
"beta": 0.3,
"filter_thresh": 0.01,
"warmup": 8,
"keyframe_thresh": 4.0,
"frontend_thresh": 16.0,
"frontend_window": 25,
"frontend_radius": 2,
"frontend_nms": 1,
"backend_thresh": 22.0,
"backend_radius": 2,
"backend_nms": 3,
# need high resolution depths
"upsample": True,
"stereo": False,
}
args = argparse.Namespace(**args)
self._args = args
self.droid = None
try:
torch.multiprocessing.set_start_method('spawn')
except Exception as e:
print(f"Error setting start method: {e}")
def compute_scores(self, rendered_images: List[str], reconstruction_path: str = None) -> float:
for (t, image, intrinsics) in tqdm(image_stream(rendered_images, self._args.stride, getattr(self._args, 'calib', None))):
if t < self._args.t0:
continue
if self.droid is None:
self._args.image_size = [image.shape[2], image.shape[3]]
self.droid = Droid(self._args)
self.droid.track(t, image, intrinsics=intrinsics)
traj_est, valid_errors = self.droid.terminate(image_stream(rendered_images, self._args.stride, getattr(self._args, 'calib', None)))
if len(valid_errors) > 0:
mean_error = valid_errors.mean().item()
if reconstruction_path is not None:
save_reconstruction(self.droid, reconstruction_path)
self.droid = None
return mean_error
|