Rewrite draw_segmentation_masks and update gallery example to illustrate both instance and semantic segmentation models #3824
Conversation
NicolasHug
changed the title
NicolasHug
changed the title
There was a problem hiding this comment.
It looks great to me. Thanks @NicolasHug.
I know that you were not convinced about removing the float support but thank you for removing it from the PR to let more discussions around the issue.
This specific enhancement looks very nice and I think it's a great idea to expand the tool to support semantic segmentation. For me this is good to merge.
There was a problem hiding this comment.
Looks great, thanks a lot @NicolasHug !
I have a minor comment, otherwise looks good to me.
gallery/plot_visualization_utils.py
Outdated
| sem_class_to_idx = {cls: idx for (idx, cls) in enumerate(sem_classes)} | ||
|
|
||
| # We normalize the masks of each image in the batch independently | ||
| normalized_masks = torch.stack([torch.nn.Softmax(dim=0)(masks) for masks in output]) |
There was a problem hiding this comment.
Can you instead do torch.nn.functional.softmax(masks, dim=0) or maybe masks.softmax(dim=0)?
Actually, you can probably just do
normalized_masks = output.softmax(dim=1)no?
I think it's an anti-pattern to instantiate the nn.Module just to perform this operation once
There was a problem hiding this comment.
Ok, I'll use the functional API.
normalized_masks = output.softmax(dim=1) I think this would compute the softmax across the batch as well (dimension 0), and we wouldn't end up with values that sum to 1 for each image independently, the values would sum to 1 across the entire batch. That's actually something I wanted to double check with you, WDYT?
Edit: Hm they do seem to be the same. I'll use the simpler version as suggested then
There was a problem hiding this comment.
I think it should do the expected thing if we use output.softmax(dim=1), but it's a good question.
We use it already in the loss for the segmentation models, which internally is handled by passing dim=1 to (log)softmax.
Internally, the softmax implementation splits the dimensions into those that are "batch dimensions" (which comes before the dim) and the "to be softmax" dimensions (which comes after `dim), see here for more details
|
So glad that the implementation is improved now! 😃 🎉 |
|
Thanks for your initial work @oke-aditya , it made this next version much easier to write and test! |
…o illustrate both instance and semantic segmentation models (#3824) Reviewed By: cpuhrsch Differential Revision: D28538765 fbshipit-source-id: c8ff00c9e86f99f52e023a8b41c845ce51337671
This PR changes the API of
draw_segmentation_masksto something more generic that can support all models. See #3820 for more details.It also significantly updates the gallery example with didactic instructions about the models outputs, and how to draw the masks with the new API.