dora-sam2 0.3.11__py3-none-any.whl → 0.3.12__py3-none-any.whl

dora_sam2/main.py

@@ -133,7 +133,9 @@ def main():
                                 )
 
             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 = (
@@ -162,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:

{dora_sam2-0.3.11.dist-info → dora_sam2-0.3.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dora-sam2
-Version: 0.3.11
+Version: 0.3.12
 Summary: dora-sam2
 Author-email: Your Name <email@email.com>
 License: MIT

dora_sam2-0.3.12.dist-info/RECORD

@@ -0,0 +1,8 @@
+dora_sam2/__init__.py,sha256=tF7WHhHiDweUUzyHsbmFe_ktphE08aA5j33E4ja1udA,381
+dora_sam2/__main__.py,sha256=NlAb6Jbmmn82K8Ahdi12sliZYdzyY7QaoCCHRuoR_Hg,90
+dora_sam2/main.py,sha256=Jo6y5vY8LafpdUfb_gHgy7zLJMZz_8sLBw0gL3HeG0M,10873
+dora_sam2-0.3.12.dist-info/METADATA,sha256=IW0SXlpL6v3CPkV9Rf7rNdeB0RUDq1TmTFwJ-tcOS1w,820
+dora_sam2-0.3.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dora_sam2-0.3.12.dist-info/entry_points.txt,sha256=eObMaDQauVA_sv4J6fsNjn8V-8syGJK7mO-LrsBu1aA,50
+dora_sam2-0.3.12.dist-info/top_level.txt,sha256=IgKcOITGe2Nlyc79J6dwh3dcp3Wsf-IipIy-1h9GcPE,10
+dora_sam2-0.3.12.dist-info/RECORD,,

{dora_sam2-0.3.11.dist-info → dora_sam2-0.3.12.dist-info}/WHEEL

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

dora_sam2-0.3.11.dist-info/RECORD

@@ -1,8 +0,0 @@
-dora_sam2/__init__.py,sha256=tF7WHhHiDweUUzyHsbmFe_ktphE08aA5j33E4ja1udA,381
-dora_sam2/__main__.py,sha256=NlAb6Jbmmn82K8Ahdi12sliZYdzyY7QaoCCHRuoR_Hg,90
-dora_sam2/main.py,sha256=aQj1zZTtpzebwaayyPqD4B82iFscLKuZ5v7COs3dTNE,8487
-dora_sam2-0.3.11.dist-info/METADATA,sha256=jTkg2BE_lLdUTroTRkZ6hdsT2xwwtF9Z2D_SYtgG8_Y,820
-dora_sam2-0.3.11.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-dora_sam2-0.3.11.dist-info/entry_points.txt,sha256=eObMaDQauVA_sv4J6fsNjn8V-8syGJK7mO-LrsBu1aA,50
-dora_sam2-0.3.11.dist-info/top_level.txt,sha256=IgKcOITGe2Nlyc79J6dwh3dcp3Wsf-IipIy-1h9GcPE,10
-dora_sam2-0.3.11.dist-info/RECORD,,