dora-sam2 0.0.0__py3-none-any.whl

dora_sam2/__init__.py

@@ -0,0 +1,11 @@
+import os
+
+# Define the path to the README file relative to the package directory
+readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md")
+
+# Read the content of the README file
+try:
+    with open(readme_path, encoding="utf-8") as f:
+        __doc__ = f.read()
+except FileNotFoundError:
+    __doc__ = "README file not found."

dora_sam2/__main__.py

@@ -0,0 +1,4 @@
+from .main import main
+
+if __name__ == "__main__":
+    main()

dora_sam2/main.py

@@ -0,0 +1,94 @@
+import cv2
+import numpy as np
+import pyarrow as pa
+import torch
+from dora import Node
+from PIL import Image
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+
+
+def main():
+    pa.array([])  # initialize pyarrow array
+    node = Node()
+    frames = {}
+    for event in node:
+        event_type = event["type"]
+
+        if event_type == "INPUT":
+            event_id = event["id"]
+
+            if "image" in event_id:
+                storage = event["value"]
+                metadata = event["metadata"]
+                encoding = metadata["encoding"]
+                width = metadata["width"]
+                height = metadata["height"]
+
+                if (
+                    encoding == "bgr8"
+                    or encoding == "rgb8"
+                    or encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]
+                ):
+                    channels = 3
+                    storage_type = np.uint8
+                else:
+                    error = f"Unsupported image encoding: {encoding}"
+                    raise RuntimeError(error)
+
+                if encoding == "bgr8":
+                    frame = (
+                        storage.to_numpy()
+                        .astype(storage_type)
+                        .reshape((height, width, channels))
+                    )
+                    frame = frame[:, :, ::-1]  # OpenCV image (BGR to RGB)
+                elif encoding == "rgb8":
+                    frame = (
+                        storage.to_numpy()
+                        .astype(storage_type)
+                        .reshape((height, width, channels))
+                    )
+                elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
+                    storage = storage.to_numpy()
+                    frame = cv2.imdecode(storage, cv2.IMREAD_COLOR)
+                    frame = frame[:, :, ::-1]  # OpenCV image (BGR to RGB)
+                else:
+                    raise RuntimeError(f"Unsupported image encoding: {encoding}")
+                image = Image.fromarray(frame)
+                frames[event_id] = image
+
+            elif "boxes2d" in event_id:
+                boxes2d = event["value"].to_numpy()
+                metadata = event["metadata"]
+                encoding = metadata["encoding"]
+                if encoding != "xyxy":
+                    raise RuntimeError(f"Unsupported boxes2d encoding: {encoding}")
+
+                image_id = metadata["image_id"]
+                with torch.inference_mode(), torch.autocast(
+                    "cuda",
+                    dtype=torch.bfloat16,
+                ):
+                    predictor.set_image(frames[image_id])
+                    masks, _, _ = predictor.predict(box=boxes2d)
+                    masks = masks[0]
+                    ## Mask to 3 channel image
+
+                    node.send_output(
+                        "masks",
+                        pa.array(masks.ravel()),
+                        metadata={
+                            "image_id": image_id,
+                            "width": frames[image_id].width,
+                            "height": frames[image_id].height,
+                        },
+                    )
+
+        elif event_type == "ERROR":
+            print("Event Error:" + event["error"])
+
+
+if __name__ == "__main__":
+    main()

dora_sam2-0.0.0.dist-info/METADATA

@@ -0,0 +1,53 @@
+Metadata-Version: 2.2
+Name: dora-sam2
+Version: 0.0.0
+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.6
+Requires-Dist: huggingface-hub>=0.29.0
+Requires-Dist: opencv-python>=4.11.0.86
+Requires-Dist: sam2>=1.1.0
+
+# dora-sam2
+
+> [!WARNING]  
+> SAM2 requires Nvidia GPU to be able to run.
+
+## Getting started
+
+- Install it with pip:
+
+```bash
+pip install -e .
+```
+
+## Contribution Guide
+
+- Format with [ruff](https://docs.astral.sh/ruff/):
+
+```bash
+ruff check . --fix
+```
+
+- Lint with ruff:
+
+```bash
+ruff check .
+```
+
+- Test with [pytest](https://github.com/pytest-dev/pytest)
+
+```bash
+pytest . # Test
+```
+
+## YAML Specification
+
+## Examples
+
+## License
+
+dora-sam2's code are released under the MIT License

dora_sam2-0.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+dora_sam2/__init__.py,sha256=HuSK3dnyI9Pb5QAuaKFwQQ3J5SIZnLcKHPJO0norGzc,353
+dora_sam2/__main__.py,sha256=Vdhw8YA1K3wPMlbJQYL5WqvRzAKVeZ16mZQFO9VRmCo,62
+dora_sam2/main.py,sha256=7z7kZyI25eOSqnbwUEL8bGXOUnoa6UQ_AXThrjHUAB0,3380
+dora_sam2-0.0.0.dist-info/METADATA,sha256=If_OUNOaxv2wAK8BIze_gb07Bh9SpeMiMksV3N0a4qY,819
+dora_sam2-0.0.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+dora_sam2-0.0.0.dist-info/entry_points.txt,sha256=eObMaDQauVA_sv4J6fsNjn8V-8syGJK7mO-LrsBu1aA,50
+dora_sam2-0.0.0.dist-info/top_level.txt,sha256=IgKcOITGe2Nlyc79J6dwh3dcp3Wsf-IipIy-1h9GcPE,10
+dora_sam2-0.0.0.dist-info/RECORD,,

dora_sam2-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

dora_sam2-0.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+dora-sam2 = dora_sam2.main:main

dora_sam2-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+dora_sam2