Visualization support for single channel images #288
Conversation
frgfm
left a comment
frgfm
left a comment
There was a problem hiding this comment.
Thanks for the PR!
Good fix, let's only add one test to avoid regression and I added a comment to make this more robust. Could you add a test case here: https://github.com/frgfm/torch-cam/blob/main/tests/test_utils.py ?
torchcam/utils.py
Outdated
| if len(img.getbands()) == 1: | ||
| overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, 0]).astype(np.uint8) | ||
| else: | ||
| overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :3]).astype(np.uint8) | ||
|
|
There was a problem hiding this comment.
Let's make this even more robust:
- the visualization can only work properly for single channel & 3 channel colormaps
- so let's raise an AssertionError if we don't have one of those two even before creating the
cmap. - since we're only doing broadcasting, I suggest to do:
overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :len(img.getbands())]).astype(np.uint8)There was a problem hiding this comment.
overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, 2 if len(img.getbands())==1 else slice(0, 3)]).astype(np.uint8)
this will be more correct
frgfm
added
module: utils
|
PR updated, test added |
frgfm
left a comment
frgfm
left a comment
There was a problem hiding this comment.
Thanks! Just a few adjustments left, and could you run make style or the precommit to fix the linter errors?
frgfm
left a comment
frgfm
left a comment
There was a problem hiding this comment.
Thanks for the edits, only a final question :)
torchcam/utils.py
Outdated
| overlay = mask.resize(img.size, resample=Resampling.BICUBIC) | ||
| overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :3]).astype(np.uint8) | ||
|
|
||
| overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, 2 if len(img.getbands()) == 1 else slice(0, 3)]).astype( |
There was a problem hiding this comment.
Final question: why taking the last channel, not first or second? Is any of those better for viz in grayscale? Or should we average accross the channel dimension?
frgfm
left a comment
frgfm
left a comment
There was a problem hiding this comment.
I just edited the final detail, thanks again for the PR!
2 participants
What does this PR do?
Added support for single channel images into overlay_mask