track_fighters_with_roboflow_kd

gistfile1.txt

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Optional

if "ORT_PROVIDERS" not in os.environ:
    os.environ["ORT_PROVIDERS"] = "CPUExecutionProvider"

import cv2
import numpy as np
import supervision as sv
from dotenv import load_dotenv
from inference.models.utils import get_roboflow_model

load_dotenv()

_ROBOFLOW_PROJECT = os.getenv("ROBOFLOW_PROJECT")
_ROBOFLOW_VERSION = os.getenv("ROBOFLOW_VERSION")
model = get_roboflow_model(
    model_id=(
        f"{_ROBOFLOW_PROJECT}/{_ROBOFLOW_VERSION}"
        if _ROBOFLOW_PROJECT and _ROBOFLOW_VERSION
        else None
    ),
    api_key=os.getenv("ROBOFLOW_API_KEY"),
)
tracker = sv.ByteTrack()
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()

DEFAULT_BELL_FPS = 30.0


def parse_bell_timecode(tc: str, fps: float = DEFAULT_BELL_FPS) -> float:
    """Парсит таймкод из bells в секунды. Формат: HH:MM:SS:FF."""
    parts = tc.strip().split(":")
    if len(parts) != 4:
        raise ValueError(f"Ожидается формат HH:MM:SS:FF, получено: {tc!r}")
    h, m, s, f = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
    return h * 3600 + m * 60 + s + f / fps


def fetch_bells(fight_id: str) -> Optional[list[dict]]:
    """Загружает таймкоды гонгов (начало/конец раундов) для боя из БД. fights.bells."""
    try:
        from src.database import get_db_cursor
        with get_db_cursor() as cur:
            cur.execute("SELECT bells FROM fights WHERE id = %s", (fight_id,))
            row = cur.fetchone()
        if row is None:
            return None
        bells = row.get("bells")
        if bells is None:
            return None
        if isinstance(bells, str):
            bells = json.loads(bells)
        return bells
    except Exception as e:
        print(f"❌ Не удалось загрузить bells для fight_id={fight_id}: {e}", file=sys.stderr)
        return None


def _box_iou(a: np.ndarray, b: np.ndarray) -> float:
    """IoU одного бокса a (4,) с одним боксом b (4,) xyxy."""
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    if ix2 <= ix1 or iy2 <= iy1:
        return 0.0
    inter = (ix2 - ix1) * (iy2 - iy1)
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    return inter / (area_a + area_b - inter + 1e-6)


