415 lines
13 KiB
Python
415 lines
13 KiB
Python
import os
|
|
|
|
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
|
SRC_DIR = os.path.join(ROOT_DIR, "third-party/DROID-SLAM")
|
|
print("PROJ SRC", ROOT_DIR)
|
|
print("DROID SRC", SRC_DIR)
|
|
|
|
import sys
|
|
|
|
sys.path.append(SRC_DIR)
|
|
sys.path.append(f"{SRC_DIR}/droid_slam")
|
|
|
|
import glob
|
|
import shutil
|
|
import argparse
|
|
from tqdm import tqdm
|
|
|
|
import json
|
|
|
|
import cv2
|
|
import trimesh
|
|
import numpy as np
|
|
import torch
|
|
|
|
from lietorch import SE3
|
|
from droid import Droid
|
|
import droid_backends
|
|
|
|
|
|
def get_image(image):
|
|
return image.permute(1, 2, 0).cpu().numpy()[..., ::-1]
|
|
|
|
|
|
def show_image(image):
|
|
image = image.permute(1, 2, 0).cpu().numpy()
|
|
cv2.imshow("image", image / 255.0)
|
|
cv2.waitKey(1)
|
|
|
|
|
|
def get_hwf(test_path):
|
|
test_img = cv2.imread(test_path)
|
|
H, W, _ = test_img.shape
|
|
F = 0.5 * (H + W)
|
|
return H, W, F
|
|
|
|
|
|
def isimage(path):
|
|
ext = os.path.splitext(path)[-1].lower()
|
|
return ext == ".png" or ext == ".jpg" or ext == ".jpeg"
|
|
|
|
|
|
def get_image_files(img_dir, stride, start=0, end=-1):
|
|
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))
|
|
end = len(image_list) + 1 + end if end < 0 else end
|
|
return [os.path.join(img_dir, name) for name in image_list[start:end:stride]]
|
|
|
|
|
|
def load_intrins(intrins_path, image_files, start=0, end=-1):
|
|
N = len(image_files)
|
|
H, W, F = get_hwf(image_files[0])
|
|
# intrinsics for default FOV
|
|
default = torch.tensor([F, F, W / 2, H / 2, W, H])[None].repeat(N, 1)
|
|
if intrins_path is None:
|
|
return default
|
|
assert os.path.isfile(intrins_path), f"{intrins_path} DOES NOT EXIST"
|
|
assert intrins_path.endswith(".txt"), f"{intrins_path} INCORRECT EXT"
|
|
|
|
print(f"LOADING INTRINSICS FROM {intrins_path}")
|
|
intrins = torch.from_numpy(np.loadtxt(intrins_path).astype(np.float32))
|
|
print(intrins.shape, N)
|
|
if intrins.ndim < 2: # assume intrins is same for all frames
|
|
intrins = intrins[None].repeat(N, 1)
|
|
else:
|
|
end = N + 1 + end if end < 0 else end
|
|
intrins = intrins[start:end]
|
|
print(intrins.shape)
|
|
if intrins.shape[1] == 4: # fx, fy, cx, cy, adding on W and H
|
|
ones = torch.ones_like(intrins[:, :1])
|
|
intrins = torch.cat([intrins, ones * W, ones * H], dim=-1)
|
|
return intrins
|
|
|
|
|
|
def image_stream(image_files, intrins_all):
|
|
"""image generator"""
|
|
|
|
N = len(image_files)
|
|
|
|
assert intrins_all.shape[0] == N
|
|
assert intrins_all.shape[1] >= 4
|
|
|
|
for t, imfile in enumerate(image_files):
|
|
image = cv2.imread(imfile)
|
|
if image is None:
|
|
print(imfile, "is none, exiting")
|
|
sys.exit(1)
|
|
|
|
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)
|
|
|
|
fx, fy, cx, cy = intrins_all[t, :4]
|
|
sx, sy = w1 / w0, h1 / h0
|
|
frame_intrins = torch.as_tensor([fx * sx, fy * sy, cx * sx, cy * sy])
|
|
yield t, image[None], frame_intrins
|
|
|
|
|
|
def unpack_video(video, depth_filter_thresh=0.005):
|
|
"""
|
|
extract the video of keyframe disps and poses
|
|
returns:
|
|
(T) timestamps
|
|
(T, 7) world2cam pose for each keyframe
|
|
(T, H, W) disps
|
|
(T, H, W, 3) images
|
|
(T, H, W) mask of valid pixels
|
|
(4) intrinsics (fx, fy, cx, cy)
|
|
"""
|
|
with torch.no_grad():
|
|
with video.get_lock():
|
|
# dirty frames have disparity normalized over all previous frames
|
|
# should probably change this in visualizer
|
|
t = video.counter.value
|
|
|
|
print(f"{t} keyframes in map")
|
|
|
|
# normalization happens after all frames have been processed
|
|
if t < 1:
|
|
print("no keyframes")
|
|
return
|
|
|
|
# taken from visualization.py#80
|
|
tstamps = video.tstamp[:t]
|
|
dirty_index = torch.arange(t, device=video.poses.device)
|
|
|
|
poses = torch.index_select(video.poses, 0, dirty_index) # (T, 7)
|
|
disps = torch.index_select(video.disps, 0, dirty_index) # (T, H, W)
|
|
images = torch.index_select(video.images, 0, dirty_index) # (T, 3, H, W)
|
|
|
|
images = images[:, [2, 1, 0], 3::8, 3::8].permute(0, 2, 3, 1) / 255.0
|
|
intrins = video.intrinsics[0] # (4,)
|
|
|
|
# get mask of static points
|
|
thresh = depth_filter_thresh * torch.ones_like(disps.mean(dim=[1, 2]))
|
|
count = droid_backends.depth_filter(
|
|
video.poses, video.disps, intrins, dirty_index, thresh
|
|
)
|
|
disp_mean = disps.mean(dim=(1, 2), keepdim=True).clamp_min(0.2)
|
|
valid = (count >= 2) & (disps > 0.5 * disp_mean) # (T, H, W)
|
|
return {
|
|
"tstamps": tstamps.cpu(),
|
|
"w2c": poses.cpu(),
|
|
"disps": disps.cpu(),
|
|
"images": images.cpu(),
|
|
"valid": valid.cpu(),
|
|
"intrins": intrins.cpu(),
|
|
}
|
|
|
|
|
|
def get_keyframe_map(video_dict):
|
|
"""
|
|
:param video_dict (dict) unpacked droid slam video
|
|
:returns
|
|
(T, 4, 4) camera to world matrix
|
|
(N, 3) valid points in world
|
|
(N, 3) colors of valid points
|
|
"""
|
|
poses = video_dict["w2c"]
|
|
images = video_dict["images"]
|
|
valid = video_dict["valid"]
|
|
t = len(poses)
|
|
c2w = SE3(poses).inv()
|
|
|
|
# `iproj` only has a cuda backend; if we pass in CPU tensors, it will silently
|
|
# return zeros
|
|
points = droid_backends.iproj(
|
|
c2w.data.cuda(),
|
|
video_dict["disps"].cuda(),
|
|
video_dict["intrins"].cuda(),
|
|
)
|
|
valid_pts = torch.cat([points[i, valid[i]] for i in range(t)], dim=0)
|
|
valid_rgb = torch.cat([images[i, valid[i]] for i in range(t)], dim=0)
|
|
return (
|
|
c2w.matrix().detach().cpu(),
|
|
valid_pts.detach().cpu(),
|
|
valid_rgb.detach().cpu(),
|
|
)
|
|
|
|
|
|
def get_frame_cameras(droid, img_paths, intrins_all):
|
|
N = len(img_paths)
|
|
|
|
t = droid.video.counter.value
|
|
print(f"{t} keyframes in map")
|
|
|
|
if t > 1:
|
|
with torch.no_grad():
|
|
# localize all frames and get edges into keyframe graph
|
|
# returns 7D tensor (3D trans, 4D quat)
|
|
c2w = droid.terminate(image_stream(img_paths, intrins_all))
|
|
c2w = torch.from_numpy(c2w.astype(np.float32))
|
|
return SE3(c2w).inv().matrix()
|
|
|
|
return torch.eye(4)[None].repeat(N, 1, 1)
|
|
|
|
|
|
def save_camera_json(path, extrins, intrins, yup=False):
|
|
"""
|
|
:param path
|
|
:param extrins (N, 4, 4)
|
|
:param intrins (N, 4)
|
|
"""
|
|
N = extrins.shape[0]
|
|
if yup:
|
|
T = torch.tensor(
|
|
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
|
dtype=torch.float32,
|
|
)
|
|
extrins = torch.matmul(T[None], extrins)
|
|
with open(path, "w") as f:
|
|
json.dump(
|
|
{
|
|
"rotation": extrins[:, :3, :3].reshape((N, 9)).tolist(),
|
|
"translation": extrins[:, :3, 3].tolist(),
|
|
"intrinsics": intrins.tolist(),
|
|
},
|
|
f,
|
|
indent=1,
|
|
)
|
|
|
|
|
|
def save_pcl(path, points, colors, yup=False):
|
|
assert len(points) == len(colors)
|
|
if yup:
|
|
T = torch.tensor(
|
|
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
|
dtype=torch.float32,
|
|
)
|
|
points = torch.einsum("ij,...j->...i", T, points) # (*, 3)
|
|
points = points.reshape(-1, 3).numpy()
|
|
colors = colors.reshape(-1, 3).numpy()
|
|
obj = trimesh.Trimesh(vertices=points, vertex_colors=colors)
|
|
with open(path, "wb") as f:
|
|
f.write(trimesh.exchange.ply.export_ply(obj))
|
|
|
|
|
|
def save_keyframe_map(out_dir, kf_nodes):
|
|
"""
|
|
Saves everything with x right y up z back
|
|
(converts from x right y down z forward)
|
|
:param out_dir
|
|
:param kf_nodes (dict)
|
|
"""
|
|
c2w, points, colors = get_keyframe_map(kf_nodes)
|
|
intrins = kf_nodes["intrins"].detach().cpu()
|
|
|
|
print("c2w", c2w.shape)
|
|
print("points", points.shape)
|
|
print("colors", colors.shape)
|
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
save_pcl(f"{out_dir}/map_points.ply", points, colors)
|
|
save_camera_json(f"{out_dir}/map_cameras.json", c2w, intrins)
|
|
print(f"saved {len(c2w)} keyframes to {out_dir}")
|
|
|
|
|
|
def save_cameras(map_dir, frame_w2c, intrins):
|
|
os.makedirs(map_dir, exist_ok=True)
|
|
print("SAVING CAMERA INTRINS AND EXTRINS IN", map_dir)
|
|
|
|
# save all camera extrinsics and intrinsic
|
|
W, H = intrins[0, 4], intrins[0, 5]
|
|
focal = intrins[:, :2].mean()
|
|
print("INTRINS", intrins[0])
|
|
np.savez(
|
|
f"{map_dir}/cameras.npz",
|
|
height=H,
|
|
width=W,
|
|
focal=focal,
|
|
intrins=intrins[:, :4],
|
|
w2c=frame_w2c,
|
|
)
|
|
|
|
# save visualization
|
|
frame_intrins = torch.tensor([focal, focal, W / 2, H / 2, W, H])
|
|
frame_c2w = torch.linalg.inv(frame_w2c)
|
|
save_camera_json(f"{map_dir}/frame_cameras.json", frame_c2w, intrins)
|
|
|
|
|
|
def main(args):
|
|
args.stereo = False
|
|
torch.multiprocessing.set_start_method("spawn")
|
|
|
|
droid = None
|
|
|
|
img_paths = get_image_files(
|
|
args.img_dir, args.stride, start=args.start, end=args.end
|
|
)
|
|
if args.map_dir is not None:
|
|
os.makedirs(args.map_dir, exist_ok=True)
|
|
camera_file = os.path.join(args.map_dir, "cameras.npz")
|
|
if os.path.isfile(camera_file) and not args.overwrite:
|
|
print(glob.glob(f"{args.map_dir}/*.npz"), "already exist, skipping")
|
|
return
|
|
|
|
intrins_all = load_intrins(
|
|
args.intrins_path, img_paths, start=args.start, end=args.end
|
|
)
|
|
print("intrins shape", intrins_all.shape, "num images", len(img_paths))
|
|
|
|
for t, image, intrinsics in tqdm(image_stream(img_paths, intrins_all)):
|
|
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.map_dir is None:
|
|
return
|
|
|
|
# save cameras
|
|
frame_w2c = get_frame_cameras(droid, img_paths, intrins_all)
|
|
save_cameras(args.map_dir, frame_w2c, intrins_all)
|
|
|
|
# save keyframe cameras and points
|
|
kf_nodes = unpack_video(droid.video)
|
|
save_keyframe_map(args.map_dir, kf_nodes)
|
|
|
|
|
|
def get_slam_parser():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--stride", default=1, type=int, help="frame stride")
|
|
parser.add_argument("--t0", default=0, type=int, help="starting frame")
|
|
parser.add_argument("--weights", default=f"{ROOT_DIR}/_DATA/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=1.0,
|
|
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")
|
|
return parser
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = get_slam_parser()
|
|
parser.add_argument(
|
|
"-i", "--img_dir", required=True, type=str, help="path to image directory"
|
|
)
|
|
parser.add_argument(
|
|
"-o", "--map_dir", type=str, default=None, help="path to output directory"
|
|
)
|
|
parser.add_argument(
|
|
"-y", "--overwrite", action="store_true", help="overwrite existing cameras"
|
|
)
|
|
parser.add_argument(
|
|
"-g", "--save_graph", action="store_true", help="save the frame graph"
|
|
)
|
|
parser.add_argument(
|
|
"--intrins_path", default=None, help="path to camera intrinsics"
|
|
)
|
|
parser.add_argument("--start", default=0, type=int, help="start frame")
|
|
parser.add_argument("--end", default=-1, type=int, help="end frame")
|
|
|
|
args = parser.parse_args()
|
|
main(args)
|