lizhengwei / aigc-embedding-service

Go to a project
Toggle navigation Toggle navigation pinning
wangdongxu
2fd09ae7 1 parent 40f4c17d

jira:NYJ-1550 desc:m3u8 images

Inline Side-by-side
Showing 1 changed file with 391 additions and 0 deletions
import base64
import hashlib
import subprocess
from pathlib import Path
import cv2
try:
from .m3u8_to_mp4_sei import (
download_m3u8_to_mp4,
extract_h264_es_from_mp4,
parse_h264_sei,
extract_utc_from_sei,
)
except ImportError:
from m3u8_to_mp4_sei import (
download_m3u8_to_mp4,
extract_h264_es_from_mp4,
parse_h264_sei,
extract_utc_from_sei,
)
# 用于区分传入的 start_utc/end_utc 是绝对 UTC 时间戳还是秒偏移
_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000
class Video2Frame:
def __init__(self, cache_dir):
self.cache_dir = Path(cache_dir)
self.video_dir = self.cache_dir / "videos"
self.frames_dir = self.cache_dir / "frames"
self.video_dir.mkdir(parents=True, exist_ok=True)
self.frames_dir.mkdir(parents=True, exist_ok=True)
def to_frames(self, url, start_utc, end_utc, fps, roi=None, max_px_area=None) -> list:
"""
下载 m3u8 视频流的指定片段并提取帧，保存到 cache_dir 下。
- video_dir 中保存的是 start_utc ~ end_utc 截取后的视频片段（而非完整视频）。
- 如果视频中包含 SEI UTC 信息且传入的 start_utc/end_utc 为绝对时间戳，
则会基于 SEI UTC 进行帧级定位；否则将 start_utc/end_utc 视为秒偏移量。
Args:
url: m3u8 视频流地址。
start_utc: 截取开始时间（秒偏移或绝对 UTC 时间戳）。
end_utc: 截取结束时间（秒偏移或绝对 UTC 时间戳）。
fps: 目标抽帧帧率。
roi: 感兴趣区域 (x, y, w, h)，可选。
max_px_area: 最大像素面积，超过则等比例缩小，可选。
Returns:
提取的帧图片路径列表。
"""
if end_utc <= start_utc:
raise ValueError("end_utc 必须大于 start_utc")
unique_str = f"{url}_{start_utc}_{end_utc}"
unique_id = hashlib.md5(unique_str.encode("utf-8")).hexdigest()
clip_path = self.video_dir / f"{unique_id}.mp4"
frame_output_dir = self.frames_dir / unique_id
frame_output_dir.mkdir(parents=True, exist_ok=True)
# 判断是否为绝对 UTC 时间戳
is_absolute_utc = (
start_utc > _IS_ABSOLUTE_UTC_THRESHOLD
and end_utc > _IS_ABSOLUTE_UTC_THRESHOLD
)
# 若片段未缓存，先截取目标片段再保存
if not clip_path.exists():
if is_absolute_utc:
# 绝对 UTC 模式：先下载完整视频到临时文件，解析 SEI 后截取片段
temp_path = self.video_dir / f"{unique_id}_full.mp4"
download_m3u8_to_mp4(url, str(temp_path))
# 提取 SEI UTC 信息
utc_records = []
try:
es_data = extract_h264_es_from_mp4(str(temp_path))
sei_list = parse_h264_sei(es_data)
utc_records = extract_utc_from_sei(sei_list)
except Exception:
pass
if utc_records:
utc_list = [r["utc"] for r in utc_records]
start_frame_idx = self._find_frame_idx(
utc_list, start_utc, find_last_not_greater=False
)
end_frame_idx = self._find_frame_idx(
utc_list, end_utc, find_last_not_greater=True
)
cap_temp = cv2.VideoCapture(str(temp_path))
video_fps = cap_temp.get(cv2.CAP_PROP_FPS)
cap_temp.release()
if video_fps <= 0:
video_fps = 25.0
start_sec = start_frame_idx / video_fps
duration = (end_frame_idx - start_frame_idx) / video_fps
self._ffmpeg_extract_clip(
str(temp_path), str(clip_path), start_sec, duration
)
else:
# 无 SEI 无法定位，直接将完整视频重命名为片段
temp_path.rename(clip_path)
# 清理临时完整视频
if temp_path.exists():
temp_path.unlink(missing_ok=True)
else:
# 秒偏移模式：直接用 ffmpeg 从 URL 截取片段
duration = end_utc - start_utc
self._ffmpeg_download_clip(
url, str(clip_path), start_utc, duration
)
# 从截取后的片段中提取帧
cap = cv2.VideoCapture(str(clip_path))
if not cap.isOpened():
raise ValueError(f"无法打开视频片段: {clip_path}")
# total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
video_fps = cap.get(cv2.CAP_PROP_FPS)
if video_fps <= 0:
video_fps = 25.0
# 计算抽帧间隔
frame_interval = int(video_fps / fps) if fps > 0 else int(video_fps)
if frame_interval < 1:
frame_interval = 1
frames = []
saved_count = 0
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break
if frame_count % frame_interval == 0:
if roi is not None:
x, y, w, h = roi
frame = frame[y:y + h, x:x + w]
frame = self.reisize_frame(frame, max_px_area)
frames.append(frame)
frame_path = frame_output_dir / f"frame_{saved_count:06d}.jpg"
cv2.imwrite(str(frame_path), frame)
saved_count += 1
frame_count += 1
cap.release()
return frames
def reisize_frame(self, frame, max_px_area):
if max_px_area is None:
return frame
h, w = frame.shape[:2]
area = h * w
if area > max_px_area:
scale = (max_px_area / area) ** 0.5
new_w = int(w * scale)
new_h = int(h * scale)
frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
return frame
@staticmethod
def _find_frame_idx(utc_list, target_utc, find_last_not_greater=True):
"""
在 utc_list 中查找最接近 target_utc 的帧索引。
Args:
utc_list: 按帧顺序排列的 UTC 列表。
target_utc: 目标 UTC。
find_last_not_greater: True 返回最后一个 <= target_utc 的索引；
False 返回第一个 >= target_utc 的索引。
"""
if not utc_list:
return 0
if find_last_not_greater:
idx = 0
for i, utc in enumerate(utc_list):
if utc > target_utc:
break
idx = i
return idx
else:
for i, utc in enumerate(utc_list):
if utc >= target_utc:
return i
return len(utc_list) - 1
@staticmethod
def _ffmpeg_extract_clip(input_path, output_path, start_sec, duration):
"""使用 ffmpeg 从本地视频截取指定片段。"""
cmd = [
"ffmpeg", "-y",
"-i", input_path,
"-ss", str(start_sec),
"-t", str(duration),
"-c", "copy",
"-bsf:a", "aac_adtstoasc",
"-movflags", "+faststart",
output_path,
]
try:
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
raise RuntimeError(
f"ffmpeg 截取片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
)
@staticmethod
def _ffmpeg_download_clip(url, output_path, start_sec, duration):
"""使用 ffmpeg 从 URL 下载并截取指定片段。"""
cmd = [
"ffmpeg", "-y",
"-ss", str(start_sec),
"-fflags", "+discardcorrupt",
"-i", url,
"-t", str(duration),
"-c", "copy",
"-bsf:a", "aac_adtstoasc",
"-movflags", "+faststart",
output_path,
]
try:
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
raise RuntimeError(
f"ffmpeg 下载片段失败: {e.stderr.decode('utf-8', errors='ignore')}"
)
def frames2content(frames):
"""
将帧列表（图片路径列表）转为 LLM content 列表。
格式与 llm_video_content.py 中的 contents 返回的 content 一致。
"""
contents = []
video_prompt = (
f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面，"
f"请将它们视为一个连续的视频进行分析。"
)
contents.append({"type": "text", "text": video_prompt})
for frame in frames:
_, buffer = cv2.imencode('.jpg', frame)
b64_str = base64.b64encode(buffer).decode('utf-8')
# with open(frame_path, "rb") as f:
# img_bytes = f.read()
# b64_str = base64.b64encode(img_bytes).decode("utf-8")
contents.append({
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{b64_str}"
}
})
return contents
if __name__ == "__main__":
# pass
import argparse
import sys
import tempfile
parser = argparse.ArgumentParser(description="Video2Frame 测试脚本")
parser.add_argument(
"--url",
default=r"D:\pythonProject\learn\3b3e99cf4ca84c3782503d8817242de2.mp4",
help="测试用 m3u8 地址（默认使用公开测试流）",
)
parser.add_argument(
"--cache-dir",
default=r"C:\Users\lzw\AppData\Local\Temp\video2frame_test_mn26tcy_my",
help="缓存目录（默认自动创建临时目录）",
)
parser.add_argument(
"--start",
type=float,
default=0,
help="截取开始时间，单位：秒（默认 0）",
)
parser.add_argument(
"--end",
type=float,
default=10,
help="截取结束时间，单位：秒（默认 10）",
)
parser.add_argument(
"--fps",
type=float,
default=2.0,
help="抽帧帧率（默认 1.0）",
)
args = parser.parse_args()
cache_dir = args.cache_dir or tempfile.mkdtemp(prefix="video2frame_test_")
print(f"[INFO] 缓存目录: {cache_dir}")
v2f = Video2Frame(cache_dir)
# TEST 1: 基本抽帧
print(
f"\n[TEST 1] to_frames: url={args.url}, "
f"start={args.start}s, end={args.end}s, fps={args.fps}"
)
try:
frames = v2f.to_frames(
url=args.url,
start_utc=args.start,
end_utc=args.end,
fps=args.fps,
roi=None,
max_px_area=200_000,
)
print(f" 成功提取 {len(frames)} 帧")
if frames:
print(f" 首帧: {frames[0].shape}")
print(f" 末帧: {frames[-1].shape}")
except Exception as e:
print(f" 失败: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
# TEST 2: frames2content
print("\n[TEST 2] frames2content")
content = frames2content(frames)
print(f" content 列表长度: {len(content)}")
for item in content[:3]:
preview = (
item["text"][:60] + "..."
if item["type"] == "text"
else item["image_url"]["url"][:60] + "..."
)
print(f" - type={item['type']}: {preview}")
if len(content) > 3:
print(f" ... 还有 {len(content) - 3} 个元素")
print("\n[TEST 3] FootballReplayVideoEvent 进球识别（可选）")
try:
from ..util.football_replay_video_event_by_llm import FootballReplayVideoEvent
except ImportError:
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from util.football_replay_video_event_by_llm import FootballReplayVideoEvent
fbrv = FootballReplayVideoEvent(
base_url="http://192.168.1.59:11434",
model="qwen3.6:35b-a3b-q8_0",
temperature=0.7,
)
event_json = fbrv.video_event_for_contents(content, None)
print(f" 识别结果: {event_json}")
# # TEST 3: ROI + max_px_area
# print("\n[TEST 3] to_frames with roi + max_px_area")
# try:
# end_crop = min(args.start + 5, args.end)
# frames_cropped = v2f.to_frames(
# url=args.url,
# start_utc=args.start,
# end_utc=end_crop,
# fps=1.0,
# roi=None,
# max_px_area=200_000,
# )
# print(f" 成功提取 {len(frames_cropped)} 帧（带 ROI 裁剪和面积缩放）")
# except Exception as e:
# print(f" 跳过/失败: {e}")
#
# # TEST 4: SEI UTC 模式说明（仅演示，需真实含 SEI 的流才能运行）
# print("\n[TEST 4] SEI UTC 绝对时间戳模式（示例代码，需替换为含 SEI 的流）")
# print(
# " # 示例：假设视频第一帧 UTC 为 1715600000，提取 10 秒片段\n"
# " # frames = v2f.to_frames(url, start_utc=1715600000, end_utc=1715600010, fps=1)"
# )
#
# print("\n[INFO] 所有测试完成")
... ...