50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
"""Tests for pet ID classifier."""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from vigilar.detection.pet_id import PetIDClassifier, PetIDResult
|
|
|
|
|
|
class TestPetIDResult:
|
|
def test_identified(self):
|
|
r = PetIDResult(pet_id="pet-1", pet_name="Angel", confidence=0.9)
|
|
assert r.is_identified
|
|
assert not r.is_low_confidence
|
|
|
|
def test_low_confidence(self):
|
|
r = PetIDResult(pet_id="pet-1", pet_name="Angel", confidence=0.6)
|
|
assert r.is_identified
|
|
assert r.is_low_confidence
|
|
|
|
def test_unknown(self):
|
|
r = PetIDResult(pet_id=None, pet_name=None, confidence=0.3)
|
|
assert not r.is_identified
|
|
|
|
|
|
class TestPetIDClassifier:
|
|
def test_not_loaded_returns_unknown(self):
|
|
classifier = PetIDClassifier(model_path="nonexistent.pt")
|
|
assert not classifier.is_loaded
|
|
crop = np.zeros((224, 224, 3), dtype=np.uint8)
|
|
result = classifier.identify(crop, species="cat")
|
|
assert not result.is_identified
|
|
|
|
def test_no_pets_registered_returns_unknown(self):
|
|
classifier = PetIDClassifier(model_path="nonexistent.pt")
|
|
assert classifier.pet_count == 0
|
|
|
|
def test_register_pet(self):
|
|
classifier = PetIDClassifier(model_path="nonexistent.pt")
|
|
classifier.register_pet("pet-1", "Angel", "cat")
|
|
classifier.register_pet("pet-2", "Milo", "dog")
|
|
assert classifier.pet_count == 2
|
|
|
|
def test_species_filter(self):
|
|
classifier = PetIDClassifier(model_path="nonexistent.pt")
|
|
classifier.register_pet("pet-1", "Angel", "cat")
|
|
classifier.register_pet("pet-2", "Taquito", "cat")
|
|
classifier.register_pet("pet-3", "Milo", "dog")
|
|
assert len(classifier.get_pets_by_species("cat")) == 2
|
|
assert len(classifier.get_pets_by_species("dog")) == 1
|