fix mask shape bug

This commit is contained in:
rentainhe
2024-08-09 01:54:40 +08:00
parent 9c5786fc09
commit 80676e866b
4 changed files with 5 additions and 17 deletions

View File

@@ -14,7 +14,7 @@ Grounded SAM 2 does not introduce significant methodological changes compared to
## Contents ## Contents
- [Installation](#installation) - [Installation](#installation)
- [Grounded SAM 2 Demo](#grounded-sam-2-demo) - [Grounded SAM 2 Demo](#grounded-sam-2-demos)
- [Grounded SAM 2 Image Demo](#grounded-sam-2-image-demo-with-grounding-dino) - [Grounded SAM 2 Image Demo](#grounded-sam-2-image-demo-with-grounding-dino)
- [Grounded SAM 2 Image Demo (with Grounding DINO 1.5 & 1.6)](#grounded-sam-2-image-demo-with-grounding-dino-15--16) - [Grounded SAM 2 Image Demo (with Grounding DINO 1.5 & 1.6)](#grounded-sam-2-image-demo-with-grounding-dino-15--16)
- [Grounded SAM 2 Video Object Tracking Demo](#grounded-sam-2-video-object-tracking-demo) - [Grounded SAM 2 Video Object Tracking Demo](#grounded-sam-2-video-object-tracking-demo)
@@ -65,7 +65,7 @@ wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alp
wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha2/groundingdino_swinb_cogcoor.pth wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha2/groundingdino_swinb_cogcoor.pth
``` ```
## Grounded SAM 2 Demo ## Grounded SAM 2 Demos
### Grounded SAM 2 Image Demo (with Grounding DINO) ### Grounded SAM 2 Image Demo (with Grounding DINO)
Note that `Grounding DINO` has already been supported in [Huggingface](https://huggingface.co/IDEA-Research/grounding-dino-tiny), so we provide two choices for running `Grounded SAM 2` model: Note that `Grounding DINO` has already been supported in [Huggingface](https://huggingface.co/IDEA-Research/grounding-dino-tiny), so we provide two choices for running `Grounded SAM 2` model:
- Use huggingface API to inference Grounding DINO (which is simple and clear) - Use huggingface API to inference Grounding DINO (which is simple and clear)

View File

@@ -89,11 +89,7 @@ masks, scores, logits = sam2_predictor.predict(
Post-process the output of the model to get the masks, scores, and logits for visualization Post-process the output of the model to get the masks, scores, and logits for visualization
""" """
# convert the shape to (n, H, W) # convert the shape to (n, H, W)
if masks.ndim == 3: if masks.ndim == 4:
masks = masks[None]
scores = scores[None]
logits = logits[None]
elif masks.ndim == 4:
masks = masks.squeeze(1) masks = masks.squeeze(1)

View File

@@ -81,11 +81,7 @@ masks, scores, logits = sam2_predictor.predict(
Post-process the output of the model to get the masks, scores, and logits for visualization Post-process the output of the model to get the masks, scores, and logits for visualization
""" """
# convert the shape to (n, H, W) # convert the shape to (n, H, W)
if masks.ndim == 3: if masks.ndim == 4:
masks = masks[None]
scores = scores[None]
logits = logits[None]
elif masks.ndim == 4:
masks = masks.squeeze(1) masks = masks.squeeze(1)

View File

@@ -68,11 +68,7 @@ masks, scores, logits = sam2_predictor.predict(
Post-process the output of the model to get the masks, scores, and logits for visualization Post-process the output of the model to get the masks, scores, and logits for visualization
""" """
# convert the shape to (n, H, W) # convert the shape to (n, H, W)
if masks.ndim == 3: if masks.ndim == 4:
masks = masks[None]
scores = scores[None]
logits = logits[None]
elif masks.ndim == 4:
masks = masks.squeeze(1) masks = masks.squeeze(1)