def forward(self, locs, idxs, poses, scales):
""" Compute a forward pass of the SDF Convolution Layer.
Inputs:
-locs: A BxNxD tensor where B is the batch size, N is the number
of query locations, and D is the dimensionality of the
coordinate space.
-idxs: A BxM tensor where M is the number of SDfs in the scene. Each
element of idxs is an index of an SDF in the sdfs list passed
to the constructor.
-poses: A BxMx(D+R) tensor with the pose of each of the M tensors in
idxs in each row. Each row is D+R in length, where the first
D elements are the translation and the remaining R elements
are the rotation. R will vary in size and format depending on
D as follows:
-D=1: R is 0 (no rotation in 1D)
-D=2: R is 1 and is the angle in radians, where 0 is +x
and pi/2 is +y.
-D=3: R is 4 and is the quaternion representation of the
3D rotation in xyzw format and normalized.
All other values for D are not supported.
-scales: A BxM tensor with the scale
Returns: A BxNxO tensor where O is the number of output features.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
M = idxs.size()[1]
R = {1: 0, 2: 1, 3: 4}[self.ndim]
ec.check_tensor_dims(locs, "locs", (batch_size, N, self.ndim))
ec.check_tensor_dims(idxs, "idxs", (
batch_size,
M,
))
ec.check_tensor_dims(poses, "poses", (batch_size, M, self.ndim + R))
ec.check_tensor_dims(scales, "scales", (
batch_size,
M,
))
locs = locs.contiguous()
idxs = idxs.contiguous()
poses = poses.contiguous()
scales = scales.contiguous()
# Do the compution.
# convsdf = _ConvSDFFunction(self.sdfs, self.sdf_offsets, self.sdf_shapes,
# self.kernel_size, self.dilation, self.max_distance, self.compute_pose_grads)
# pdb.set_trace()
return _ConvSDFFunction.apply(locs, idxs, poses, scales, self.weight,
self.bias, self.sdfs, self.sdf_offsets,
self.sdf_shapes, self.kernel_size,
self.dilation, self.max_distance,
self.compute_pose_grads)
def forward(self, idxs, locs, data=None):
""" TODO
Inputs:
-locs: A BxNxD tensor where B is the batch size, N is the number
of particles, and D is the dimensionality of the particles'
coordinate space.
-data: [optional] A BxNxC tensor where C is the number of channels.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
ec.check_tensor_dims(locs, "locs", (batch_size, N, -1))
ec.check_tensor_dims(idxs, "idxs", (batch_size, N))
if data is not None:
ec.check_tensor_dims(data, "data", (batch_size, N, -1))
data = data.contiguous()
no_data = False
else:
data = torch.autograd.Variable(locs.data.new(0, 0, 0),
requires_grad=False)
no_data = True
locs = locs.contiguous()
idxs = idxs.contiguous()
# Do the compution.
coll = _ReorderDataFunction(self.reverse)
locs, data = coll(idxs, locs, data)
if no_data:
return locs
else:
return locs, data
def forward(self, locs, data, neighbors, qlocs=None):
""" Compute a forward pass of the Smooth Particle Convolution Layer.
Inputs:
-locs: A BxNxD tensor where B is the batch size, N is the number
of particles, and D is the dimensionality of the particles'
coordinate space.
-data: A BxNxC tensor where C is the number of input features.
-neighbors: A BxMxK tensor where K is the maximum number of neighbors.
For each particle, this is the list of particle indices
that are considered neighbors. If the number of neighbors
is less than K, the less is ended with a -1.
-qlocs: A BxMxD of query locations. Specify this if you want the query
locations to be different from the particle locations.
Returns: A BxMxO tensor where O is the number of output features. Computed at qlocs.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
ec.check_tensor_dims(locs, "locs", (batch_size, N, self.ndim))
ec.check_tensor_dims(data, "data", (batch_size, N, self.nchannels))
locs = locs.contiguous()
data = data.contiguous()
if qlocs is not None:
ec.check_tensor_dims(qlocs, "qlocs", (batch_size, -1, self.ndim))
qlocs = qlocs.contiguous()
if locs.is_cuda:
if self.device_id != torch.cuda.current_device():
self.device_id = torch.cuda.current_device()
self.nshared_device_mem = _extc.spnc_get_shared_mem_size(
self.device_id)
# Do the compution.
# convsp = _ConvSPFunction(self.radius, self.kernel_size, self.dilation, self.dis_norm, self.kernel_fn, self.ncells, self.nshared_device_mem)
# data.shape = BxCxN
# data = convsp(qlocs if qlocs is not None else locs,
# locs, data, neighbors, self.weight, self.bias)
# data.shape = BxOxN
return _ConvSPFunction.apply(qlocs if qlocs is not None else locs,
locs, data, neighbors, self.weight,
self.bias, self.radius, self.kernel_size,
self.dilation, self.dis_norm,
self.kernel_fn, self.ncells,
self.nshared_device_mem)
def forward(self, locs, image, camera_pose, camera_rot, depth_mask=None):
""" Forwad pass for the particle projection. Takes in the set of
particles and outputs an image.
TODO
Arguments:
-locs: A BxNx3 tensor where B is the batch size, N is the number
of particles, and 3 is the dimensionality of the
particles' coordinate space (this layer currently only
supports 3D projections).
-camera_pose: A Bx3 tensor containing the camera translation.
-camera_rot: A Bx4 tensor containing the camera rotation as a
quaternion in xyzw format.
-depth_mask: An optional BxHxW tensor where W and H are the
camera image width and height respectively. If not
None, then this is used to compute occlusions. The
value in each pixel in the depth_mask should be
the distance to the first object. Any particles
further away than that value will not be projected
onto the output image.
Returns: A BxHxW tensor of the projected particles.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
width = image.size()[3]
height = image.size()[2]
channels = image.size()[1]
ec.check_tensor_dims(locs, "locs", (batch_size, N, 3))
ec.check_tensor_dims(image, "image",
(batch_size, channels, height, width))
ec.check_tensor_dims(camera_pose, "camera_pose", (batch_size, 3))
ec.check_tensor_dims(camera_rot, "camera_rot", (batch_size, 4))
ec.check_nans(locs, "locs")
ec.check_nans(image, "image")
ec.check_nans(camera_pose, "camera_pose")
ec.check_nans(camera_rot, "camera_rot")
if depth_mask is not None:
ec.check_tensor_dims(depth_mask, "depth_mask",
(batch_size, height, width))
ec.check_nans(depth_mask, "depth_mask")
depth_mask = depth_mask.contiguous()
else:
if (self.empty_depth_mask.size()[0] != batch_size
or self.empty_depth_mask.size()[1] != height
or self.empty_depth_mask.size()[2] != width):
self.empty_depth_mask.resize_(batch_size, height, width)
self.empty_depth_mask.fill_(MAX_FLOAT)
depth_mask = torch.autograd.Variable(self.empty_depth_mask,
requires_grad=False)
if locs.is_cuda:
depth_mask = depth_mask.cuda()
# Let's transform the particles to camera space here.
locs = locs - camera_pose.unsqueeze(1)
# Ensure the rotation quaternion is normalized.
camera_rot = camera_rot / \
torch.sqrt(torch.sum(camera_rot**2, 1, keepdim=True))
# Invert the rotation.
inv = camera_rot.data.new(1, 4)
inv[0, 0] = -1
inv[0, 1] = -1
inv[0, 2] = -1
inv[0, 3] = 1
inv = torch.autograd.Variable(inv, requires_grad=False)
camera_rot = camera_rot * inv
rot = self._rotationMatrixFromQuaternion(camera_rot)
if (rot != rot).data.any():
raise ValueError(
"No NaNs found in camera_rot argument, but NaNs created when"
" constructing a rotation matrix from it.")
# Rotate the locs into camera space.
try:
# There's a bug that causes this to fail on the first call when using cuda.
# To fix that, just call it again.
locs = torch.bmm(locs, rot)
except RuntimeError:
locs = torch.bmm(locs, rot)
if (locs != locs).data.any():
raise ValueError(
"Rotating locs by rotation matrix resulted in NaNs.")
locs = locs.contiguous()
image = image.contiguous()
proj = _ImageProjectionFunction(self.camera_fl)
ret = proj(locs, image, depth_mask)
return ret
def forward(self, locs, data=None, qlocs=None):
""" Compute the neighbors of each location. Reorders the locs and data tensors
in place and returns the list of indices in their new order and the list of
neighbors for each location.
Inputs:
-locs: A BxNxD tensor where B is the batch size, N is the number
of particles, and D is the dimensionality of the particles'
coordinate space.
-data: [optional] A BxNxC tensor where C is the number of channels.
Add this to have it reordered alongside locs.
-qlocs: [optional] A BxMxD tensor of query locations. The neighbors
list in the output will be a list of all particles in locs
that neighbor each query location. If not provided, locs is
used instead.
Returns:
-locs: A BxNxD tensor identical to the input locs, except reordered
for optimized memory access.
-data: [optional] A BxNxC tensor identical to the input data reordered
in the same order as locs. If the input data was not provided,
then this is not returned.
-Idxs: BxN tensor with the original index of each particle location in their
new order, e.g., idxs[b, i] = j where b is the batch index, j is
the original index in locs, and i is the new index.
-Neighbors: BxMxK where K is max_neighbors. This lists the indices of
all particles within radius of each query location, up to K. If
there are fewer than K neighbors, -1 is used to indicate
the end of the neighbor list. The indices are with respect to
the reordered locs tensor. If qlocs is not specified, then
locs is used as the query points and it is reordered before
being queried, so the neighbors tensor is also reorderd.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
ec.check_tensor_dims(locs, "locs", (batch_size, N, self.ndim))
if data is not None:
ec.check_tensor_dims(data, "data", (batch_size, N, -1))
data = data.contiguous()
has_data = True
else:
has_data = False
if qlocs is not None:
ec.check_tensor_dims(qlocs, "qlocs", (batch_size, -1, self.ndim))
qlocs = qlocs.contiguous()
locs = locs.contiguous()
# Resize the internal buffers to be the right size.
buffers = [self.cellIDs, self.cellStarts, self.cellEnds]
for buf in buffers:
if buf.size()[0] != batch_size:
ns = (batch_size, ) + buf.size()[1:]
buf.resize_(ns)
if self.cellIDs.size()[1] != N or self.cellIDs.size(
)[0] != batch_size + 2:
# Allocate 2 extra batches on cellIDs for sorting.
self.cellIDs.resize_(batch_size + 2, N, 1)
if locs.is_cuda:
if self.radixsort_buffer_size < 0:
self.radixsort_buffer_size = _ext.spnc_get_radixsort_buffer_size(
)
bufsize = max(
self.radixsort_buffer_size,
int(
np.prod(locs.size()) +
(np.prod(data.size()) if has_data else 0)))
if self.cuda_buffer.size()[0] != bufsize:
self.cuda_buffer.resize_(bufsize)
# Compute grid bounds.
lower_bounds, _ = locs.min(1)
upper_bounds, _ = locs.max(1)
grid_dims = torch.ceil(
torch.clamp((upper_bounds - lower_bounds) / self.radius, 0,
self.max_grid_dim))
center = (lower_bounds + upper_bounds) / 2
lower_bounds = center - grid_dims * self.radius / 2
lower_bounds = lower_bounds.contiguous()
grid_dims = grid_dims.contiguous()
# Get the new hashgrid order.
hashorder = _HashgridOrderFunction(self.radius, self.max_grid_dim,
self.cellIDs, self.cuda_buffer)
idxs = hashorder(locs, lower_bounds, grid_dims)
# Reorder the locs and data.
if has_data:
locs, data = self.reorder(idxs, locs, data)
else:
locs = self.reorder(idxs, locs)
# Do the collision compution.
coll = _ParticleCollisionFunction(self.radius, self.max_collisions,
self.cellIDs, self.cellStarts,
self.cellEnds, self.include_self)
neighbors = coll(qlocs if qlocs is not None else locs, locs,
lower_bounds, grid_dims)
if has_data:
return locs, data, idxs, neighbors
else:
return locs, idxs, neighbors
def forward(self, locs, camera_pose, camera_rot, depth_mask=None):
""" Forwad pass for the particle projection. Takes in the set of
particles and outputs an image.
Arguments:
-locs: A BxNx3 tensor where B is the batch size, N is the number
of particles, and 3 is the dimensionality of the
particles' coordinate space (this layer currently only
supports 3D projections).
-camera_pose: A Bx3 tensor containing the camera translation.
-camera_rot: A Bx4 tensor containing the camera rotation as a
quaternion in xyzw format.
-depth_mask: An optional BxHxW tensor where W and H are the
camera image width and height respectively. If not
None, then this is used to compute occlusions. The
value in each pixel in the depth_mask should be
the distance to the first object. Any particles
further away than that value will not be projected
onto the output image.
Returns: A BxHxW tensor of the projected particles.
"""
# Error checking.
batch_size = locs.size()[0]
N = locs.size()[1]
ec.check_tensor_dims(locs, "locs", (batch_size, N, 3))
ec.check_tensor_dims(camera_pose, "camera_pose", (batch_size, 3))
ec.check_tensor_dims(camera_rot, "camera_rot", (batch_size, 4))
if depth_mask is not None:
ec.check_tensor_dims(
depth_mask, "depth_mask",
(batch_size, self.camera_size[1], self.camera_size[0]))
depth_mask = depth_mask.contiguous()
else:
if self.empty_depth_mask.size()[0] != batch_size:
self.empty_depth_mask.resize_(batch_size, self.camera_size[1],
self.camera_size[0])
self.empty_depth_mask.fill_(MAX_FLOAT)
depth_mask = torch.autograd.Variable(self.empty_depth_mask,
requires_grad=False)
if locs.is_cuda:
depth_mask = depth_mask.cuda()
# Let's transform the particles to camera space here.
locs = locs - camera_pose.unsqueeze(1)
# Ensure the rotation quaternion is normalized.
camera_rot = camera_rot / torch.sqrt(
torch.sum(camera_rot**2, 1, keepdim=True))
# Invert the rotation.
inv = camera_rot.data.new(1, 4)
inv[0, 0] = -1
inv[0, 1] = -1
inv[0, 2] = -1
inv[0, 3] = 1
inv = torch.autograd.Variable(inv, requires_grad=False)
camera_rot = camera_rot * inv
rot = self._rotationMatrixFromQuaternion(camera_rot)
# Rotate the locs into camera space.
locs = torch.bmm(locs, rot)
locs = locs.contiguous()
proj = _ParticleProjectionFunction(self.camera_fl, self.camera_size,
self.filter_std, self.filter_scale)
ret = proj(locs, depth_mask)
return ret
