Skip to content

Fix masks_to_boxes for empty masks and vectorize for performance#9347

Closed
raimbekovm wants to merge 2 commits intopytorch:mainfrom
raimbekovm:fix/masks-to-boxes-empty-masks
Closed

Fix masks_to_boxes for empty masks and vectorize for performance#9347
raimbekovm wants to merge 2 commits intopytorch:mainfrom
raimbekovm:fix/masks-to-boxes-empty-masks

Conversation

@raimbekovm
Copy link
Copy Markdown
Contributor

@raimbekovm raimbekovm commented Jan 21, 2026

Fixes #9346

masks_to_boxes raises RuntimeError when a mask is all zeros because torch.min/max on empty tensor requires explicit dim.

This PR:

  • Handles empty masks by returning [0, 0, 0, 0] bounding boxes
  • Vectorizes the implementation (removes Python loop) for better performance on large batches
  • Adds tests for empty and mixed empty/non-empty mask batches

This is based on the approach from #8194 but fixes the empty mask handling issue that was identified there (the vectorized argmax returns [0, 0, W-1, H-1] for empty masks instead of [0, 0, 0, 0]). We detect empty masks explicitly and set their boxes to zeros.

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Jan 21, 2026 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9347

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 0f8f3f1 with merge base 6acedba (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla bot added the cla signed label Jan 21, 2026
@raimbekovm raimbekovm mentioned this pull request Jan 21, 2026
Closed
@zy1git
Copy link
Copy Markdown
Contributor

zy1git commented Jan 27, 2026

@raimbekovm Thanks a lot for your contribution! Could you please separate this into two PRs?
First PR: The bugfix for empty masks (keeping the loop-based approach)
Second PR: The vectorization for performance
The reason is that the bugfix is clearly needed, but the vectorization comes with a memory tradeoff that we'll need to evaluate separately. Splitting them allows us to merge the bugfix quickly while we assess the performance optimization.
We will review each PR accordingly. Thanks!

This was referenced Jan 27, 2026
Merged
Merged
@raimbekovm
Copy link
Copy Markdown
Contributor Author

raimbekovm commented Jan 27, 2026

Thanks for the feedback! I've split this into two PRs as suggested:

Closing this one.

@raimbekovm raimbekovm closed this Jan 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

cla signed

Projects

None yet

Milestone

No milestone

2 participants

@raimbekovm @zy1git