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.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
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user