llm_image.py
13.3 KB
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
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/v1",
model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf",
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] 所有测试完成")