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