-
-
Notifications
You must be signed in to change notification settings - Fork 211
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Visualization support for single channel images #288
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
PR updated, test added |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks! Just a few adjustments left, and could you run make style
or the precommit to fix the linter errors?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the edits, only a final question :)
cmap = cm.get_cmap(colormap) | ||
# Resize mask and apply colormap | ||
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
What does this PR do?
Added support for single channel images into overlay_mask