dora-sam2 0.3.10rc1__tar.gz → 0.3.12rc0__tar.gz

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/PKG-INFO

@@ -1,12 +1,12 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dora-sam2
-Version: 0.3.10rc1
+Version: 0.3.12rc0
 Summary: dora-sam2
 Author-email: Your Name <email@email.com>
 License: MIT
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: dora-rs>=0.3.10rc1
+Requires-Dist: dora-rs>=0.3.9
 Requires-Dist: huggingface-hub>=0.29.0
 Requires-Dist: opencv-python>=4.11.0.86
 Requires-Dist: sam2>=1.1.0

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/dora_sam2/__init__.py

@@ -1,3 +1,5 @@
+"""TODO: Add docstring."""
+
 import os
 
 # Define the path to the README file relative to the package directory

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/dora_sam2/__main__.py

@@ -1,3 +1,5 @@
+"""TODO: Add docstring."""
+
 from .main import main
 
 if __name__ == "__main__":

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/dora_sam2/main.py

@@ -1,3 +1,5 @@
+"""TODO: Add docstring."""
+
 import cv2
 import numpy as np
 import pyarrow as pa
@@ -10,6 +12,7 @@ predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
 
 
 def main():
+    """TODO: Add docstring."""
     pa.array([])  # initialize pyarrow array
     node = Node()
     frames = {}
@@ -119,8 +122,8 @@ def main():
                                             {
                                                 "masks": masks.ravel(),
                                                 "labels": event["value"]["labels"],
-                                            }
-                                        ]
+                                            },
+                                        ],
                                     ),
                                     metadata={
                                         "image_id": image_id,
@@ -129,8 +132,10 @@ def main():
                                     },
                                 )
 
-            elif "boxes2d" in event_id:
-
+            if "boxes2d" in event_id:
+                if len(event["value"]) == 0:
+                    node.send_output("masks", pa.array([]))
+                    continue
                 if isinstance(event["value"], pa.StructArray):
                     boxes2d = event["value"][0].get("bbox").values.to_numpy()
                     labels = (
@@ -159,7 +164,59 @@ def main():
                 ):
                     predictor.set_image(frames[image_id])
                     masks, _scores, last_pred = predictor.predict(
-                        box=boxes2d, point_labels=labels, multimask_output=False
+                        box=boxes2d,
+                        point_labels=labels,
+                        multimask_output=False,
+                    )
+
+                    if len(masks.shape) == 4:
+                        masks = masks[:, 0, :, :]
+                        last_pred = last_pred[:, 0, :, :]
+                    else:
+                        masks = masks[0, :, :]
+                        last_pred = last_pred[0, :, :]
+
+                    masks = masks > 0
+                    metadata["image_id"] = image_id
+                    metadata["width"] = frames[image_id].width
+                    metadata["height"] = frames[image_id].height
+                    ## Mask to 3 channel image
+                    match return_type:
+                        case pa.Array:
+                            node.send_output("masks", pa.array(masks.ravel()), metadata)
+                        case pa.StructArray:
+                            node.send_output(
+                                "masks",
+                                pa.array(
+                                    [
+                                        {
+                                            "masks": masks.ravel(),
+                                            "labels": event["value"]["labels"],
+                                        },
+                                    ],
+                                ),
+                                metadata,
+                            )
+            elif "points" in event_id:
+                points = event["value"].to_numpy().reshape((-1, 2))
+                return_type = pa.Array
+                if len(frames) == 0:
+                    continue
+                first_image = next(iter(frames.keys()))
+                image_id = event["metadata"].get("image_id", first_image)
+                with (
+                    torch.inference_mode(),
+                    torch.autocast(
+                        "cuda",
+                        dtype=torch.bfloat16,
+                    ),
+                ):
+                    predictor.set_image(frames[image_id])
+                    labels = [i for i in range(len(points))]
+                    masks, _scores, last_pred = predictor.predict(
+                        points,
+                        point_labels=labels,
+                        multimask_output=False,
                     )
 
                     if len(masks.shape) == 4:
@@ -190,8 +247,8 @@ def main():
                                         {
                                             "masks": masks.ravel(),
                                             "labels": event["value"]["labels"],
-                                        }
-                                    ]
+                                        },
+                                    ],
                                 ),
                                 metadata={
                                     "image_id": image_id,

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/dora_sam2.egg-info/PKG-INFO

@@ -1,12 +1,12 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dora-sam2
-Version: 0.3.10rc1
+Version: 0.3.12rc0
 Summary: dora-sam2
 Author-email: Your Name <email@email.com>
 License: MIT
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: dora-rs>=0.3.10rc1
+Requires-Dist: dora-rs>=0.3.9
 Requires-Dist: huggingface-hub>=0.29.0
 Requires-Dist: opencv-python>=4.11.0.86
 Requires-Dist: sam2>=1.1.0

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/dora_sam2.egg-info/requires.txt

@@ -1,4 +1,4 @@
-dora-rs>=0.3.10rc1
+dora-rs>=0.3.9
 huggingface-hub>=0.29.0
 opencv-python>=4.11.0.86
 sam2>=1.1.0

dora_sam2-0.3.12rc0/pyproject.toml

@@ -0,0 +1,36 @@
+[project]
+name = "dora-sam2"
+version = "0.3.12-rc0"
+authors = [{ name = "Your Name", email = "email@email.com" }]
+description = "dora-sam2"
+license = { text = "MIT" }
+readme = "README.md"
+requires-python = ">=3.10"
+
+dependencies = [
+  "dora-rs >= 0.3.9",
+  "huggingface-hub>=0.29.0",
+  "opencv-python>=4.11.0.86",
+  "sam2>=1.1.0",
+]
+
+[tool.uv]
+no-build-isolation-package = ['sam2']
+
+[dependency-groups]
+dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
+
+[project.scripts]
+dora-sam2 = "dora_sam2.main:main"
+
+[tool.ruff.lint]
+extend-select = [
+  "D",    # pydocstyle
+  "UP",   # Ruff's UP rule
+  "PERF", # Ruff's PERF rule
+  "RET",  # Ruff's RET rule
+  "RSE",  # Ruff's RSE rule
+  "NPY",  # Ruff's NPY rule
+  "N",    # Ruff's N rule
+  "I",    # Ruff's I rule
+]

{dora_sam2-0.3.10rc1 → dora_sam2-0.3.12rc0}/tests/test_dora_sam2.py

@@ -1,3 +1,5 @@
+"""TODO: Add docstring."""
+
 import os
 
 import pytest
@@ -6,6 +8,7 @@ CI = os.getenv("CI", "false") in ["True", "true"]
 
 
 def test_import_main():
+    """TODO: Add docstring."""
     if CI:
         # Skip test as test requires Nvidia GPU
         return

dora_sam2-0.3.10rc1/pyproject.toml

@@ -1,21 +0,0 @@
-[project]
-name = "dora-sam2"
-version = "0.3.10-rc1"
-authors = [{ name = "Your Name", email = "email@email.com" }]
-description = "dora-sam2"
-license = { text = "MIT" }
-readme = "README.md"
-requires-python = ">=3.10"
-
-dependencies = [
-    "dora-rs >= 0.3.10rc1",
-    "huggingface-hub>=0.29.0",
-    "opencv-python>=4.11.0.86",
-    "sam2>=1.1.0",
-]
-
-[dependency-groups]
-dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
-
-[project.scripts]
-dora-sam2 = "dora_sam2.main:main"