263 lines
8.6 KiB
Python
263 lines
8.6 KiB
Python
import torch
|
|
import math
|
|
import bisect
|
|
|
|
|
|
## Extracted from https://github.com/TimoStoff/event_utils
|
|
def events_to_voxel_torch(
|
|
xs, ys, ts, ps, B, device=None, sensor_size=(180, 240), temporal_bilinear=True
|
|
):
|
|
"""
|
|
Turn set of events to a voxel grid tensor, using temporal bilinear interpolation
|
|
@param xs List of event x coordinates (torch tensor)
|
|
@param ys List of event y coordinates (torch tensor)
|
|
@param ts List of event timestamps (torch tensor)
|
|
@param ps List of event polarities (torch tensor)
|
|
@param B Number of bins in output voxel grids (int)
|
|
@param device Device to put voxel grid. If left empty, same device as events
|
|
@param sensor_size The size of the event sensor/output voxels
|
|
@param temporal_bilinear Whether the events should be naively
|
|
accumulated to the voxels (faster), or properly
|
|
temporally distributed
|
|
@returns Voxel of the events between t0 and t1
|
|
"""
|
|
if device is None:
|
|
device = xs.device
|
|
assert len(xs) == len(ys) and len(ys) == len(ts) and len(ts) == len(ps)
|
|
bins = []
|
|
dt = ts[-1] - ts[0]
|
|
t_norm = (ts - ts[0]) / dt * (B - 1)
|
|
zeros = torch.zeros_like(t_norm)
|
|
for bi in range(B):
|
|
assert temporal_bilinear, "no other option not supported"
|
|
bilinear_weights = torch.max(zeros, 1.0 - torch.abs(t_norm - bi))
|
|
weights = ps * bilinear_weights
|
|
vb = events_to_image_torch(
|
|
xs,
|
|
ys,
|
|
weights,
|
|
device,
|
|
sensor_size=sensor_size,
|
|
clip_out_of_range=False,
|
|
)
|
|
bins.append(vb)
|
|
bins = torch.stack(bins)
|
|
return bins
|
|
|
|
|
|
## Extracted from https://github.com/TimoStoff/event_utils
|
|
def events_to_image_torch(
|
|
xs,
|
|
ys,
|
|
ps,
|
|
device=None,
|
|
sensor_size=(180, 240),
|
|
clip_out_of_range=True,
|
|
interpolation=None,
|
|
padding=True,
|
|
default=0,
|
|
):
|
|
"""
|
|
Method to turn event tensor to image. Allows for bilinear interpolation.
|
|
@param xs Tensor of x coords of events
|
|
@param ys Tensor of y coords of events
|
|
@param ps Tensor of event polarities/weights
|
|
@param device The device on which the image is. If none, set to events device
|
|
@param sensor_size The size of the image sensor/output image
|
|
@param clip_out_of_range If the events go beyond the desired image size,
|
|
clip the events to fit into the image
|
|
@param interpolation Which interpolation to use. Options=None,'bilinear'
|
|
@param padding If bilinear interpolation, allow padding the image by 1 to allow events to fit:
|
|
@returns Event image from the events
|
|
"""
|
|
if device is None:
|
|
device = xs.device
|
|
if interpolation == "bilinear" and padding:
|
|
img_size = (sensor_size[0] + 1, sensor_size[1] + 1)
|
|
else:
|
|
img_size = list(sensor_size)
|
|
|
|
mask = torch.ones(xs.size(), device=device)
|
|
if clip_out_of_range:
|
|
zero_v = torch.tensor([0.0], device=device)
|
|
ones_v = torch.tensor([1.0], device=device)
|
|
clipx = (
|
|
img_size[1]
|
|
if interpolation is None and padding == False
|
|
else img_size[1] - 1
|
|
)
|
|
clipy = (
|
|
img_size[0]
|
|
if interpolation is None and padding == False
|
|
else img_size[0] - 1
|
|
)
|
|
mask = torch.where(xs >= clipx, zero_v, ones_v) * torch.where(
|
|
ys >= clipy, zero_v, ones_v
|
|
)
|
|
|
|
img = (torch.ones(img_size) * default).to(device)
|
|
if (
|
|
interpolation == "bilinear"
|
|
and xs.dtype is not torch.long
|
|
and xs.dtype is not torch.long
|
|
):
|
|
pxs = (xs.floor()).float()
|
|
pys = (ys.floor()).float()
|
|
dxs = (xs - pxs).float()
|
|
dys = (ys - pys).float()
|
|
pxs = (pxs * mask).long()
|
|
pys = (pys * mask).long()
|
|
masked_ps = ps.squeeze() * mask
|
|
interpolate_to_image(pxs, pys, dxs, dys, masked_ps, img)
|
|
else:
|
|
if xs.dtype is not torch.long:
|
|
xs = xs.long().to(device)
|
|
if ys.dtype is not torch.long:
|
|
ys = ys.long().to(device)
|
|
try:
|
|
mask = mask.long().to(device)
|
|
xs, ys = xs * mask, ys * mask
|
|
img.index_put_((ys, xs), ps, accumulate=True)
|
|
except Exception as e:
|
|
print(
|
|
"Unable to put tensor {} positions ({}, {}) into {}. Range = {},{}".format(
|
|
ps.shape,
|
|
ys.shape,
|
|
xs.shape,
|
|
img.shape,
|
|
torch.max(ys),
|
|
torch.max(xs),
|
|
)
|
|
)
|
|
raise e
|
|
return img
|
|
|
|
|
|
## Extracted from https://github.com/TimoStoff/event_utils
|
|
def interpolate_to_image(pxs, pys, dxs, dys, weights, img):
|
|
"""
|
|
Accumulate x and y coords to an image using bilinear interpolation
|
|
@param pxs Numpy array of integer typecast x coords of events
|
|
@param pys Numpy array of integer typecast y coords of events
|
|
@param dxs Numpy array of residual difference between x coord and int(x coord)
|
|
@param dys Numpy array of residual difference between y coord and int(y coord)
|
|
@returns Image
|
|
"""
|
|
img.index_put_((pys, pxs), weights * (1.0 - dxs) * (1.0 - dys), accumulate=True)
|
|
img.index_put_((pys, pxs + 1), weights * dxs * (1.0 - dys), accumulate=True)
|
|
img.index_put_((pys + 1, pxs), weights * (1.0 - dxs) * dys, accumulate=True)
|
|
img.index_put_((pys + 1, pxs + 1), weights * dxs * dys, accumulate=True)
|
|
return img
|
|
|
|
|
|
def voxel_normalization(voxel):
|
|
"""
|
|
normalize the voxel same as https://arxiv.org/abs/1912.01584 Section 3.1
|
|
Params:
|
|
voxel: torch.Tensor, shape is [num_bins, H, W]
|
|
|
|
return:
|
|
normalized voxel
|
|
"""
|
|
# check if voxel all element is 0
|
|
tmp = torch.zeros_like(voxel)
|
|
if torch.equal(voxel, tmp):
|
|
return voxel
|
|
abs_voxel, _ = torch.sort(torch.abs(voxel).view(-1, 1).squeeze(1))
|
|
first_non_zero_idx = torch.nonzero(abs_voxel)[0].item()
|
|
non_zero_voxel = abs_voxel[first_non_zero_idx:]
|
|
norm_idx = math.floor(non_zero_voxel.shape[0] * 0.98)
|
|
ones = torch.ones_like(voxel)
|
|
normed_voxel = torch.where(
|
|
torch.abs(voxel) < non_zero_voxel[norm_idx],
|
|
voxel / non_zero_voxel[norm_idx],
|
|
voxel,
|
|
)
|
|
normed_voxel = torch.where(
|
|
normed_voxel >= non_zero_voxel[norm_idx], ones, normed_voxel
|
|
)
|
|
normed_voxel = torch.where(
|
|
normed_voxel <= -non_zero_voxel[norm_idx], -ones, normed_voxel
|
|
)
|
|
return normed_voxel
|
|
|
|
|
|
# Taken and modified from https://github.com/DachunKai/EvTexture/issues/12#issuecomment-2198243470
|
|
def package_bidirectional_event_voxels(
|
|
x,
|
|
y,
|
|
t,
|
|
p,
|
|
timestamp_list,
|
|
backward,
|
|
bins,
|
|
sensor_size,
|
|
):
|
|
"""
|
|
params:
|
|
x: ndarray, x-position of events
|
|
y: ndarray, y-position of events
|
|
t: ndarray, timestamp of events
|
|
p: ndarray, polarity of events
|
|
backward: bool, if forward or backward
|
|
timestamp_list: list, to split events via timestamp
|
|
bins: voxel num_bins
|
|
returns:
|
|
no return.
|
|
"""
|
|
assert x.shape == y.shape == t.shape == p.shape
|
|
|
|
# Step 2: select events between two frames according to timestamp
|
|
temp = t.cpu().numpy().tolist()
|
|
output = [
|
|
temp[
|
|
bisect.bisect_left(temp, timestamp_list[i]) : bisect.bisect_left(
|
|
temp, timestamp_list[i + 1]
|
|
)
|
|
]
|
|
for i in range(len(timestamp_list) - 1)
|
|
]
|
|
|
|
# Debug: Check if data error!!!
|
|
assert (
|
|
len(output) == len(timestamp_list) - 1
|
|
), f"len(output) is {len(output)}, but len(timestamp_list) is {len(timestamp_list)}"
|
|
sum_output = []
|
|
sum = 0
|
|
for i in range(len(output)):
|
|
if len(output[i]) <= 1:
|
|
raise ValueError(f"len(output[{i}] == 0)")
|
|
sum += len(output[i])
|
|
sum_output.append(sum)
|
|
|
|
assert len(sum_output) == len(output)
|
|
|
|
# Step 3: After checking data, continue.
|
|
start_idx = 0
|
|
out_voxels = []
|
|
for voxel_idx in range(len(timestamp_list) - 1):
|
|
end_idx = start_idx + len(output[voxel_idx])
|
|
|
|
xs = x[start_idx:end_idx]
|
|
ys = y[start_idx:end_idx]
|
|
ts = t[start_idx:end_idx]
|
|
ps = p[start_idx:end_idx]
|
|
|
|
if backward:
|
|
t_start = timestamp_list[voxel_idx]
|
|
t_end = timestamp_list[voxel_idx + 1]
|
|
xs = torch.flip(xs, dims=[0])
|
|
ys = torch.flip(ys, dims=[0])
|
|
ts = torch.flip(t_end - ts + t_start, dims=[0])
|
|
ps = torch.flip(-ps, dims=[0])
|
|
|
|
voxel = events_to_voxel_torch(
|
|
xs, ys, ts, ps, bins, device=None, sensor_size=sensor_size
|
|
)
|
|
normed_voxel = voxel_normalization(voxel)
|
|
|
|
out_voxels.append(normed_voxel)
|
|
start_idx = end_idx
|
|
|
|
return out_voxels
|