Integrate YOLOv8 detection and pet ID into camera worker

This commit is contained in:
Aaron D. Lee 2026-04-03 13:22:25 -04:00
parent d0acf7703c
commit 4c9ebe029d

View File

@ -22,6 +22,9 @@ from vigilar.camera.recorder import AdaptiveRecorder
from vigilar.camera.ring_buffer import RingBuffer from vigilar.camera.ring_buffer import RingBuffer
from vigilar.config import CameraConfig, MQTTConfig, RemoteConfig from vigilar.config import CameraConfig, MQTTConfig, RemoteConfig
from vigilar.constants import Topics from vigilar.constants import Topics
from vigilar.detection.yolo import YOLODetector
from vigilar.detection.pet_id import PetIDClassifier
from vigilar.detection.wildlife import classify_wildlife_threat
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -46,6 +49,7 @@ def run_camera_worker(
recordings_dir: str, recordings_dir: str,
hls_dir: str, hls_dir: str,
remote_cfg: RemoteConfig | None = None, remote_cfg: RemoteConfig | None = None,
pets_cfg: "PetsConfig | None" = None,
) -> None: ) -> None:
"""Main entry point for a camera worker process.""" """Main entry point for a camera worker process."""
camera_id = camera_cfg.id camera_id = camera_cfg.id
@ -107,6 +111,21 @@ def run_camera_worker(
bitrate_kbps=remote_cfg.remote_hls_bitrate_kbps, bitrate_kbps=remote_cfg.remote_hls_bitrate_kbps,
) )
# Object detection (YOLOv8 unified detector)
yolo_detector = None
pet_classifier = None
if pets_cfg and pets_cfg.enabled:
yolo_detector = YOLODetector(
model_path=pets_cfg.model_path,
confidence_threshold=pets_cfg.confidence_threshold,
)
if pets_cfg.pet_id_enabled:
pet_classifier = PetIDClassifier(
model_path=pets_cfg.pet_id_model_path,
high_threshold=pets_cfg.pet_id_threshold,
low_threshold=pets_cfg.pet_id_low_confidence,
)
state = CameraState() state = CameraState()
shutdown = False shutdown = False
@ -243,6 +262,49 @@ def run_camera_worker(
if state.frame_count % idle_skip_factor == 0: if state.frame_count % idle_skip_factor == 0:
recorder.write_frame(frame) recorder.write_frame(frame)
# Run object detection on motion frames
if state.motion_active and yolo_detector and yolo_detector.is_loaded:
detections = yolo_detector.detect(frame)
for det in detections:
category = YOLODetector.classify(det)
if category == "domestic_animal":
# Crop for pet ID
x, y, w, h = det.bbox
crop = frame[max(0, y):y + h, max(0, x):x + w]
pet_result = None
if pet_classifier and pet_classifier.is_loaded and crop.size > 0:
pet_result = pet_classifier.identify(crop, species=det.class_name)
payload = {
"species": det.class_name,
"confidence": round(det.confidence, 3),
"camera_location": camera_cfg.location,
}
if pet_result and pet_result.is_identified:
payload["pet_id"] = pet_result.pet_id
payload["pet_name"] = pet_result.pet_name
payload["pet_confidence"] = round(pet_result.confidence, 3)
bus.publish_event(
Topics.pet_location(pet_result.pet_name.lower()),
camera_id=camera_id,
camera_location=camera_cfg.location,
)
bus.publish_event(Topics.camera_pet_detected(camera_id), **payload)
elif category == "wildlife" and pets_cfg:
frame_area = frame.shape[0] * frame.shape[1]
threat_level, species = classify_wildlife_threat(
det, pets_cfg.wildlife, frame_area,
)
bus.publish_event(
Topics.camera_wildlife_detected(camera_id),
species=species,
threat_level=threat_level,
confidence=round(det.confidence, 3),
camera_location=camera_cfg.location,
)
# Heartbeat every 10 seconds # Heartbeat every 10 seconds
if now - last_heartbeat >= 10: if now - last_heartbeat >= 10:
last_heartbeat = now last_heartbeat = now