model/video_utils.py

# video_utils.py
# -*- coding: utf-8 -*-
"""
视频可视化：
- 原视频检查：采样帧上画红色矩形框（线宽可见）
- 黑底 mask：采样帧填充白色 bbox 区域，其它帧纯黑
- 编解码兜底：mp4v 打不开则自动降级到 XVID + .avi
"""

import cv2
import numpy as np
from typing import List, Dict

RED   = (0, 0, 255)     # BGR
WHITE = (255, 255, 255)

def _to_xyxy(bbox, W, H):
    x1, y1, x2, y2 = bbox
    return int(round(x1 * W)), int(round(y1 * H)), int(round(x2 * W)), int(round(y2 * H))

def _make_writer(path: str, W: int, H: int, fps: float):
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    vw = cv2.VideoWriter(path, fourcc, fps, (W, H))
    if not vw.isOpened():
        # fallback
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        path = path.rsplit('.', 1)[0] + '.avi'
        vw = cv2.VideoWriter(path, fourcc, fps, (W, H))
    return vw, path

def draw_inspection_video(src_path: str, out_path: str, sampled_idx: List[int],
                          tasks_json: Dict[str, List[Dict]], thickness: int = 4) -> None:
    cap = cv2.VideoCapture(src_path)
    if not cap.isOpened():
        raise RuntimeError(f"cannot open video: {src_path}")
    W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0

    vw, out_path = _make_writer(out_path, W, H, fps)

    # [帧号 -> 多任务 bbox 列表]（按采样顺序对齐：seq 下标 k 对应 sampled_idx[k]）
    bbox_by_frame = {}
    for _, seq in tasks_json.items():
        for k, e in enumerate(seq):
            if k >= len(sampled_idx):
                break
            bb = e.get("bbox")
            if bb is None:
                continue
            frm = sampled_idx[k]
            bbox_by_frame.setdefault(frm, []).append(bb)

    cur = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if cur in bbox_by_frame:
            for bb in bbox_by_frame[cur]:
                x1, y1, x2, y2 = _to_xyxy(bb, W, H)
                cv2.rectangle(frame, (x1, y1), (x2, y2), RED, thickness)
        vw.write(frame)
        cur += 1

    cap.release()
    vw.release()

def draw_mask_video(src_path: str, out_path: str, sampled_idx: List[int],
                    tasks_json: Dict[str, List[Dict]]) -> None:
    cap = cv2.VideoCapture(src_path)
    if not cap.isOpened():
        raise RuntimeError(f"cannot open video: {src_path}")
    W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    vw, out_path = _make_writer(out_path, W, H, fps)
    black = np.zeros((H, W, 3), dtype=np.uint8)

    bbox_by_frame = {}
    for _, seq in tasks_json.items():
        for k, e in enumerate(seq):
            if k >= len(sampled_idx):
                break
            bb = e.get("bbox")
            if bb is None:
                continue
            frm = sampled_idx[k]
            bbox_by_frame.setdefault(frm, []).append(bb)

    for cur in range(total):
        frame = black.copy()
        if cur in bbox_by_frame:
            for bb in bbox_by_frame[cur]:
                x1, y1, x2, y2 = _to_xyxy(bb, W, H)
                # 白色填充作为 mask
                cv2.rectangle(frame, (x1, y1), (x2, y2), WHITE, thickness=-1)
        vw.write(frame)

    cap.release()
    vw.release()