Integrate YOLOv8 detection and pet ID into camera worker
This commit is contained in:
parent
d0acf7703c
commit
4c9ebe029d
@ -22,6 +22,9 @@ from vigilar.camera.recorder import AdaptiveRecorder
|
||||
from vigilar.camera.ring_buffer import RingBuffer
|
||||
from vigilar.config import CameraConfig, MQTTConfig, RemoteConfig
|
||||
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__)
|
||||
|
||||
@ -46,6 +49,7 @@ def run_camera_worker(
|
||||
recordings_dir: str,
|
||||
hls_dir: str,
|
||||
remote_cfg: RemoteConfig | None = None,
|
||||
pets_cfg: "PetsConfig | None" = None,
|
||||
) -> None:
|
||||
"""Main entry point for a camera worker process."""
|
||||
camera_id = camera_cfg.id
|
||||
@ -107,6 +111,21 @@ def run_camera_worker(
|
||||
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()
|
||||
shutdown = False
|
||||
|
||||
@ -243,6 +262,49 @@ def run_camera_worker(
|
||||
if state.frame_count % idle_skip_factor == 0:
|
||||
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
|
||||
if now - last_heartbeat >= 10:
|
||||
last_heartbeat = now
|
||||
|
||||
Loading…
Reference in New Issue
Block a user