fix: correct dims in case of automatic mask generation

pyproject.toml

@@ -3,7 +3,7 @@ requires = ["flit_core >=3.2,<4"]
 build-backend = "flit_core.buildapi"
 
 [project]
-description = "fork of the official SAMv2 implementation with cpu support"
+description = "CPU compatible fork of the official SAMv2 implementation"
 name = "samv2"
 version = "0.0.4"
 authors = [{ name = "Saurav Maheshkar", email = "[email protected]" }]

sam2/utils/visualization.py

@@ -32,8 +32,8 @@ def show_masks(
             output_image = Image.new(
                 mode="RGBA",
                 size=(
-                    masks[0]["segmentation"].shape[0],
                     masks[0]["segmentation"].shape[1],
+                    masks[0]["segmentation"].shape[0],
                 ),
                 color=(0, 0, 0),
             )

tests/test_image_inference.py

@@ -62,4 +62,14 @@ def test_prompt_mode(load_image, image_predictor, point, label, box, num_masks)
 @pytest.mark.full
 def test_mask_generation(load_image, mask_generator) -> None:
     masks = mask_generator.generate(load_image)
+
+    output_mask = show_masks(
+        image=load_image,
+        masks=masks,  # type: ignore
+        scores=None,
+        only_best=False,
+        autogenerated_mask=True,
+    )
+
     assert len(masks) > 0
+    assert isinstance(output_mask, Image.Image)