42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
"""Per-track image cropping + capture directory setup."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
import cv2
|
|
|
|
from saqr.core.detection import STATUSES
|
|
from saqr.core.geometry import clamp_bbox
|
|
from saqr.core.paths import CAPTURES_DIR
|
|
|
|
|
|
def setup_capture_dirs() -> Dict[str, Path]:
|
|
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 save_track_image(frame, track, capture_dirs: Dict[str, Path]) -> Optional[Path]:
|
|
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
|