Skip to content
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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

DmitriyValetov
Copy link

What does this PR do?

Added support for single channel images into overlay_mask

Copy link
Owner

@frgfm frgfm left a 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 ?

Comment on lines 45 to 49
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)

Copy link
Owner

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)

Copy link
Author

@DmitriyValetov DmitriyValetov Nov 16, 2024

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

@frgfm frgfm self-assigned this Nov 13, 2024
@frgfm frgfm added module: utils Related to torchcam.utils type: improvement New feature or request labels Nov 13, 2024
@DmitriyValetov
Copy link
Author

PR updated, test added

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: utils Related to torchcam.utils type: improvement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants