Spaces:
Running
Running
File size: 4,064 Bytes
b7eedf7 |
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 |
import sys
sys.path.append('droid_slam')
sys.path.append('thirdparty/tartanair_tools')
from tqdm import tqdm
import numpy as np
import torch
import lietorch
import cv2
import os
import glob
import time
import yaml
import argparse
from droid import Droid
def image_stream(datapath, image_size=[384, 512], intrinsics_vec=[320.0, 320.0, 320.0, 240.0], stereo=False):
""" image generator """
# read all png images in folder
ht0, wd0 = [480, 640]
images_left = sorted(glob.glob(os.path.join(datapath, 'image_left/*.png')))
images_right = sorted(glob.glob(os.path.join(datapath, 'image_right/*.png')))
data = []
for t in range(len(images_left)):
images = [ cv2.resize(cv2.imread(images_left[t]), (image_size[1], image_size[0])) ]
if stereo:
images += [ cv2.resize(cv2.imread(images_right[t]), (image_size[1], image_size[0])) ]
images = torch.from_numpy(np.stack(images, 0)).permute(0,3,1,2)
intrinsics = .8 * torch.as_tensor(intrinsics_vec)
data.append((t, images, intrinsics))
return data
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--datapath", default="datasets/TartanAir")
parser.add_argument("--weights", default="droid.pth")
parser.add_argument("--buffer", type=int, default=1000)
parser.add_argument("--image_size", default=[384,512])
parser.add_argument("--stereo", action="store_true")
parser.add_argument("--disable_vis", action="store_true")
parser.add_argument("--plot_curve", action="store_true")
parser.add_argument("--id", type=int, default=-1)
parser.add_argument("--beta", type=float, default=0.3)
parser.add_argument("--filter_thresh", type=float, default=2.4)
parser.add_argument("--warmup", type=int, default=12)
parser.add_argument("--keyframe_thresh", type=float, default=3.5)
parser.add_argument("--frontend_thresh", type=float, default=15)
parser.add_argument("--frontend_window", type=int, default=20)
parser.add_argument("--frontend_radius", type=int, default=1)
parser.add_argument("--frontend_nms", type=int, default=1)
parser.add_argument("--backend_thresh", type=float, default=20.0)
parser.add_argument("--backend_radius", type=int, default=2)
parser.add_argument("--backend_nms", type=int, default=3)
args = parser.parse_args()
torch.multiprocessing.set_start_method('spawn')
from data_readers.tartan import test_split
from evaluation.tartanair_evaluator import TartanAirEvaluator
if not os.path.isdir("figures"):
os.mkdir("figures")
if args.id >= 0:
test_split = [ test_split[args.id] ]
ate_list = []
for scene in test_split:
print("Performing evaluation on {}".format(scene))
torch.cuda.empty_cache()
droid = Droid(args)
scenedir = os.path.join(args.datapath, scene)
for (tstamp, image, intrinsics) in tqdm(image_stream(scenedir, stereo=args.stereo)):
droid.track(tstamp, image, intrinsics=intrinsics)
# fill in non-keyframe poses + global BA
traj_est = droid.terminate(image_stream(scenedir))
### do evaluation ###
evaluator = TartanAirEvaluator()
gt_file = os.path.join(scenedir, "pose_left.txt")
traj_ref = np.loadtxt(gt_file, delimiter=' ')[:, [1, 2, 0, 4, 5, 3, 6]] # ned -> xyz
# usually stereo should not be scale corrected, but we are comparing monocular and stereo here
results = evaluator.evaluate_one_trajectory(
traj_ref, traj_est, scale=True, title=scenedir[-20:].replace('/', '_'))
print(results)
ate_list.append(results["ate_score"])
print("Results")
print(ate_list)
if args.plot_curve:
import matplotlib.pyplot as plt
ate = np.array(ate_list)
xs = np.linspace(0.0, 1.0, 512)
ys = [np.count_nonzero(ate < t) / ate.shape[0] for t in xs]
plt.plot(xs, ys)
plt.xlabel("ATE [m]")
plt.ylabel("% runs")
plt.show()
|