def _sec_to_smpte_timecode(time_sec: float) -> str:
    """Переводит время в секундах в SMPTE-таймкод HH:MM:SS.mmm."""
    h = int(time_sec // 3600)
    m = int((time_sec % 3600) // 60)
    s = time_sec % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}"


def frame_to_records(key_points: sv.KeyPoints, detections: sv.Detections, frame_index: int, fps: float | None) -> list[dict]:
    """
    Собирает записи трекинга по одному кадру из sv.KeyPoints и Detections (после ByteTrack).
    Не зависит от infer_keypoints; после трекера detections — подмножество, сопоставление с key_points по IoU bbox.
    """
    records = []
    if detections.tracker_id is None or len(detections) == 0:
        return records

    time_sec = (frame_index / fps) if fps and fps > 0 else None
    timecode = _sec_to_smpte_timecode(time_sec) if time_sec is not None else None
    kp_boxes = key_points.as_detections().xyxy
    kp_objects = list(key_points)

    for i in range(len(detections)):
        bbox = detections.xyxy[i]
        tracker_id = int(detections.tracker_id[i])

        # Сопоставляем с key_point по максимальному IoU (трекер возвращает подмножество)
        best_j = 0
        best_iou = 0.0
        for j in range(len(kp_boxes)):
            iou = _box_iou(bbox, kp_boxes[j])
            if iou > best_iou:
                best_iou, best_j = iou, j

        kp_list = []
        if best_j < len(kp_objects):
            xy, conf, kp_class_id, _ = kp_objects[best_j]
            kp_class_id_arr = np.atleast_1d(kp_class_id) if kp_class_id is not None else None
            for j in range(xy.shape[0]):
                kp = {"x": float(xy[j, 0]), "y": float(xy[j, 1]), "confidence": float(conf[j]) if conf is not None else None}
                if kp_class_id_arr is not None and j < len(kp_class_id_arr):
                    kp["class_id"] = int(kp_class_id_arr[j])
                kp_list.append(kp)

        record = {
            "first_frame_index": frame_index,
            "tracker_id": tracker_id,
            "bbox": {"x_min": float(bbox[0]), "y_min": float(bbox[1]), "x_max": float(bbox[2]), "y_max": float(bbox[3])},
            "keypoints": kp_list,
        }
        if timecode is not None:
            record["timecode"] = timecode
        records.append(record)

    return records


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Трекинг боксёров на видео (pose + ByteTrack), вывод в JSONL."
    )
    parser.add_argument(
        "-i", "--input_video",
        type=str,
        required=True,
        help="Путь к входному видео.",
    )
    parser.add_argument(
        "-o", "--output_path",
        type=str,
        default=None,
        help="Путь к выходному JSONL (по умолчанию: директория входного видео, имя <input>_tracking.jsonl).",
    )
    parser.add_argument(
        "-f", "--fight_id",
        type=str,
        default=None,
        help="ID боя в БД (fights.id); по нему из fights.bells берётся начало 1-го раунда для вывода кадра с bbox и ByteTrack ID.",
    )
    args = parser.parse_args()

    input_path = Path(args.input_video)
    if not input_path.is_file():
        raise FileNotFoundError(f"Видео не найдено: {input_path}")

    if args.output_path is not None:
        output_path = Path(args.output_path)
    else:
        output_path = input_path.parent / f"{input_path.stem}_tracking.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    bells = fetch_bells(args.fight_id) if args.fight_id else None
    if args.fight_id and bells is None:
        print("⚠️ Bells для fight_id не найдены или запрос не удался; кадр первого раунда не будет показан.", file=sys.stderr)

    cap = cv2.VideoCapture(str(input_path))
    if not cap.isOpened():
        raise RuntimeError(f"Не удалось открыть видео: {input_path}")
    fps = cap.get(cv2.CAP_PROP_FPS) or None
    if fps is not None and fps <= 0:
        fps = None

    first_round_start_frame: Optional[int] = None
    if bells and fps:
        start_bell = bells[0].get("start_bell")
        if start_bell:
            start_sec = parse_bell_timecode(start_bell, fps)
            first_round_start_frame = int(round(start_sec * fps))

    tracker.reset()
    frame_index = 0
    first_round_frame_data: Optional[tuple[np.ndarray, sv.Detections]] = None

    with open(output_path, "w", encoding="utf-8") as f:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frame_bgr = frame
            results = model.infer(frame_bgr)[0]
            key_points = sv.KeyPoints.from_inference(results)
            detections = key_points.as_detections()
            detections = tracker.update_with_detections(detections)

            if first_round_start_frame is not None and frame_index == first_round_start_frame and len(detections) > 0:
                first_round_frame_data = (frame_bgr.copy(), detections)

            for record in frame_to_records(key_points, detections, frame_index, fps):
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

            frame_index += 1

    cap.release()
    print(f"Готово: {output_path} ({frame_index} кадров)")

    if first_round_frame_data is not None:
        img_bgr, det = first_round_frame_data
        labels = [f"boxer #{int(tid)}" for tid in det.tracker_id]
        annotated = box_annotator.annotate(scene=img_bgr.copy(), detections=det)
        annotated = label_annotator.annotate(scene=annotated, detections=det, labels=labels)
        img_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
        import matplotlib.pyplot as plt
        plt.figure(figsize=(12, 8))
        plt.imshow(img_rgb)
        plt.axis("off")
        plt.title(f"Первый раунд (кадр {first_round_start_frame}), ByteTrack ID → боец A/B")
        plt.tight_layout()
        plt.show()


if __name__ == "__main__":
    main()