wangdongxu

jira:NYJ-1550 desc:m3u8 images

  1 +import base64
  2 +import hashlib
  3 +import subprocess
  4 +from pathlib import Path
  5 +
  6 +import cv2
  7 +
  8 +try:
  9 + from .m3u8_to_mp4_sei import (
  10 + download_m3u8_to_mp4,
  11 + extract_h264_es_from_mp4,
  12 + parse_h264_sei,
  13 + extract_utc_from_sei,
  14 + )
  15 +except ImportError:
  16 + from m3u8_to_mp4_sei import (
  17 + download_m3u8_to_mp4,
  18 + extract_h264_es_from_mp4,
  19 + parse_h264_sei,
  20 + extract_utc_from_sei,
  21 + )
  22 +
  23 +# 用于区分传入的 start_utc/end_utc 是绝对 UTC 时间戳还是秒偏移
  24 +_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000
  25 +
  26 +
  27 +class Video2Frame:
  28 + def __init__(self, cache_dir):
  29 + self.cache_dir = Path(cache_dir)
  30 + self.video_dir = self.cache_dir / "videos"
  31 + self.frames_dir = self.cache_dir / "frames"
  32 + self.video_dir.mkdir(parents=True, exist_ok=True)
  33 + self.frames_dir.mkdir(parents=True, exist_ok=True)
  34 +
  35 + def to_frames(self, url, start_utc, end_utc, fps, roi=None, max_px_area=None) -> list:
  36 + """
  37 + 下载 m3u8 视频流的指定片段并提取帧,保存到 cache_dir 下。
  38 +
  39 + - video_dir 中保存的是 start_utc ~ end_utc 截取后的视频片段(而非完整视频)。
  40 + - 如果视频中包含 SEI UTC 信息且传入的 start_utc/end_utc 为绝对时间戳,
  41 + 则会基于 SEI UTC 进行帧级定位;否则将 start_utc/end_utc 视为秒偏移量。
  42 +
  43 + Args:
  44 + url: m3u8 视频流地址。
  45 + start_utc: 截取开始时间(秒偏移或绝对 UTC 时间戳)。
  46 + end_utc: 截取结束时间(秒偏移或绝对 UTC 时间戳)。
  47 + fps: 目标抽帧帧率。
  48 + roi: 感兴趣区域 (x, y, w, h),可选。
  49 + max_px_area: 最大像素面积,超过则等比例缩小,可选。
  50 +
  51 + Returns:
  52 + 提取的帧图片路径列表。
  53 + """
  54 + if end_utc <= start_utc:
  55 + raise ValueError("end_utc 必须大于 start_utc")
  56 +
  57 + unique_str = f"{url}_{start_utc}_{end_utc}"
  58 + unique_id = hashlib.md5(unique_str.encode("utf-8")).hexdigest()
  59 +
  60 + clip_path = self.video_dir / f"{unique_id}.mp4"
  61 + frame_output_dir = self.frames_dir / unique_id
  62 + frame_output_dir.mkdir(parents=True, exist_ok=True)
  63 +
  64 + # 判断是否为绝对 UTC 时间戳
  65 + is_absolute_utc = (
  66 + start_utc > _IS_ABSOLUTE_UTC_THRESHOLD
  67 + and end_utc > _IS_ABSOLUTE_UTC_THRESHOLD
  68 + )
  69 +
  70 + # 若片段未缓存,先截取目标片段再保存
  71 + if not clip_path.exists():
  72 + if is_absolute_utc:
  73 + # 绝对 UTC 模式:先下载完整视频到临时文件,解析 SEI 后截取片段
  74 + temp_path = self.video_dir / f"{unique_id}_full.mp4"
  75 + download_m3u8_to_mp4(url, str(temp_path))
  76 +
  77 + # 提取 SEI UTC 信息
  78 + utc_records = []
  79 + try:
  80 + es_data = extract_h264_es_from_mp4(str(temp_path))
  81 + sei_list = parse_h264_sei(es_data)
  82 + utc_records = extract_utc_from_sei(sei_list)
  83 + except Exception:
  84 + pass
  85 +
  86 + if utc_records:
  87 + utc_list = [r["utc"] for r in utc_records]
  88 + start_frame_idx = self._find_frame_idx(
  89 + utc_list, start_utc, find_last_not_greater=False
  90 + )
  91 + end_frame_idx = self._find_frame_idx(
  92 + utc_list, end_utc, find_last_not_greater=True
  93 + )
  94 +
  95 + cap_temp = cv2.VideoCapture(str(temp_path))
  96 + video_fps = cap_temp.get(cv2.CAP_PROP_FPS)
  97 + cap_temp.release()
  98 + if video_fps <= 0:
  99 + video_fps = 25.0
  100 +
  101 + start_sec = start_frame_idx / video_fps
  102 + duration = (end_frame_idx - start_frame_idx) / video_fps
  103 + self._ffmpeg_extract_clip(
  104 + str(temp_path), str(clip_path), start_sec, duration
  105 + )
  106 + else:
  107 + # 无 SEI 无法定位,直接将完整视频重命名为片段
  108 + temp_path.rename(clip_path)
  109 +
  110 + # 清理临时完整视频
  111 + if temp_path.exists():
  112 + temp_path.unlink(missing_ok=True)
  113 + else:
  114 + # 秒偏移模式:直接用 ffmpeg 从 URL 截取片段
  115 + duration = end_utc - start_utc
  116 + self._ffmpeg_download_clip(
  117 + url, str(clip_path), start_utc, duration
  118 + )
  119 +
  120 + # 从截取后的片段中提取帧
  121 + cap = cv2.VideoCapture(str(clip_path))
  122 + if not cap.isOpened():
  123 + raise ValueError(f"无法打开视频片段: {clip_path}")
  124 +
  125 + # total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
  126 + video_fps = cap.get(cv2.CAP_PROP_FPS)
  127 + if video_fps <= 0:
  128 + video_fps = 25.0
  129 +
  130 + # 计算抽帧间隔
  131 + frame_interval = int(video_fps / fps) if fps > 0 else int(video_fps)
  132 + if frame_interval < 1:
  133 + frame_interval = 1
  134 +
  135 + frames = []
  136 + saved_count = 0
  137 + frame_count = 0
  138 + while True:
  139 + ret, frame = cap.read()
  140 + if not ret:
  141 + break
  142 +
  143 + if frame_count % frame_interval == 0:
  144 + if roi is not None:
  145 + x, y, w, h = roi
  146 + frame = frame[y:y + h, x:x + w]
  147 + frame = self.reisize_frame(frame, max_px_area)
  148 + frames.append(frame)
  149 + frame_path = frame_output_dir / f"frame_{saved_count:06d}.jpg"
  150 + cv2.imwrite(str(frame_path), frame)
  151 + saved_count += 1
  152 + frame_count += 1
  153 + cap.release()
  154 + return frames
  155 +
  156 + def reisize_frame(self, frame, max_px_area):
  157 + if max_px_area is None:
  158 + return frame
  159 + h, w = frame.shape[:2]
  160 + area = h * w
  161 + if area > max_px_area:
  162 + scale = (max_px_area / area) ** 0.5
  163 + new_w = int(w * scale)
  164 + new_h = int(h * scale)
  165 + frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
  166 + return frame
  167 +
  168 + @staticmethod
  169 + def _find_frame_idx(utc_list, target_utc, find_last_not_greater=True):
  170 + """
  171 + 在 utc_list 中查找最接近 target_utc 的帧索引。
  172 +
  173 + Args:
  174 + utc_list: 按帧顺序排列的 UTC 列表。
  175 + target_utc: 目标 UTC。
  176 + find_last_not_greater: True 返回最后一个 <= target_utc 的索引;
  177 + False 返回第一个 >= target_utc 的索引。
  178 + """
  179 + if not utc_list:
  180 + return 0
  181 +
  182 + if find_last_not_greater:
  183 + idx = 0
  184 + for i, utc in enumerate(utc_list):
  185 + if utc > target_utc:
  186 + break
  187 + idx = i
  188 + return idx
  189 + else:
  190 + for i, utc in enumerate(utc_list):
  191 + if utc >= target_utc:
  192 + return i
  193 + return len(utc_list) - 1
  194 +
  195 + @staticmethod
  196 + def _ffmpeg_extract_clip(input_path, output_path, start_sec, duration):
  197 + """使用 ffmpeg 从本地视频截取指定片段。"""
  198 + cmd = [
  199 + "ffmpeg", "-y",
  200 + "-i", input_path,
  201 + "-ss", str(start_sec),
  202 + "-t", str(duration),
  203 + "-c", "copy",
  204 + "-bsf:a", "aac_adtstoasc",
  205 + "-movflags", "+faststart",
  206 + output_path,
  207 + ]
  208 + try:
  209 + subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  210 + except subprocess.CalledProcessError as e:
  211 + raise RuntimeError(
  212 + f"ffmpeg 截取片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
  213 + )
  214 +
  215 + @staticmethod
  216 + def _ffmpeg_download_clip(url, output_path, start_sec, duration):
  217 + """使用 ffmpeg 从 URL 下载并截取指定片段。"""
  218 + cmd = [
  219 + "ffmpeg", "-y",
  220 + "-ss", str(start_sec),
  221 + "-fflags", "+discardcorrupt",
  222 + "-i", url,
  223 + "-t", str(duration),
  224 + "-c", "copy",
  225 + "-bsf:a", "aac_adtstoasc",
  226 + "-movflags", "+faststart",
  227 + output_path,
  228 + ]
  229 + try:
  230 + subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  231 + except subprocess.CalledProcessError as e:
  232 + raise RuntimeError(
  233 + f"ffmpeg 下载片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
  234 + )
  235 +
  236 +
  237 +def frames2content(frames):
  238 + """
  239 + 将帧列表(图片路径列表)转为 LLM content 列表。
  240 + 格式与 llm_video_content.py 中的 contents 返回的 content 一致。
  241 + """
  242 + contents = []
  243 + video_prompt = (
  244 + f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面,"
  245 + f"请将它们视为一个连续的视频进行分析。"
  246 + )
  247 + contents.append({"type": "text", "text": video_prompt})
  248 +
  249 + for frame in frames:
  250 + _, buffer = cv2.imencode('.jpg', frame)
  251 + b64_str = base64.b64encode(buffer).decode('utf-8')
  252 +
  253 + # with open(frame_path, "rb") as f:
  254 + # img_bytes = f.read()
  255 + # b64_str = base64.b64encode(img_bytes).decode("utf-8")
  256 +
  257 + contents.append({
  258 + "type": "image_url",
  259 + "image_url": {
  260 + "url": f"data:image/jpeg;base64,{b64_str}"
  261 + }
  262 + })
  263 +
  264 + return contents
  265 +
  266 +
  267 +if __name__ == "__main__":
  268 + # pass
  269 +
  270 + import argparse
  271 + import sys
  272 + import tempfile
  273 +
  274 + parser = argparse.ArgumentParser(description="Video2Frame 测试脚本")
  275 + parser.add_argument(
  276 + "--url",
  277 + default=r"D:\pythonProject\learn\3b3e99cf4ca84c3782503d8817242de2.mp4",
  278 + help="测试用 m3u8 地址(默认使用公开测试流)",
  279 + )
  280 + parser.add_argument(
  281 + "--cache-dir",
  282 + default=r"C:\Users\lzw\AppData\Local\Temp\video2frame_test_mn26tcy_my",
  283 + help="缓存目录(默认自动创建临时目录)",
  284 + )
  285 + parser.add_argument(
  286 + "--start",
  287 + type=float,
  288 + default=0,
  289 + help="截取开始时间,单位:秒(默认 0)",
  290 + )
  291 + parser.add_argument(
  292 + "--end",
  293 + type=float,
  294 + default=10,
  295 + help="截取结束时间,单位:秒(默认 10)",
  296 + )
  297 + parser.add_argument(
  298 + "--fps",
  299 + type=float,
  300 + default=2.0,
  301 + help="抽帧帧率(默认 1.0)",
  302 + )
  303 + args = parser.parse_args()
  304 +
  305 + cache_dir = args.cache_dir or tempfile.mkdtemp(prefix="video2frame_test_")
  306 + print(f"[INFO] 缓存目录: {cache_dir}")
  307 +
  308 + v2f = Video2Frame(cache_dir)
  309 +
  310 + # TEST 1: 基本抽帧
  311 + print(
  312 + f"\n[TEST 1] to_frames: url={args.url}, "
  313 + f"start={args.start}s, end={args.end}s, fps={args.fps}"
  314 + )
  315 + try:
  316 + frames = v2f.to_frames(
  317 + url=args.url,
  318 + start_utc=args.start,
  319 + end_utc=args.end,
  320 + fps=args.fps,
  321 + roi=None,
  322 + max_px_area=200_000,
  323 + )
  324 + print(f" 成功提取 {len(frames)} 帧")
  325 + if frames:
  326 + print(f" 首帧: {frames[0].shape}")
  327 + print(f" 末帧: {frames[-1].shape}")
  328 + except Exception as e:
  329 + print(f" 失败: {e}")
  330 + import traceback
  331 +
  332 + traceback.print_exc()
  333 + sys.exit(1)
  334 +
  335 + # TEST 2: frames2content
  336 + print("\n[TEST 2] frames2content")
  337 + content = frames2content(frames)
  338 + print(f" content 列表长度: {len(content)}")
  339 + for item in content[:3]:
  340 + preview = (
  341 + item["text"][:60] + "..."
  342 + if item["type"] == "text"
  343 + else item["image_url"]["url"][:60] + "..."
  344 + )
  345 + print(f" - type={item['type']}: {preview}")
  346 + if len(content) > 3:
  347 + print(f" ... 还有 {len(content) - 3} 个元素")
  348 +
  349 + print("\n[TEST 3] FootballReplayVideoEvent 进球识别(可选)")
  350 +
  351 + try:
  352 + from ..util.football_replay_video_event_by_llm import FootballReplayVideoEvent
  353 + except ImportError:
  354 + import sys
  355 + from pathlib import Path
  356 + sys.path.insert(0, str(Path(__file__).parent.parent))
  357 + from util.football_replay_video_event_by_llm import FootballReplayVideoEvent
  358 +
  359 + fbrv = FootballReplayVideoEvent(
  360 + base_url="http://192.168.1.59:11434",
  361 + model="qwen3.6:35b-a3b-q8_0",
  362 + temperature=0.7,
  363 + )
  364 + event_json = fbrv.video_event_for_contents(content, None)
  365 + print(f" 识别结果: {event_json}")
  366 +
  367 +
  368 + # # TEST 3: ROI + max_px_area
  369 + # print("\n[TEST 3] to_frames with roi + max_px_area")
  370 + # try:
  371 + # end_crop = min(args.start + 5, args.end)
  372 + # frames_cropped = v2f.to_frames(
  373 + # url=args.url,
  374 + # start_utc=args.start,
  375 + # end_utc=end_crop,
  376 + # fps=1.0,
  377 + # roi=None,
  378 + # max_px_area=200_000,
  379 + # )
  380 + # print(f" 成功提取 {len(frames_cropped)} 帧(带 ROI 裁剪和面积缩放)")
  381 + # except Exception as e:
  382 + # print(f" 跳过/失败: {e}")
  383 + #
  384 + # # TEST 4: SEI UTC 模式说明(仅演示,需真实含 SEI 的流才能运行)
  385 + # print("\n[TEST 4] SEI UTC 绝对时间戳模式(示例代码,需替换为含 SEI 的流)")
  386 + # print(
  387 + # " # 示例:假设视频第一帧 UTC 为 1715600000,提取 10 秒片段\n"
  388 + # " # frames = v2f.to_frames(url, start_utc=1715600000, end_utc=1715600010, fps=1)"
  389 + # )
  390 + #
  391 + # print("\n[INFO] 所有测试完成")