File size: 3,452 Bytes
aa975a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | # 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()
|