UPD: added setup.py for installation

SegmentAnything2AssistApp.py

@@ -4,7 +4,7 @@ import gradio_imageslider
 import spaces
 import torch
 
-import src.SegmentAnything2Assist as SegmentAnything2Assist
+import src.SegmentAnything2Assist.SegmentAnything2Assist as SegmentAnything2Assist
 
 example_image_annotation = {
     "image": "assets/cars.jpg",

setup.py

@@ -0,0 +1,25 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="SegmentAnything2Assist",
+    version="0.1",
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
+    install_requires=[
+        "SAM-2 @ git+https://github.com/facebookresearch/segment-anything-2.git@7e1596c0b6462eb1d1ba7e1492430fed95023598",
+        "ultralytics @ git+https://github.com/THU-MIG/yolov10.git@cd2f79c70299c9041fb6d19617ef1296f47575b1",
+        "opencv-python==4.10.0.84",
+    ],
+    author="xqt",
+    author_email="[email protected]",
+    description="A package to segment anything and assist in the process",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="https://huggingface.co/spaces/xqt/Segment-Anything-2-Assist",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires=">=3.8.0",
+)

src/{SegmentAnything2Assist.py → SegmentAnything2Assist/SegmentAnything2Assist.py}

@@ -5,12 +5,11 @@ import tqdm
 import requests
 import torch
 import numpy
-import pickle
 
 import sam2.build_sam
 import sam2.automatic_mask_generator
 
-from . import YOLOv10Plugin
+from .Plugin import YOLOv10Plugin
 
 import cv2
 
@@ -122,14 +121,17 @@ class SegmentAnything2Assist:
             print(f"SegmentAnything2Assist::is_model_available::{ret}")
         return ret
 
-    def load_model(self) -> None:
+    def load_model(self) -> bool:
         if self.is_model_available():
             self.sam2 = sam2.build_sam(checkpoint=self.model_path)
+            return True
 
-    def download_model(self, force: bool = False) -> None:
+        return False
+
+    def download_model(self, force: bool = False) -> bool:
         if not force and self.is_model_available():
             print(f"{self.model_path} already exists. Skipping download.")
-            return
+            return False
 
         response = requests.get(self.download_url, stream=True)
         total_size = int(response.headers.get("content-length", 0))
@@ -141,10 +143,12 @@ class SegmentAnything2Assist:
                 file.write(data)
                 progress_bar.update(len(data))
 
+        return True
+
     def generate_automatic_masks(
         self,
-        image,
-        points_per_side=32,
+        image: numpy.ndarray,
+        points_per_side=10,
         points_per_batch=32,
         pred_iou_thresh=0.8,
         stability_score_thresh=0.95,

test/test_module.py

@@ -0,0 +1,59 @@
+import unittest
+import src.SegmentAnything2Assist.SegmentAnything2Assist as SegmentAnything2Assist
+import cv2
+
+
+class TestSegmentAnything2Assist(unittest.TestCase):
+    def setUp(self) -> None:
+        return super().setUp()
+
+    def tearDown(self) -> None:
+        return super().tearDown()
+
+    def _loading_all_sam_model_types(self):
+        # Test loading all types of SAM2 models.
+        all_sam_models_type = [
+            "sam2_hiera_tiny",
+            "sam2_hiera_small",
+            "sam2_hiera_base_plus",
+            "sam2_hiera_large",
+        ]
+        for sam_model_type in all_sam_models_type:
+            sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
+                sam_model_name=sam_model_type, download=True, device="cpu"
+            )
+            self.assertEqual(sam_model.is_model_available(), True)
+
+            sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
+                sam_model_name=sam_model_type,
+                download=False,
+                model_path=f".tmp/checkpoints/{sam_model_type}.pth",
+                device="cpu",
+            )
+
+            with self.assertRaises(Exception):
+                sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
+                    sam_model_name=sam_model_type,
+                    download=False,
+                    model_path=".",
+                    device="cpu",
+                )
+
+    def test_generate_automatic_mask(self):
+        image = cv2.imread("test/assets/liberty.jpg")
+
+        sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
+            sam_model_name="sam2_hiera_tiny", download=True, device="cpu"
+        )
+
+        masks, segmentation_masks, bboxes = sam_model.generate_automatic_masks(image)
+
+        print(type(masks[0]))
+        print(type(segmentation_masks[0]))
+        print(type(bboxes[0]))
+
+        self.assertEqual(len(masks), len(segmentation_masks))
+        self.assertEqual(len(masks), len(bboxes))
+
+        # for mask, segmentation_mask, bbox in zip(masks, segmentation_masks, bboxes):
+        self.assertEqual(segmentation_masks[0].shape, image.shape)