impose_skeletons

gistfile1.txt

"""
Навешивает на кадры видео keypoints и рёбра скелета из результатов инференса.
Цвет (красный/синий) задаётся через bbox_indices по object_index на опорном кадре.
Входные данные читаются из CSV или JSONL (например [fight_id]_cleaned_keypoints.csv).
"""

import argparse
import csv
import json
import shutil
import subprocess
import tempfile
from pathlib import Path

import cv2
import numpy as np

DEBUG_LOG_PATH = Path(__file__).resolve().parent.parent / ".cursor" / "debug-b383c2.log"


def _debug_log(location: str, hypothesis_id: str, data: dict) -> None:
    try:
        payload = {"sessionId": "b383c2", "hypothesisId": hypothesis_id, "location": location, "message": "impose_bboxes", "data": data, "timestamp": __import__("time").time() * 1000}
        with open(DEBUG_LOG_PATH, "a", encoding="utf-8") as f:
            f.write(json.dumps(payload, ensure_ascii=False) + "\n")
    except Exception:
        pass


# BGR
COLOR_RED = (0, 0, 255)
COLOR_BLUE = (255, 0, 0)
COLOR_GRAY = (128, 128, 128)

# Рёбра скелета: пары (body_part_a, body_part_b) для отрезков между ключевыми точками.
SKELETON_EDGES: list[tuple[str, str]] = [
    ("crown", "atlas"),
    ("crown", "nose"),
    ("nose", "jaw"),
    ("atlas", "left-shoulder"),
    ("atlas", "right-shoulder"),
    ("left-shoulder", "right-shoulder"),
    ("left-shoulder", "left-elbow"),
    ("left-shoulder", "left-hip-bone"),
    ("left-elbow", "left-fist"),
    ("right-shoulder", "right-elbow"),
    ("right-shoulder", "right-hip-bone"),
    ("right-elbow", "right-fist"),
    ("left-hip-bone", "right-hip-bone"),
    ("left-hip-bone", "left-knee"),
    ("right-hip-bone", "right-knee"),
    ("left-knee", "left-heel"),
    ("left-heel", "left-forefoot"),
    ("right-knee", "right-heel"),
    ("right-heel", "right-forefoot"),
]

# Фиксированные столбцы CSV (не ключевые точки)
CSV_META_COLUMNS = {
    "first_frame_index",
    "object_index",
    "trunks",
    "class_id",
    "confidence",
    "timecode",
    "bbox_x_min",
    "bbox_y_min",
    "bbox_x_max",
    "bbox_y_max",
}


def load_bbox_indices(path: Path) -> dict:
    """Загружает bbox_indices.json: first_frame_index, red_fighter_first_frame_idx, blue_fighter_first_frame_idx."""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    for key in ("first_frame_index", "red_fighter_first_frame_idx", "blue_fighter_first_frame_idx"):
        if key not in data:
            raise SystemExit(f"bbox_indices должен содержать ключ '{key}'.")
    return data


def _parse_float(value: str | None) -> float | None:
    if value is None or (isinstance(value, str) and value.strip() == ""):
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        return None


def _parse_int(value: str | None) -> int | None:
    v = _parse_float(value)
    return int(v) if v is not None else None


def _keypoints_from_csv_row(row: dict) -> list[dict]:
    """Собирает список ключевых точек из столбцов CSV (каждый столбец — часть тела, значение — JSON-массив)."""
    keypoints: list[dict] = []
    for col, raw in row.items():
        if col in CSV_META_COLUMNS or not raw:
            continue
        if isinstance(raw, str):
            raw = raw.strip()
            if not raw:
                continue
            try:
                arr = json.loads(raw)
            except (json.JSONDecodeError, TypeError):
                continue
        else:
            arr = raw
        if not isinstance(arr, list) or not arr:
            continue
        for item in arr:
            if not isinstance(item, dict):
                continue
            x = item.get("x")
            y = item.get("y")
            if x is None or y is None:
                continue
            keypoints.append({
                "x": float(x),
                "y": float(y),
                "confidence": item.get("confidence", 0.0),
                "body_part": col,
            })
    return keypoints


