Spatial Soft-argmax (+optional scale adaptive)

[bhack]

Describe the feature and the current behavior/state.
Recover Spatial Soft-argmax from contrib + scale adaptive extension

Relevant information

• Are you willing to contribute it (yes/no): no
• Are you willing to maintain it going forward? (yes/no): yes
• Is there a relevant academic paper? (if so, where): https://arxiv.org/abs/2003.07543 (for the scale adaptive version)
• Is there already an implementation in another framework? (if so, where): Kornia (Scale adaptation extended formulation is still not released)
• Was it part of tf.contrib? (if so, where): SpatialSoftmax tensorflow#23331 (only the non-scale adaptive version)

Which API type would this fall under (layer, metric, optimizer, etc.)
layer

Who will benefit with this feature?
Any application that requires heatmap based keypoints/landmarks regression

Any other info.

[bhack]

/cc @ericjang for https://github.com/google-research/tensor2robot/blob/master/layers/spatial_softmax.py#L19

[0xshreyash]

I can take a look at this in the weekend if you guys still need help?

[bhack]

@shreyashpatodia How do you want to organize the scale parameter API? With the scale adaptable extension the spatial soft argmax ops need to work inside ROI windows in every channel.

[bhack]

It partially relate on how we want to design basic keypoints<->heatmap ops. See #1366

[bhack]

I think that something like tf.image.crop_to_bounding_box but that works with a batch of offsets it is needed as parameter to work with Scales.

[bhack]

Probably the ROIs could be formulated like boxes in tf.image.draw_bounding_boxes

[bhack]

@shreyashpatodia Do you plan to make a PR?

[bhack]

@gabrieldemarmiesse @seanpmorgan Do you remember why we lost the tf.contrib version in the migration? In the table I see:

One OSS project found / Needs refactored as base Layer subclass / Uses get_variable_collections

[seanpmorgan]

@gabrieldemarmiesse @seanpmorgan Do you remember why we lost the tf.contrib version in the migration? In the table I see:

One OSS project found / Needs refactored as base Layer subclass / Uses get_variable_collections

Just means there wasn't enough OSS usage to justify moving it. During the RFC period all candidates to be dropped were welcome to be picked up if there was community support, but no one did for spatial_softmax. Happy to review a PR for it, but the tf.contrib version used variable scoping etc. so there are some changes to be made.

[bhack]

What do you think about the Google version that I mentioned early in this ticket?