You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import time
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from helpers import filter_yolo_boxes_with_scores
|
|
from yolo_worker import INFERENCE_DEVICE, INFERENCE_HALF, INFERENCE_SIZE_FULL, YOLOWorker, fixed_letterbox, raw_yolo_boxes
|
|
|
|
|
|
class FailingModel:
|
|
def __call__(self, *args, **kwargs):
|
|
raise RuntimeError("test failure")
|
|
|
|
|
|
class Boxes:
|
|
xyxy = torch.tensor([[0.0, 160.0, 640.0, 480.0]])
|
|
conf = torch.tensor([0.9])
|
|
cls = torch.tensor([0.0])
|
|
|
|
def __len__(self):
|
|
return 1
|
|
|
|
|
|
class Result:
|
|
boxes = Boxes()
|
|
|
|
|
|
class YOLOWorkerLifecycleTests(unittest.TestCase):
|
|
def test_cpu_inference_uses_cpu_precision_and_workable_size(self):
|
|
if torch.cuda.is_available():
|
|
self.skipTest("CPU fallback is inactive on CUDA hosts")
|
|
self.assertEqual(INFERENCE_DEVICE, "cpu")
|
|
self.assertFalse(INFERENCE_HALF)
|
|
self.assertLessEqual(INFERENCE_SIZE_FULL, 640)
|
|
|
|
def test_fixed_letterbox_maps_boxes_back_without_distortion(self):
|
|
image = np.zeros((100, 200, 3), dtype=np.uint8)
|
|
boxed, scale, pad_x, pad_y = fixed_letterbox(image, 640)
|
|
self.assertEqual(boxed.shape, (640, 640, 3))
|
|
raw = raw_yolo_boxes(Result(), scale=1.0 / scale, pad_x=pad_x, pad_y=pad_y)
|
|
np.testing.assert_allclose(raw[0][:4], [0.0, 0.0, 200.0, 100.0])
|
|
filtered = filter_yolo_boxes_with_scores(
|
|
Result(),
|
|
frame_w=200,
|
|
frame_h=100,
|
|
min_conf=0.1,
|
|
input_scale=scale,
|
|
pad_x=pad_x,
|
|
pad_y=pad_y,
|
|
content_w=200,
|
|
content_h=100,
|
|
)
|
|
np.testing.assert_allclose(filtered[0][:4], [0.0, 0.0, 200.0, 100.0])
|
|
|
|
def test_worker_reports_failed_inference_without_hanging(self):
|
|
worker = YOLOWorker(FailingModel())
|
|
with patch("yolo_worker.torch.cuda.is_available", return_value=False):
|
|
worker.start()
|
|
try:
|
|
worker.submit(np.zeros((16, 16, 3), dtype=np.uint8), None, "FULL", 1.0)
|
|
deadline = time.monotonic() + 2.0
|
|
result = None
|
|
while result is None and time.monotonic() < deadline:
|
|
result = worker.try_get()
|
|
time.sleep(0.01)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result[0], [])
|
|
self.assertEqual(result[1], 1.0)
|
|
finally:
|
|
worker.stop()
|
|
self.assertFalse(worker.thread.is_alive())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|