68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""Per-track image cropping, capture dirs, and full-frame event snapshots."""
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
import cv2
|
|
|
|
from core.detection import STATUSES
|
|
from core.geometry import clamp_bbox
|
|
from core.paths import CAPTURES_DIR, SNAPSHOTS_DIR
|
|
|
|
|
|
def setup_capture_dirs() -> Dict[str, Path]:
|
|
"""runtime/captures/{SAFE,PARTIAL,UNSAFE}/ — latest crop per track."""
|
|
dirs: Dict[str, Path] = {}
|
|
for s in STATUSES:
|
|
d = CAPTURES_DIR / s
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
dirs[s] = d
|
|
return dirs
|
|
|
|
|
|
def setup_snapshot_dirs() -> Dict[str, Path]:
|
|
"""runtime/snapshots/{SAFE,PARTIAL,UNSAFE}/ — annotated full frame per transition."""
|
|
dirs: Dict[str, Path] = {}
|
|
for s in STATUSES:
|
|
d = SNAPSHOTS_DIR / s
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
dirs[s] = d
|
|
return dirs
|
|
|
|
|
|
def save_track_image(frame, track, capture_dirs: Dict[str, Path]) -> Optional[Path]:
|
|
"""Save the latest crop for a track. Overwritten each frame."""
|
|
h, w = frame.shape[:2]
|
|
x1, y1, x2, y2 = clamp_bbox(track.bbox, w, h)
|
|
if x2 <= x1 or y2 <= y1:
|
|
return None
|
|
crop = frame[y1:y2, x1:x2]
|
|
if crop.size == 0:
|
|
return None
|
|
|
|
target = capture_dirs[track.status] / f"track_{track.track_id:04d}.jpg"
|
|
if track.photo_path and track.photo_path != target and track.photo_path.exists():
|
|
try:
|
|
track.photo_path.unlink()
|
|
except OSError:
|
|
pass
|
|
|
|
cv2.imwrite(str(target), crop)
|
|
track.photo_path = target
|
|
return target
|
|
|
|
|
|
def save_event_snapshot(annotated_frame, track, snapshot_dirs: Dict[str, Path]) -> Optional[Path]:
|
|
"""Save the full annotated frame at the moment of a NEW / STATUS_CHANGE event.
|
|
|
|
Timestamped filename so a history is preserved across events.
|
|
"""
|
|
if annotated_frame is None or annotated_frame.size == 0:
|
|
return None
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
|
target = snapshot_dirs[track.status] / f"track_{track.track_id:04d}_{ts}.jpg"
|
|
cv2.imwrite(str(target), annotated_frame)
|
|
return target
|