194 lines
6.4 KiB
Python
194 lines
6.4 KiB
Python
import numpy as np
|
|
import torch
|
|
import lietorch
|
|
import droid_backends
|
|
|
|
from torch.multiprocessing import Process, Queue, Lock, Value
|
|
from collections import OrderedDict
|
|
|
|
from droid_net import cvx_upsample
|
|
import geom.projective_ops as pops
|
|
|
|
class DepthVideo:
|
|
def __init__(self, image_size=[480, 640], buffer=1024, stereo=False, device="cuda:0"):
|
|
|
|
# current keyframe count
|
|
self.counter = Value('i', 0)
|
|
self.ready = Value('i', 0)
|
|
self.ht = ht = image_size[0]
|
|
self.wd = wd = image_size[1]
|
|
|
|
### state attributes ###
|
|
self.tstamp = torch.zeros(buffer, device="cuda", dtype=torch.float).share_memory_()
|
|
self.images = torch.zeros(buffer, 3, ht, wd, device="cuda", dtype=torch.uint8)
|
|
self.dirty = torch.zeros(buffer, device="cuda", dtype=torch.bool).share_memory_()
|
|
self.red = torch.zeros(buffer, device="cuda", dtype=torch.bool).share_memory_()
|
|
self.poses = torch.zeros(buffer, 7, device="cuda", dtype=torch.float).share_memory_()
|
|
self.disps = torch.ones(buffer, ht//8, wd//8, device="cuda", dtype=torch.float).share_memory_()
|
|
self.disps_sens = torch.zeros(buffer, ht//8, wd//8, device="cuda", dtype=torch.float).share_memory_()
|
|
self.disps_up = torch.zeros(buffer, ht, wd, device="cuda", dtype=torch.float).share_memory_()
|
|
self.intrinsics = torch.zeros(buffer, 4, device="cuda", dtype=torch.float).share_memory_()
|
|
|
|
self.stereo = stereo
|
|
c = 1 if not self.stereo else 2
|
|
|
|
### feature attributes ###
|
|
self.fmaps = torch.zeros(buffer, c, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
|
|
self.nets = torch.zeros(buffer, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
|
|
self.inps = torch.zeros(buffer, 128, ht//8, wd//8, dtype=torch.half, device="cuda").share_memory_()
|
|
|
|
# initialize poses to identity transformation
|
|
self.poses[:] = torch.as_tensor([0, 0, 0, 0, 0, 0, 1], dtype=torch.float, device="cuda")
|
|
|
|
def get_lock(self):
|
|
return self.counter.get_lock()
|
|
|
|
def __item_setter(self, index, item):
|
|
if isinstance(index, int) and index >= self.counter.value:
|
|
self.counter.value = index + 1
|
|
|
|
elif isinstance(index, torch.Tensor) and index.max().item() > self.counter.value:
|
|
self.counter.value = index.max().item() + 1
|
|
|
|
# self.dirty[index] = True
|
|
self.tstamp[index] = item[0]
|
|
self.images[index] = item[1]
|
|
|
|
if item[2] is not None:
|
|
self.poses[index] = item[2]
|
|
|
|
if item[3] is not None:
|
|
self.disps[index] = item[3]
|
|
|
|
if item[4] is not None:
|
|
depth = item[4][3::8,3::8]
|
|
self.disps_sens[index] = torch.where(depth>0, 1.0/depth, depth)
|
|
|
|
if item[5] is not None:
|
|
self.intrinsics[index] = item[5]
|
|
|
|
if len(item) > 6:
|
|
self.fmaps[index] = item[6]
|
|
|
|
if len(item) > 7:
|
|
self.nets[index] = item[7]
|
|
|
|
if len(item) > 8:
|
|
self.inps[index] = item[8]
|
|
|
|
def __setitem__(self, index, item):
|
|
with self.get_lock():
|
|
self.__item_setter(index, item)
|
|
|
|
def __getitem__(self, index):
|
|
""" index the depth video """
|
|
|
|
with self.get_lock():
|
|
# support negative indexing
|
|
if isinstance(index, int) and index < 0:
|
|
index = self.counter.value + index
|
|
|
|
item = (
|
|
self.poses[index],
|
|
self.disps[index],
|
|
self.intrinsics[index],
|
|
self.fmaps[index],
|
|
self.nets[index],
|
|
self.inps[index])
|
|
|
|
return item
|
|
|
|
def append(self, *item):
|
|
with self.get_lock():
|
|
self.__item_setter(self.counter.value, item)
|
|
|
|
|
|
### geometric operations ###
|
|
|
|
@staticmethod
|
|
def format_indicies(ii, jj):
|
|
""" to device, long, {-1} """
|
|
|
|
if not isinstance(ii, torch.Tensor):
|
|
ii = torch.as_tensor(ii)
|
|
|
|
if not isinstance(jj, torch.Tensor):
|
|
jj = torch.as_tensor(jj)
|
|
|
|
ii = ii.to(device="cuda", dtype=torch.long).reshape(-1)
|
|
jj = jj.to(device="cuda", dtype=torch.long).reshape(-1)
|
|
|
|
return ii, jj
|
|
|
|
def upsample(self, ix, mask):
|
|
""" upsample disparity """
|
|
|
|
disps_up = cvx_upsample(self.disps[ix].unsqueeze(-1), mask)
|
|
self.disps_up[ix] = disps_up.squeeze()
|
|
|
|
def normalize(self):
|
|
""" normalize depth and poses """
|
|
|
|
with self.get_lock():
|
|
s = self.disps[:self.counter.value].mean()
|
|
self.disps[:self.counter.value] /= s
|
|
self.poses[:self.counter.value,:3] *= s
|
|
self.dirty[:self.counter.value] = True
|
|
|
|
|
|
def reproject(self, ii, jj):
|
|
""" project points from ii -> jj """
|
|
ii, jj = DepthVideo.format_indicies(ii, jj)
|
|
Gs = lietorch.SE3(self.poses[None])
|
|
|
|
coords, valid_mask = \
|
|
pops.projective_transform(Gs, self.disps[None], self.intrinsics[None], ii, jj)
|
|
|
|
return coords, valid_mask
|
|
|
|
def distance(self, ii=None, jj=None, beta=0.3, bidirectional=True):
|
|
""" frame distance metric """
|
|
|
|
return_matrix = False
|
|
if ii is None:
|
|
return_matrix = True
|
|
N = self.counter.value
|
|
ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
|
|
|
|
ii, jj = DepthVideo.format_indicies(ii, jj)
|
|
|
|
if bidirectional:
|
|
|
|
poses = self.poses[:self.counter.value].clone()
|
|
|
|
d1 = droid_backends.frame_distance(
|
|
poses, self.disps, self.intrinsics[0], ii, jj, beta)
|
|
|
|
d2 = droid_backends.frame_distance(
|
|
poses, self.disps, self.intrinsics[0], jj, ii, beta)
|
|
|
|
d = .5 * (d1 + d2)
|
|
|
|
else:
|
|
d = droid_backends.frame_distance(
|
|
self.poses, self.disps, self.intrinsics[0], ii, jj, beta)
|
|
|
|
if return_matrix:
|
|
return d.reshape(N, N)
|
|
|
|
return d
|
|
|
|
def ba(self, target, weight, eta, ii, jj, t0=1, t1=None, itrs=2, lm=1e-4, ep=0.1, motion_only=False):
|
|
""" dense bundle adjustment (DBA) """
|
|
|
|
with self.get_lock():
|
|
|
|
# [t0, t1] window of bundle adjustment optimization
|
|
if t1 is None:
|
|
t1 = max(ii.max().item(), jj.max().item()) + 1
|
|
|
|
droid_backends.ba(self.poses, self.disps, self.intrinsics[0], self.disps_sens,
|
|
target, weight, eta, ii, jj, t0, t1, itrs, lm, ep, motion_only)
|
|
|
|
self.disps.clamp_(min=0.001)
|