def load_bboxes_by_frame_from_csv(csv_path: Path) -> dict[int, list[dict]]:
    """Группирует записи CSV по first_frame_index. Каждая запись: object_index, bbox, keypoints (если есть столбцы)."""
    frames: dict[int, list[dict]] = {}
    with open(csv_path, "r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames:
            return frames
        # #region agent log
        _debug_log("impose_bboxes.py:load_csv", "H2", {"fieldnames": list(reader.fieldnames)[:15], "path_suffix": csv_path.suffix})
        # #endregion
        for row in reader:
            fi = _parse_int(row.get("first_frame_index"))
            if fi is None:
                continue
            x_min = _parse_float(row.get("bbox_x_min"))
            y_min = _parse_float(row.get("bbox_y_min"))
            x_max = _parse_float(row.get("bbox_x_max"))
            y_max = _parse_float(row.get("bbox_y_max"))
            if x_min is None or y_min is None or x_max is None or y_max is None:
                continue
            obj_idx = _parse_int(row.get("object_index"))
            record: dict = {
                "object_index": obj_idx if obj_idx is not None else -1,
                "bbox": {
                    "x_min": x_min,
                    "y_min": y_min,
                    "x_max": x_max,
                    "y_max": y_max,
                },
            }
            kps = _keypoints_from_csv_row(row)
            if kps:
                record["keypoints"] = kps
            if fi not in frames:
                frames[fi] = []
            frames[fi].append(record)
    return frames


def load_bboxes_by_frame_from_jsonl(jsonl_path: Path) -> dict[int, list[dict]]:
    """
    Группирует записи JSONL по frame_index (или first_frame_index).

    Поддерживает 2 формата:
    1) Объектный (macOS): каждая строка уже содержит object_index, bbox{x_min..}, keypoints.
    2) Кадровый (Windows): каждая строка содержит frame_index и массив predictions, где bbox задан как x,y,width,height.
    """
    frames: dict[int, list[dict]] = {}
    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except (json.JSONDecodeError, TypeError):
                continue
            fi = _parse_int(rec.get("first_frame_index") or rec.get("frame_index"))
            if fi is None:
                continue
            records_for_frame: list[dict] = []

            # Формат 1: одна строка = один объект (уже с bbox в x_min..y_max)
            bbox_raw = rec.get("bbox")
            if isinstance(bbox_raw, dict):
                x_min = _parse_float(bbox_raw.get("x_min"))
                y_min = _parse_float(bbox_raw.get("y_min"))
                x_max = _parse_float(bbox_raw.get("x_max"))
                y_max = _parse_float(bbox_raw.get("y_max"))
                if x_min is not None and y_min is not None and x_max is not None and y_max is not None:
                    obj_idx = _parse_int(rec.get("object_index"))
                    record: dict = {
                        "object_index": obj_idx if obj_idx is not None else -1,
                        "bbox": {"x_min": x_min, "y_min": y_min, "x_max": x_max, "y_max": y_max},
                    }
                    kps_raw = rec.get("keypoints")
                    if isinstance(kps_raw, list):
                        kps = []
                        for item in kps_raw:
                            if not isinstance(item, dict):
                                continue
                            x, y = item.get("x"), item.get("y")
                            if x is None or y is None:
                                continue
                            kps.append({
                                "x": float(x),
                                "y": float(y),
                                "confidence": item.get("confidence", 0.0),
                                "body_part": item.get("body_part", ""),
                            })
                        if kps:
                            record["keypoints"] = kps
                    records_for_frame.append(record)

            # Формат 2: одна строка = кадр, объекты в predictions (bbox как center-size)
            preds = rec.get("predictions")
            if isinstance(preds, list):
                for obj_idx, pred in enumerate(preds):
                    if not isinstance(pred, dict):
                        continue
                    cx = _parse_float(pred.get("x"))
                    cy = _parse_float(pred.get("y"))
                    bw = _parse_float(pred.get("width"))
                    bh = _parse_float(pred.get("height"))
                    if cx is None or cy is None or bw is None or bh is None:
                        continue
                    x_min = cx - bw / 2.0
                    y_min = cy - bh / 2.0
                    x_max = cx + bw / 2.0
                    y_max = cy + bh / 2.0
                    record: dict = {
                        "object_index": obj_idx,
                        "bbox": {"x_min": x_min, "y_min": y_min, "x_max": x_max, "y_max": y_max},
                    }
                    kps_raw = pred.get("keypoints")
                    if isinstance(kps_raw, list):
                        kps = []
                        for item in kps_raw:
                            if not isinstance(item, dict):
                                continue
                            x, y = item.get("x"), item.get("y")
                            if x is None or y is None:
                                continue
                            kps.append({
                                "x": float(x),
                                "y": float(y),
                                "confidence": item.get("confidence", 0.0),
                                "body_part": item.get("body_part") or item.get("class") or "",
                            })
                        if kps:
                            record["keypoints"] = kps
                    records_for_frame.append(record)

            if records_for_frame:
                if fi not in frames:
                    frames[fi] = []
                frames[fi].extend(records_for_frame)
    return frames


def bbox_color(
    object_index: int,
    red_idx: int,
    blue_idx: int,
) -> tuple[int, int, int]:
    """Возвращает BGR-цвет для bbox по object_index и привязке красный/синий угол."""
    if object_index == red_idx:
        return COLOR_RED
    if object_index == blue_idx:
        return COLOR_BLUE
    return COLOR_GRAY


def _keypoints_by_body_part(
    keypoints: list[dict],
    min_confidence: float,
) -> dict[str, tuple[int, int]]:
    """Строит словарь body_part -> (x, y) для точек с confidence >= min_confidence."""
    out: dict[str, tuple[int, int]] = {}
    for kp in keypoints:
        if kp.get("confidence") is None or float(kp.get("confidence", 0)) < min_confidence:
            continue
        part = (kp.get("body_part") or kp.get("class") or "").strip()
        if not part:
            continue
        x = int(round(kp.get("x", 0)))
        y = int(round(kp.get("y", 0)))
        out[part] = (x, y)
    return out


def draw_skeleton_edges(
    frame: np.ndarray,
    keypoints: list[dict],
    color: tuple[int, int, int],
    min_confidence: float = 0.1,
    thickness: int = 2,
) -> None:
    """Рисует рёбра скелета (отрезки между ключевыми точками) по SKELETON_EDGES."""
    by_part = _keypoints_by_body_part(keypoints, min_confidence)
    for a, b in SKELETON_EDGES:
        if a not in by_part or b not in by_part:
            continue
        pt1 = by_part[a]
        pt2 = by_part[b]
        cv2.line(frame, pt1, pt2, color, thickness, cv2.LINE_AA)


def draw_keypoints(
    frame: np.ndarray,
    keypoints: list[dict],
    color: tuple[int, int, int],
    min_confidence: float = 0.1,
    radius: int = 2,
) -> None:
    """Рисует ключевые точки на кадре (кружки)."""
    for kp in keypoints:
        conf = kp.get("confidence", 0)
        if conf is None or conf < min_confidence:
            continue
        x = int(round(kp.get("x", 0)))
        y = int(round(kp.get("y", 0)))
        cv2.circle(frame, (x, y), radius, color, -1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Накладывает keypoints и рёбра скелета на видео по привязке red/blue из bbox_indices."
    )
    parser.add_argument(
        "-i", "--input_video",
        type=Path,
        required=True,
        help="Входное видео (.mp4).",
    )
    parser.add_argument(
        "--input_bboxes",
        type=Path,
        required=True,
        help="CSV или JSONL с bbox'ами: [fight_id]_cleaned_keypoints.csv или [fight_id]_bboxes_keypoints.jsonl.",
    )
    parser.add_argument(
        "-b", "--bbox_indices",
        type=Path,
        required=True,
        help="Ссылка на JSON-файл [fight_id]_bbox_indices.json: first_frame_index, red_fighter_first_frame_idx, blue_fighter_first_frame_idx.",
    )
    parser.add_argument(
        "--output_video",
        type=Path,
        default=None,
        help="Выходное видео с наложенными keypoints и скелетом. По умолчанию: рядом с входным, суффикс _no_tracking.mp4.",
    )
    parser.add_argument(
        "--no_keypoints",
        action="store_true",
        help="Не накладывать ключевые точки и рёбра скелета.",
    )
    parser.add_argument(
        "--keypoint_confidence",
        type=float,
        default=0.2,
        help="Минимальный confidence для отрисовки ключевой точки (по умолчанию 0.2).",
    )
    args = parser.parse_args()

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

    input_bboxes = args.input_bboxes.resolve()
    if not input_bboxes.is_file():
        raise SystemExit(f"Файл bbox'ов не найден: {input_bboxes}")

    # #region agent log
    _debug_log("impose_bboxes.py:main", "H1", {"input_bboxes_suffix": input_bboxes.suffix, "input_bboxes_name": input_bboxes.name})
    # #endregion

    bbox_indices_path = args.bbox_indices.resolve()
    if not bbox_indices_path.is_file():
        raise SystemExit(f"Файл bbox_indices не найден: {bbox_indices_path}")

    indices = load_bbox_indices(bbox_indices_path)
    red_idx = indices["red_fighter_first_frame_idx"]
    blue_idx = indices["blue_fighter_first_frame_idx"]

    if input_bboxes.suffix.lower() == ".jsonl":
        bboxes_by_frame = load_bboxes_by_frame_from_jsonl(input_bboxes)
    else:
        bboxes_by_frame = load_bboxes_by_frame_from_csv(input_bboxes)

    # #region agent log
    _debug_log("impose_bboxes.py:main", "H1", {"num_frames": len(bboxes_by_frame), "sample_frame_keys": list(bboxes_by_frame.keys())[:5], "runId": "post-fix"})
    # #endregion

    if not bboxes_by_frame:
        raise SystemExit(
            "Нет записей с bbox'ами. Для CSV проверьте столбцы first_frame_index, bbox_x_min/y_min/x_max/y_max. "
            "Для JSONL — наличие frame_index (или first_frame_index), bbox с x_min/y_min/x_max/y_max."
        )

    # CSV может использовать 1-based номера кадров (1, 2, 3, …); видео в OpenCV — 0-based (0, 1, 2, …).
    frame_keys = set(bboxes_by_frame.keys())
    csv_frame_is_1_based = bool(frame_keys and min(frame_keys) == 1 and 0 not in frame_keys)
    if csv_frame_is_1_based:
        print("Обнаружена 1-based нумерация кадров в CSV — сопоставление с видео скорректировано.")

    if args.output_video is not None:
        output_video = args.output_video.resolve()
    else:
        output_video = input_video.parent / f"{input_video.stem}_no_tracking.mp4"

    cap = cv2.VideoCapture(str(input_video))
    if not cap.isOpened():
        raise SystemExit(f"Не удалось открыть видео: {input_video}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        temp_path = tmp.name
    try:
        fourcc = cv2.VideoWriter_fourcc("m", "p", "4", "v")
        out = cv2.VideoWriter(temp_path, fourcc, fps, (w, h))
        if not out.isOpened():
            cap.release()
            Path(temp_path).unlink(missing_ok=True)
            raise SystemExit(f"Не удалось создать временное видео: {temp_path}")

        frame_index = 0
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            lookup_key = (frame_index + 1) if csv_frame_is_1_based else frame_index
            records = bboxes_by_frame.get(lookup_key, [])
            for rec in records:
                obj_idx = rec.get("object_index", -1)
                color = bbox_color(obj_idx, red_idx, blue_idx)
                if not args.no_keypoints:
                    kps = rec.get("keypoints") or []
                    draw_skeleton_edges(frame, kps, color, min_confidence=args.keypoint_confidence)
                    draw_keypoints(frame, kps, color, min_confidence=args.keypoint_confidence)
            out.write(frame)
            frame_index += 1

        cap.release()
        out.release()

        # Подмешиваем аудио из исходного видео через ffmpeg
        ffmpeg_cmd = [
            "ffmpeg", "-y",
            "-i", str(input_video),
            "-i", temp_path,
            "-map", "0:a",
            "-map", "1:v",
            "-c:a", "copy",
            "-c:v", "libx264",
            "-shortest",
            str(output_video),
        ]
        try:
            result = subprocess.run(ffmpeg_cmd, capture_output=True, text=True)
            if result.returncode != 0:
                shutil.copy2(temp_path, output_video)
                print("В исходном видео нет аудио или ffmpeg недоступен — сохранено только видео.")
        except FileNotFoundError:
            shutil.copy2(temp_path, output_video)
            print("ffmpeg не найден в PATH — сохранено только видео без аудио.")
        Path(temp_path).unlink(missing_ok=True)
    except Exception:
        Path(temp_path).unlink(missing_ok=True)
        raise

    print(f"Записано кадров: {frame_index}, выход: {output_video}")


if __name__ == "__main__":
    main()