39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
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()
|