lizhengwei / aigc-embedding-service

Go to a project
Toggle navigation Toggle navigation pinning
lizhengwei
1396dd74 1 parent 1ad47aa5

jira:NYJ-1540 desc: update llm_image.py

Inline Side-by-side
Showing 2 changed files with 735 additions and 3 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):
pass
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:
pass
"""
下载 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
frame_paths = []
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]
if max_px_area is not None:
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
)
frame_path = frame_output_dir / f"frame_{saved_count:06d}.jpg"
cv2.imwrite(str(frame_path), frame)
frame_paths.append(str(frame_path))
saved_count += 1
frame_count += 1
cap.release()
return frame_paths
@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):
return
\ No newline at end of file
"""
将帧列表（图片路径列表）转为 LLM content 列表。
格式与 llm_video_content.py 中的 contents 返回的 content 一致。
"""
content = []
video_prompt = (
f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面，"
f"请将它们视为一个连续的视频进行分析。"
)
content.append({"type": "text", "text": video_prompt})
for frame_path in frames:
with open(frame_path, "rb") as f:
img_bytes = f.read()
b64_str = base64.b64encode(img_bytes).decode("utf-8")
content.append({
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{b64_str}"
}
})
return content
if __name__ == "__main__":
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=None,
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=1.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,
)
print(f" 成功提取 {len(frames)} 帧")
if frames:
print(f" 首帧: {frames[0]}")
print(f" 末帧: {frames[-1]}")
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} 个元素")
# 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] 所有测试完成")
... ...
#!/usr/bin/env python3
"""
m3u8_to_mp4_sei.py
下载 m3u8 为 mp4，并提取视频中的 SEI UTC 信息。
依赖:
ffmpeg (需要添加到系统 PATH)
用法示例:
python m3u8_to_mp4_sei.py "http://example.com/playlist.m3u8"
python m3u8_to_mp4_sei.py "http://example.com/playlist.m3u8" -o myvideo.mp4 --json
python m3u8_to_mp4_sei.py "http://example.com/playlist.m3u8" -t 30 --json
"""
import argparse
import json
import os
import struct
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional
def check_ffmpeg() -> bool:
"""检查系统是否安装了 ffmpeg"""
try:
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
return True
except (FileNotFoundError, subprocess.CalledProcessError):
return False
def download_m3u8_to_mp4(url: str, output_path: str, duration: Optional[int] = None) -> None:
"""
使用 ffmpeg 将 m3u8 下载并封装为 mp4
:param duration: 若指定，则只下载前 N 秒
"""
cmd = [
"ffmpeg",
"-y",
"-fflags", "+discardcorrupt", # 容错：丢弃损坏包
"-i", url,
]
if duration is not None:
cmd.extend(["-t", str(duration)])
cmd.extend([
"-c", "copy",
"-bsf:a", "aac_adtstoasc",
"-movflags", "+faststart",
output_path,
])
print(f"[1/3] 正在下载并封装 mp4: {output_path}")
result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
if result.returncode != 0:
raise RuntimeError(f"ffmpeg 下载失败: {result.stderr}")
def extract_h264_es_from_mp4(mp4_path: str) -> bytes:
"""
使用 ffmpeg 从 mp4 中提取 H.264 Elementary Stream (AnnexB 格式)
返回完整的 ES 字节流
"""
cmd = [
"ffmpeg",
"-y",
"-i", mp4_path,
"-c:v", "copy",
"-an",
"-bsf:v", "h264_mp4toannexb",
"-f", "h264",
"pipe:1",
]
print("[2/3] 正在从 mp4 中提取 H.264 视频流...")
proc = subprocess.run(cmd, capture_output=True)
if proc.returncode != 0:
stderr = proc.stderr.decode("utf-8", errors="replace") if proc.stderr else ""
raise RuntimeError(f"ffmpeg 提取视频流失败: {stderr}")
return proc.stdout
def _sei_type_name(payload_type: int) -> str:
names = {
0: "buffering_period",
1: "pic_timing",
2: "pan_scan_rect",
3: "filler_payload",
4: "user_data_registered_itu_t_t35",
5: "user_data_unregistered",
6: "recovery_point",
7: "dec_ref_pic_marking_repetition",
8: "spare_pic",
9: "scene_info",
10: "sub_seq_info",
11: "sub_seq_layer_characteristics",
12: "sub_seq_characteristics",
13: "full_frame_freeze",
14: "full_frame_freeze_release",
15: "full_frame_snapshot",
16: "progressive_refinement_segment_start",
17: "progressive_refinement_segment_end",
18: "motion_constrained_slice_group_set",
19: "film_grain_characteristics",
20: "deblocking_filter_display_preference",
21: "stereo_video_info",
22: "post_filter_hint",
23: "tone_mapping_info",
24: "scalability_info",
25: "sub_pic_scalable_layer",
26: "non_required_layer_rep",
27: "priority_layer_info",
28: "layers_not_present",
29: "layer_dependency_change",
30: "scalable_nesting",
31: "base_layer_temporal_hrd",
32: "quality_layer_integrity_check",
33: "redundant_pic_property",
34: "tl0_dep_rep_index",
35: "tl_switching_point",
36: "parallel_decoding_info",
37: "mvc_scalable_nesting",
38: "view_scalability_info",
39: "multiview_scene_info",
40: "multiview_acquisition_info",
41: "non_equivalent_view_dependency",
42: "view_dependency_change",
43: "operation_points_not_present",
44: "base_view_temporal_hrd",
45: "frame_packing_arrangement",
46: "multiview_view_position",
47: "display_orientation",
48: "mvcd_scalable_nesting",
49: "mvcd_view_scalability_info",
50: "depth_representation_info",
51: "three_dimensional_reference_displays_info",
52: "depth_timing",
53: "depth_sampling_info",
54: "constrained_depth_parameter_set_identifier",
55: "green_metadata",
56: "mastering_display_colour_volume",
57: "colour_remapping_info",
58: "alternative_transfer_characteristics",
59: "alternative_depth_info",
}
return names.get(payload_type, f"unknown({payload_type})")
def parse_h264_sei(es_data: bytes) -> List[Dict]:
"""
从 H.264 ES 数据中解析 SEI NAL 单元
返回 SEI 列表，包含 payload_type, uuid, text/data
"""
sei_list = []
i = 0
def find_next_start(data: bytes, start: int) -> int:
j = start
while j + 3 < len(data):
if data[j:j + 4] == b"\x00\x00\x00\x01":
return j
if data[j:j + 3] == b"\x00\x00\x01":
return j
j += 1
return len(data)
while i < len(es_data):
if es_data[i:i + 4] == b"\x00\x00\x00\x01":
start_len = 4
elif es_data[i:i + 3] == b"\x00\x00\x01":
start_len = 3
else:
i += 1
continue
if i + start_len >= len(es_data):
break
nal_unit_type = es_data[i + start_len] & 0x1F
if nal_unit_type == 6: # SEI
payload_start = i + start_len + 1
next_start = find_next_start(es_data, payload_start)
sei_payload = es_data[payload_start:next_start]
k = 0
payload_type = 0
while k < len(sei_payload) and sei_payload[k] == 0xFF:
payload_type += 255
k += 1
if k < len(sei_payload):
payload_type += sei_payload[k]
k += 1
payload_size = 0
while k < len(sei_payload) and sei_payload[k] == 0xFF:
payload_size += 255
k += 1
if k < len(sei_payload):
payload_size += sei_payload[k]
k += 1
if k + payload_size <= len(sei_payload):
payload_data = sei_payload[k:k + payload_size]
sei_info = {
"payload_type": payload_type,
"payload_type_name": _sei_type_name(payload_type),
"payload_size": payload_size,
}
if payload_type == 5: # user_data_unregistered
if len(payload_data) >= 16:
uuid_bytes = payload_data[:16]
user_data = payload_data[16:]
while user_data and user_data[-1] == 0x00:
user_data = user_data[:-1]
if user_data and user_data[-1] == 0x80:
user_data = user_data[:-1]
sei_info["uuid_hex"] = uuid_bytes.hex()
sei_info["uuid_ascii"] = uuid_bytes.decode("ascii", errors="replace")
try:
text = user_data.decode("utf-8", errors="replace")
sei_info["text"] = text
try:
sei_info["json"] = json.loads(text)
except json.JSONDecodeError:
pass
except Exception:
sei_info["raw_hex"] = user_data.hex()
sei_list.append(sei_info)
i = next_start
else:
i += 1
return sei_list
def extract_utc_from_sei(sei_list: List[Dict]) -> List[Dict]:
"""从 SEI 列表中提取 UTC 信息"""
utc_records = []
seen = set()
for sei in sei_list:
json_data = sei.get("json")
if isinstance(json_data, dict) and "utc" in json_data:
# 去重：相同的 utc 只保留一次
key = json.dumps(json_data, sort_keys=True, ensure_ascii=False)
if key in seen:
continue
seen.add(key)
utc_records.append({
"utc": json_data["utc"],
"origin_utc": json_data.get("origin_utc"),
"origin_offset": json_data.get("origin_offset"),
"full_json": json_data,
})
return utc_records
def main():
parser = argparse.ArgumentParser(
description="将 m3u8 下载为 mp4 并提取 SEI UTC 信息",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
示例:
%(prog)s "http://example.com/playlist.m3u8"
%(prog)s "http://example.com/playlist.m3u8" -o myvideo.mp4 --json
%(prog)s "http://example.com/playlist.m3u8" -t 30 --json -o result.json
""",
)
parser.add_argument("url", help="m3u8 播放列表 URL")
parser.add_argument("-o", "--output", default=None, help="mp4 输出路径（默认根据 URL 自动命名）")
parser.add_argument("-t", "--duration", type=int, default=None, help="只下载前 N 秒（默认下载全部）")
parser.add_argument("--json", action="store_true", help="以 JSON 格式输出结果")
parser.add_argument("--no-download", action="store_true", help="如果本地已有 mp4，跳过下载")
parser.add_argument("--save-es", metavar="FILE", help="保存提取的 H.264 ES 流到指定文件（调试用）")
args = parser.parse_args()
# 强制 stdout UTF-8
import io
try:
if getattr(sys.stdout, "encoding", None) != "utf-8":
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
elif hasattr(sys.stdout, "buffer"):
sys.stdout = io.TextIOWrapper(
sys.stdout.buffer, encoding="utf-8", line_buffering=True
)
except Exception:
pass
if not check_ffmpeg():
print("错误: 未检测到 ffmpeg，请先安装并添加到 PATH。", file=sys.stderr)
print("下载地址: https://ffmpeg.org/download.html", file=sys.stderr)
sys.exit(1)
# 确定 mp4 输出路径
if args.output:
mp4_path = Path(args.output)
else:
base = args.url.split("?")[0].rsplit("/", 1)[-1]
if not base or "." not in base:
base = "output.mp4"
else:
base = base.rsplit(".", 1)[0] + ".mp4"
mp4_path = Path(base)
es_data = b""
try:
# 1. 下载/封装 mp4
if not args.no_download or not mp4_path.exists():
download_m3u8_to_mp4(args.url, str(mp4_path), duration=args.duration)
print(f" mp4 已保存: {mp4_path.absolute()}")
else:
print(f"[1/3] 使用本地 mp4: {mp4_path.absolute()}")
# 2. 提取 H.264 ES
es_data = extract_h264_es_from_mp4(str(mp4_path))
print(f" 提取到 {len(es_data)} 字节的 H.264 ES 数据")
if args.save_es:
es_path = Path(args.save_es)
es_path.write_bytes(es_data)
print(f" H.264 ES 已保存: {es_path.absolute()}")
# 3. 解析 SEI
sei_list = parse_h264_sei(es_data)
print(f"[3/3] 解析完成，发现 {len(sei_list)} 条 SEI 消息")
# 4. 提取 UTC
utc_records = extract_utc_from_sei(sei_list)
result = {
"m3u8_url": args.url,
"mp4_path": str(mp4_path.absolute()),
"sei_count": len(sei_list),
"utc_records": utc_records,
"sei_items": sei_list[:20],
}
# 输出
json_str = json.dumps(result, ensure_ascii=False, indent=2)
if args.json and args.output and args.output.endswith(".json"):
# 如果 --output 以 .json 结尾，且 --json，则保存为 JSON 文件
with open(args.output, "w", encoding="utf-8-sig") as f:
f.write(json_str)
f.write("\n")
print(f"JSON 结果已保存到: {args.output}")
elif args.json:
print(json_str)
else:
print(f"\n{'='*50}")
print("SEI UTC 信息提取结果")
print(f"{'='*50}")
print(f"mp4 文件 : {result['mp4_path']}")
print(f"SEI 消息总数 : {result['sei_count']}")
print(f"UTC 记录数 : {len(utc_records)}")
if utc_records:
print(f"\nUTC 详细信息:")
for idx, rec in enumerate(utc_records, 1):
print(f" [{idx}] utc={rec['utc']}, origin_utc={rec.get('origin_utc')}, origin_offset={rec.get('origin_offset')}")
else:
print("\n未从 SEI 中提取到 UTC 信息。")
print(f"{'='*50}")
except Exception as e:
print(f"错误: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()
# python mp4_sei_extractor.py "3b3e99cf4ca84c3782503d8817242de2.mp4" --json -o sei_result.json
\ No newline at end of file
... ...