fix(video): append models only if they exist (#1117)

Author: shcheklein

src/datachain/model/ultralytics/bbox.py

@@ -33,7 +33,7 @@ def from_result(result: "Results") -> "YoloBBox":
         name = summary[0].get("name", "")
         box = (
             BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
+            if summary[0].get("box")
             else BBox()
         )
         return YoloBBox(
@@ -69,7 +69,8 @@ def from_results(results: list["Results"]) -> "YoloBBoxes":
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
+                if s.get("box"):
+                    box.append(BBox.from_dict(s.get("box"), title=name))
         return YoloBBoxes(
             cls=cls,
             name=names,
@@ -102,7 +103,7 @@ def from_result(result: "Results") -> "YoloOBBox":
         name = summary[0].get("name", "")
         box = (
             OBBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
+            if summary[0].get("box")
             else OBBox()
         )
         return YoloOBBox(
@@ -138,7 +139,8 @@ def from_results(results: list["Results"]) -> "YoloOBBoxes":
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(OBBox.from_dict(s.get("box", {}), title=name))
+                if s.get("box"):
+                    box.append(OBBox.from_dict(s.get("box"), title=name))
         return YoloOBBoxes(
             cls=cls,
             name=names,

src/datachain/model/ultralytics/pose.py

@@ -58,12 +58,12 @@ def from_result(result: "Results") -> "YoloPose":
         name = summary[0].get("name", "")
         box = (
             BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
+            if summary[0].get("box")
             else BBox()
         )
         pose = (
             Pose3D.from_dict(summary[0]["keypoints"])
-            if "keypoints" in summary[0]
+            if summary[0].get("keypoints")
             else Pose3D()
         )
         return YoloPose(
@@ -102,8 +102,10 @@ def from_results(results: list["Results"]) -> "YoloPoses":
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                pose.append(Pose3D.from_dict(s.get("keypoints", {})))
+                if s.get("box"):
+                    box.append(BBox.from_dict(s.get("box"), title=name))
+                if s.get("keypoints"):
+                    pose.append(Pose3D.from_dict(s.get("keypoints")))
         return YoloPoses(
             cls=cls,
             name=names,

src/datachain/model/ultralytics/segment.py

@@ -36,12 +36,12 @@ def from_result(result: "Results") -> "YoloSegment":
         name = summary[0].get("name", "")
         box = (
             BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
+            if summary[0].get("box")
             else BBox()
         )
         segment = (
             Segment.from_dict(summary[0]["segments"], title=name)
-            if "segments" in summary[0]
+            if summary[0].get("segments")
             else Segment()
         )
         return YoloSegment(
@@ -80,8 +80,10 @@ def from_results(results: list["Results"]) -> "YoloSegments":
                 cls.append(s["class"])
                 names.append(name)
                 confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                segment.append(Segment.from_dict(s.get("segments", {}), title=name))
+                if s.get("box"):
+                    box.append(BBox.from_dict(s.get("box"), title=name))
+                if s.get("segments"):
+                    segment.append(Segment.from_dict(s.get("segments"), title=name))
         return YoloSegments(
             cls=cls,
             name=names,

tests/func/model/test_yolo.py

@@ -205,6 +205,45 @@ def test_yolo_segment_from_results_empty(running_img):
     }
 
 
+def test_yolo_segment_from_results_empty_segments(running_img):
+    result = Results(
+        orig_img=running_img,
+        path="running.jpeg",
+        names={0: "person"},
+        boxes=torch.tensor([[102.0, 84.0, 183.0, 238.0, 0.9078, 0.0]]),
+    )
+
+    model = YoloSegment.from_result(result)
+    assert model.model_dump() == {
+        "cls": 0,
+        "name": "person",
+        "confidence": 0.9078,
+        "box": {
+            "coords": [102, 84, 183, 238],
+            "title": "person",
+        },
+        "segment": {
+            "title": "",
+            "x": [],
+            "y": [],
+        },
+    }
+
+    model = YoloSegments.from_results([result])
+    assert model.model_dump() == {
+        "cls": [0],
+        "name": ["person"],
+        "confidence": [0.9078],
+        "box": [
+            {
+                "coords": [102, 84, 183, 238],
+                "title": "person",
+            }
+        ],
+        "segment": [],
+    }
+
+
 def test_yolo_seg_from_results(running_img, running_img_masks):
     result = Results(
         orig_img=running_img,
@@ -2239,6 +2278,40 @@ def test_yolo_pose_from_results_empty(running_img):
     }
 
 
+def test_yolo_pose_from_results_empty_poses(running_img):
+    result = Results(
+        orig_img=running_img,
+        path="running.jpeg",
+        names={0: "person"},
+        boxes=torch.tensor([[102.0, 84.0, 183.0, 238.0, 0.9078, 0.0]]),
+    )
+
+    model = YoloPose.from_result(result)
+    assert model.model_dump() == {
+        "cls": 0,
+        "name": "person",
+        "confidence": 0.9078,
+        "box": {
+            "coords": [102, 84, 183, 238],
+            "title": "person",
+        },
+        "pose": {
+            "x": [],
+            "y": [],
+            "visible": [],
+        },
+    }
+
+    model = YoloPoses.from_results([result])
+    assert model.model_dump() == {
+        "cls": [0],
+        "name": ["person"],
+        "confidence": [0.9078],
+        "box": [{"coords": [102, 84, 183, 238], "title": "person"}],
+        "pose": [],
+    }
+
+
 def test_yolo_pose_from_results(running_img, running_img_masks):
     result = Results(
         orig_img=running_img,