initial commit

tests/test_video_smoother.py

@@ -0,0 +1,38 @@
+import unittest
+
+from app.video_pipeline import VideoResultSmoother
+
+
+def result(class_id, class_name, confidence, box):
+    return {
+        "bbox_xyxy": box,
+        "cls_class_id": class_id,
+        "cls_class_name": class_name,
+        "cls_confidence": confidence,
+        "show_cls_name": True,
+    }
+
+
+class VideoResultSmootherTests(unittest.TestCase):
+    def test_one_low_confidence_label_does_not_flip_track(self):
+        smoother = VideoResultSmoother(hold_frames=4, alpha=0.5)
+        smoother.update([result(1, "白鹤", 0.95, [0, 0, 20, 20])], 0)
+        second = smoother.update(
+            [result(2, "白枕鹤", 0.20, [1, 1, 21, 21])],
+            1,
+        )
+        self.assertEqual(second[0]["cls_class_name"], "白鹤")
+        self.assertNotIn("_smoother_class_votes", second[0])
+
+    def test_matching_does_not_require_same_predicted_class(self):
+        smoother = VideoResultSmoother(hold_frames=4, alpha=0.5)
+        smoother.update([result(1, "白鹤", 0.9, [0, 0, 20, 20])], 0)
+        second = smoother.update(
+            [result(2, "白枕鹤", 0.9, [2, 2, 22, 22])],
+            1,
+        )
+        self.assertEqual(second[0]["bbox_xyxy"], [1.0, 1.0, 21.0, 21.0])
+
+
+if __name__ == "__main__":
+    unittest.main()