53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
"""Tests for YOLOv8 unified detector."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from vigilar.detection.person import Detection
|
|
from vigilar.detection.yolo import YOLODetector, ANIMAL_CLASSES, WILDLIFE_CLASSES
|
|
|
|
|
|
class TestYOLOConstants:
|
|
def test_animal_classes(self):
|
|
assert "cat" in ANIMAL_CLASSES
|
|
assert "dog" in ANIMAL_CLASSES
|
|
|
|
def test_wildlife_classes(self):
|
|
assert "bear" in WILDLIFE_CLASSES
|
|
assert "bird" in WILDLIFE_CLASSES
|
|
|
|
def test_no_overlap_animal_wildlife(self):
|
|
assert not ANIMAL_CLASSES.intersection(WILDLIFE_CLASSES)
|
|
|
|
|
|
class TestYOLODetector:
|
|
def test_initializes_without_model(self):
|
|
detector = YOLODetector(model_path="nonexistent.pt", confidence_threshold=0.5)
|
|
assert not detector.is_loaded
|
|
|
|
def test_detect_returns_empty_when_not_loaded(self):
|
|
detector = YOLODetector(model_path="nonexistent.pt")
|
|
frame = np.zeros((480, 640, 3), dtype=np.uint8)
|
|
detections = detector.detect(frame)
|
|
assert detections == []
|
|
|
|
def test_classify_detection_person(self):
|
|
d = Detection(class_name="person", class_id=0, confidence=0.9, bbox=(10, 20, 100, 200))
|
|
assert YOLODetector.classify(d) == "person"
|
|
|
|
def test_classify_detection_vehicle(self):
|
|
d = Detection(class_name="car", class_id=2, confidence=0.85, bbox=(10, 20, 100, 200))
|
|
assert YOLODetector.classify(d) == "vehicle"
|
|
|
|
def test_classify_detection_domestic_animal(self):
|
|
d = Detection(class_name="cat", class_id=15, confidence=0.9, bbox=(10, 20, 100, 200))
|
|
assert YOLODetector.classify(d) == "domestic_animal"
|
|
|
|
def test_classify_detection_wildlife(self):
|
|
d = Detection(class_name="bear", class_id=21, confidence=0.8, bbox=(10, 20, 100, 200))
|
|
assert YOLODetector.classify(d) == "wildlife"
|
|
|
|
def test_classify_detection_other(self):
|
|
d = Detection(class_name="chair", class_id=56, confidence=0.7, bbox=(10, 20, 100, 200))
|
|
assert YOLODetector.classify(d) == "other"
|