in utils.py [0:0]
def calc_local_chamfer(samples, batch, losses):
batch_size = samples.shape[0]
# a grid of point projected towards the surface of the object, starting from the same position and orientation
# as the touch sensor when the touch occurred, but 5 times its size
planes = batch['radius'].cuda().view(batch_size, 4, 100, 100, 3)
# mask indicating which point hit the surface of the object, ie, tho ones we care about
masks = batch['radius_masks'].cuda().view(batch_size, 4, 100, 100)
successful = batch['successful']
num_examples = 0
# for every grasps
for pred, gt, mask, success in zip(samples, planes, masks, successful):
# for every ring size around each touch site
for i in range(5):
# for every touch
for j in range(4):
if not success[j]:
continue
# select the right ring of points, ie 1 x size of sensor ... 5 x size of sensor
dim_mask = torch.zeros(mask[j].shape).clone()
dim_mask[40 - i * 10: 60 + i * 10, 40 - i * 10: 60 + i * 10] = 1
dim_mask[50 - i * 10: 50 + i * 10, 50 - i * 10: 50 + i * 10] = 0
# select point which are on the objects surface
dim_mask[mask[j] == 0] = 0
gt_masked = gt[j][dim_mask == 1]
if (gt_masked.shape[0] == 0):
continue
# compute the local loss between the selected points and the predicted surface
indices, _ = chamfer_dist(gt_masked.unsqueeze(0), pred.unsqueeze(0))
pred_counter = pred[indices.long()[0]]
loss = (torch.sum((pred_counter - gt_masked) ** 2, dim=1)).mean()
losses[i] += loss
if i == 0:
num_examples += 1.
return losses, num_examples