facebookresearch / facebookresearch/sam2
Bugfix: Dimension Mismatch in AutoMaskGenerator
- Dominant language
- Jupyter Notebook
- Stars
- 19.9k
- Forks
- 2.5k
- PR merge metrics
- No merged PRs in 30d
Description
## Dimension Mismatch caused by wrong if condition
In [Automatic Mask Generator, line 236](https://github.com/facebookresearch/sam2/blob/main/sam2/automatic_mask_generator.py), the wrong condition `len(crop_boxes)` may cause box area calculation of an empty box `scores = 1 / box_area(data["crop_boxes"])`.
```python
# Remove duplicate masks between crops
if len(crop_boxes) > 1: # <-- HERE
# Prefer masks from smaller crops
scores = 1 / box_area(data["crop_boxes"])
scores = scores.to(data["boxes"].device)
keep_by_nms = batched_nms(
data["boxes"].float(),
scores,
torch.zeros_like(data["boxes"][:, 0]), # categories
iou_threshold=self.crop_nms_thresh,
)
```
**The condition should be** `if len(data["crop_boxes"]) > 1:`. However, I haven't read the code thoroughly, so I'm not sure if i fixed it right.
_Problem Traceback:_
```cuda
File "/root/miniconda3/envs/segzero/lib/python3.12/site-packages/sam2/automatic_mask_generator.py", line 196, in generate
mask_data = self._generate_masks(image)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/envs/segzero/lib/python3.12/site-packages/sam2/automatic_mask_generator.py", line 239, in _generate_masks
scores = 1 / box_area(data["crop_boxes"])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/envs/segzero/lib/python3.12/site-packages/torchvision/ops/boxes.py", line 251, in box_area
return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
~~~~~^^^^^^
IndexError: too many indices for tensor of dimension 1
```
Contributor guide
Assessment
This issue has not been assessed yet.