103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Offline subject tracking for one source time range (for cinematic auto-reframe).
|
|
|
|
Usage:
|
|
python subject_track.py <source_video> <start_seconds> <end_seconds> <samples>
|
|
|
|
Samples <samples> frames evenly across [start, end], runs a local YOLO detector on each, and reports
|
|
the NORMALISED centre (cx, cy in 0..1) of the dominant subject (largest person; else largest box of any
|
|
class). Prints one JSON object to stdout:
|
|
|
|
{"path": [{"t": <s>, "cx": <0..1>, "cy": <0..1>, "area": <0..1>}, ...]}
|
|
|
|
Local/offline only: ultralytics + OpenCV, weights from ./yolov8n.pt (no download). YOLOv8 is AGPL-3.0, so
|
|
this reframe path is non-commercial — matching the repo's existing CV stance. On any failure it prints
|
|
{"path": []} and exits 0 so the caller can fail soft to a centred crop.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) < 5:
|
|
print(json.dumps({"path": []}))
|
|
return 0
|
|
source, start_s, end_s, samples_s = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
|
|
try:
|
|
start = max(0.0, float(start_s))
|
|
end = float(end_s)
|
|
samples = max(1, int(samples_s))
|
|
if end <= start:
|
|
print(json.dumps({"path": []}))
|
|
return 0
|
|
|
|
os.environ.setdefault("YOLO_OFFLINE", "1")
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
weights = os.environ.get("SUBJECT_TRACK_WEIGHTS", "./yolov8n.pt")
|
|
model = YOLO(weights)
|
|
cap = cv2.VideoCapture(source)
|
|
if not cap.isOpened():
|
|
print(json.dumps({"path": []}))
|
|
return 0
|
|
|
|
path = []
|
|
step = (end - start) / samples
|
|
for i in range(samples):
|
|
t = start + step * (i + 0.5)
|
|
cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000.0)
|
|
ok, frame = cap.read()
|
|
if not ok or frame is None:
|
|
continue
|
|
h, w = frame.shape[:2]
|
|
result = model.predict(frame, verbose=False, device="cpu")[0]
|
|
best = _dominant_box(result)
|
|
if best is None:
|
|
continue
|
|
x1, y1, x2, y2 = best
|
|
cx = ((x1 + x2) / 2.0) / w
|
|
cy = ((y1 + y2) / 2.0) / h
|
|
area = ((x2 - x1) * (y2 - y1)) / float(w * h)
|
|
path.append({"t": round(t, 3), "cx": round(_clamp01(cx), 4),
|
|
"cy": round(_clamp01(cy), 4), "area": round(area, 4)})
|
|
cap.release()
|
|
print(json.dumps({"path": path}))
|
|
return 0
|
|
except Exception as exc: # noqa: BLE001 - fail soft, never break the render
|
|
sys.stderr.write("subject_track failed: %s\n" % exc)
|
|
print(json.dumps({"path": []}))
|
|
return 0
|
|
|
|
|
|
def _dominant_box(result):
|
|
"""Largest 'person' box; else the largest box of any class. Returns (x1,y1,x2,y2) or None."""
|
|
boxes = getattr(result, "boxes", None)
|
|
if boxes is None or len(boxes) == 0:
|
|
return None
|
|
names = result.names
|
|
best = None
|
|
best_area = -1.0
|
|
best_person = None
|
|
best_person_area = -1.0
|
|
for b in boxes:
|
|
xyxy = b.xyxy[0].tolist()
|
|
area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
|
|
cls_name = names.get(int(b.cls[0]), "") if isinstance(names, dict) else ""
|
|
if cls_name == "person" and area > best_person_area:
|
|
best_person_area = area
|
|
best_person = xyxy
|
|
if area > best_area:
|
|
best_area = area
|
|
best = xyxy
|
|
return best_person if best_person is not None else best
|
|
|
|
|
|
def _clamp01(v: float) -> float:
|
|
return 0.0 if v < 0 else (1.0 if v > 1 else v)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|