116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
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()
|
|
|