TVTensor-KeyPoints with support for keypoint visibility for training Keypoint R-CNN

[pieterblok]

🚀 The feature

@Alexandre-SCHOEPP @NicolasHug @AntoineSimoulin

Thank you for releasing the tv_tensor keypoints functionality. I use it in my university teaching, where students build a custom keypoint dataloader for training Keypoint R-CNN.

The tv_tensor utilities correctly handle the geometric transformations, but unfortunately they do not yet support the visibility flag, which is required for training Keypoint R-CNN.

Ultimately, this leads to redundant and inefficient code when building custom dataloaders. See the snippet below:

with open(annotation_path) as f:
    data = json.load(f)
shapes = data["shapes"]

keypoints = []
for shape in shapes:
    cx, cy = shape["points"][0]
    keypoints.append([cx, cy])
keypoints = torch.tensor(keypoints, dtype=torch.float32)
keypoints = keypoints.view(-1, 1, 2)

target = {}
target["keypoints"] = tv_tensors.KeyPoints(keypoints, canvas_size=F.get_size(img))

if self.transforms is not None:
    img, target = self.transforms(img, target)

## add visibility flag - redundant code
rows = target['keypoints'].shape[0]
visibility = torch.full(((rows, 1, 1)), 2.0, device=target['keypoints'].device, dtype=target['keypoints'].dtype)
target['keypoints'] = torch.cat([target['keypoints'], visibility], dim=2)

Motivation, pitch

It would be much more efficient if I could include the visibility flag directly in the keypoints list and then wrap it with the tv_tensors.KeyPoints class. However, this is currently not possible due to the input tensor shape.

A useful new feature would be to apply the transformations only to the first two values (x, y) while leaving the third value (visibility) unchanged.

I’m curious to hear your thoughts on this potential feature. Thank you in advance.

Alternatives

No response

Additional context

No response

[Alexandre-SCHOEPP]

Hi @pieterblok !

Thanks for the proposal ! I think that this does make a lot of sense.

However, what I'm worried about is the possibility that people might mistake a 3rd dimension in the shape as being the 3rd spatial dimension. Perhaps a .format attribute, like was done for the BoundingBoxes, would make this clearer for the end-user.

[NicolasHug]

Hi @pieterblok , thanks for your feedback, and we're glad the KeyPoints are useful!

IIUC the visibility tensor is of shape num_keypoints (with potentially more than 1 dimension), and it's supposed to be passed-through? Instead of calling torch.cat I think you should just be able to pass the visibility tensor as-is to the transforms. E.g.:

# ...
target["keypoints"] = tv_tensors.KeyPoints(keypoints, canvas_size=F.get_size(img))
target["visibility"] = torch.full(...)

This avoids the call to torch.cat which, IIUC, is the source of the inefficiency you mentioned? Then if you ever call SanitizeKeyPoints you'd just need to make sure labels_getter can find the visibility tensor so the shapes are preserved. Hope this helps, let me know if I'm misunderstanding something

[pieterblok]

@Alexandre-SCHOEPP @NicolasHug

Thanks to both of you for the feedback! I think Alexandre’s suggestion regarding the .format input is the most appropriate direction.

According to the KeyPoints documentation, it would ideally be possible to specify the final dimension as either 2 or 3 elements—[x, y] or [x, y, v]—for tensors shaped like [N_chains, N_points, 2], [N_polygons, N_points, 2], or [N_skeletons, N_bones, 2, 2], which is currently bound by 2 points in the last dimension.

The issue is that embedding the visibility flag directly inside target["keypoints"] (as required by Keypoint R-CNN) is currently not supported. With the proposed format option in KeyPoints, the implementation could ideally look like this:

with open(annotation_path) as f:
    data = json.load(f)
shapes = data["shapes"]

keypoints = []
v = 2.0
for shape in shapes:
    cx, cy = shape["points"][0]
    keypoints.append([cx, cy, v])
keypoints = torch.tensor(keypoints, dtype=torch.float32)
keypoints = keypoints.view(-1, 1, 3)

target = {}
target["keypoints"] = tv_tensors.KeyPoints(keypoints, format="XYV", canvas_size=F.get_size(img))

if self.transforms is not None:
    img, target = self.transforms(img, target)

I hope this ideal implementation helps clarify the point I’m referring to.