171 lines
6.5 KiB
Python
171 lines
6.5 KiB
Python
import os
|
||
import warnings
|
||
import imageio
|
||
import whisper
|
||
import numpy as np
|
||
from PIL import Image
|
||
from skimage.metrics import structural_similarity as ssim
|
||
from collections import defaultdict
|
||
import subprocess
|
||
|
||
# ======================== 配置参数 ========================
|
||
warnings.filterwarnings("ignore", message="FP16 is not supported on CPU; using FP32 instead")
|
||
VIDEO_PATH = "D:/python项目文件/1/input.mp4" # 输入视频路径（需修改为实际路径）
|
||
MODEL_DIR = "D:/whisper_models" # Whisper模型目录
|
||
FFMPEG_BIN = r"D:\Program Files\ffmpeg\bin" # FFmpeg安装目录
|
||
SSIM_THRESHOLD = 0.85 # 关键帧去重阈值（值越小保留越多帧）
|
||
FRAME_INTERVAL = 2 # 抽帧间隔（秒）
|
||
OUTPUT_DIR = "output_summary" # 输出目录
|
||
# ========================================================
|
||
|
||
# 动态添加FFmpeg到环境变量（修复核心问题）
|
||
os.environ["PATH"] = FFMPEG_BIN + os.pathsep + os.environ["PATH"]
|
||
|
||
def check_ffmpeg():
|
||
"""验证FFmpeg可访问性"""
|
||
try:
|
||
subprocess.run(["ffmpeg", "-version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
print("FFmpeg验证成功！")
|
||
return True
|
||
except Exception as e:
|
||
print(f"FFmpeg验证失败: {str(e)}")
|
||
print(f"请确认FFmpeg路径是否正确: {FFMPEG_BIN}")
|
||
return False
|
||
|
||
def extract_keyframes_with_time(video_path: str) -> tuple:
|
||
"""提取关键帧及其时间戳"""
|
||
try:
|
||
# 初始化视频读取器
|
||
reader = imageio.get_reader(video_path, 'ffmpeg')
|
||
fps = reader.get_meta_data().get('fps', 30)
|
||
keyframes = []
|
||
keyframe_times = []
|
||
prev_frame = None
|
||
|
||
for i, frame in enumerate(reader):
|
||
# 按固定间隔抽帧（每FRAME_INTERVAL秒抽1帧）
|
||
if i % int(fps * FRAME_INTERVAL) != 0:
|
||
continue
|
||
|
||
current_time = i / fps
|
||
# 降采样至320x240加速处理
|
||
curr_frame = Image.fromarray(frame).resize((320, 240))
|
||
|
||
if prev_frame is None:
|
||
# 首帧强制保留
|
||
keyframes.append(curr_frame)
|
||
keyframe_times.append(current_time)
|
||
prev_frame = np.array(curr_frame.convert('L'))
|
||
else:
|
||
# 计算灰度图SSIM
|
||
curr_gray = np.array(curr_frame.convert('L'))
|
||
score = ssim(prev_frame, curr_gray, data_range=255)
|
||
|
||
# 判定为关键帧的条件
|
||
if score < SSIM_THRESHOLD:
|
||
keyframes.append(curr_frame)
|
||
keyframe_times.append(current_time)
|
||
prev_frame = curr_gray
|
||
|
||
reader.close()
|
||
print(f"关键帧提取完成，共{len(keyframes)}帧")
|
||
return keyframes, keyframe_times
|
||
except Exception as e:
|
||
print(f"关键帧提取失败: {str(e)}")
|
||
return [], []
|
||
|
||
|
||
def generate_page_intervals(keyframe_times, video_duration):
|
||
"""生成每个关键帧对应的时间段: [start, end)"""
|
||
intervals = []
|
||
n = len(keyframe_times)
|
||
for i in range(n):
|
||
start = keyframe_times[i]
|
||
end = keyframe_times[i + 1] if i < n - 1 else video_duration
|
||
intervals.append((start, end))
|
||
return intervals
|
||
|
||
|
||
def align_text_with_keyframes(video_path: str, keyframe_times: list, video_duration: float) -> list:
|
||
"""语音-关键帧对齐（修复版）"""
|
||
try:
|
||
# 加载模型（不再传递ffmpeg_path）
|
||
model = whisper.load_model("tiny", device="cpu", download_root=MODEL_DIR)
|
||
result = model.transcribe(video_path, fp16=False) # 移除ffmpeg_path参数
|
||
segments = result["segments"]
|
||
|
||
# 生成时间段并聚合文本
|
||
page_intervals = generate_page_intervals(keyframe_times, video_duration)
|
||
page_texts = defaultdict(list)
|
||
|
||
# 将语音段落分配到对应页面（按起始时间判断）
|
||
for seg in segments:
|
||
seg_start = seg["start"]
|
||
seg_end = seg["end"]
|
||
for page_idx, (page_start, page_end) in enumerate(page_intervals):
|
||
if page_start <= seg_start < page_end:
|
||
page_texts[page_idx].append(seg["text"].strip())
|
||
break # 每个段落只属于一个页面
|
||
|
||
# 合并页面文本
|
||
merged = []
|
||
for page_idx in sorted(page_texts.keys()):
|
||
full_text = " ".join(page_texts[page_idx])
|
||
merged.append({
|
||
"page": page_idx,
|
||
"start_time": page_intervals[page_idx][0],
|
||
"end_time": page_intervals[page_idx][1],
|
||
"text": full_text
|
||
})
|
||
return merged
|
||
except Exception as e:
|
||
print(f"语音处理失败: {str(e)}")
|
||
return []
|
||
|
||
|
||
def save_summary(merged_texts, keyframes, output_dir):
|
||
"""保存关键帧和对应的完整语音文本"""
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
for idx, item in enumerate(merged_texts):
|
||
# 保存关键帧图像
|
||
img_path = os.path.join(output_dir, f"page_{idx}.jpg")
|
||
keyframes[idx].save(img_path)
|
||
|
||
# 保存文本
|
||
txt_path = os.path.join(output_dir, f"page_{idx}.txt")
|
||
with open(txt_path, "w", encoding="utf-8") as f:
|
||
f.write(f"页面时间段: {item['start_time']:.1f}s - {item['end_time']:.1f}s\n\n")
|
||
f.write("完整讲解内容:\n")
|
||
f.write(item["text"])
|
||
print(f"已保存: {txt_path}")
|
||
|
||
if __name__ == "__main__":
|
||
# 步骤0: 验证FFmpeg
|
||
if not check_ffmpeg():
|
||
exit()
|
||
|
||
# 步骤1: 检查视频文件存在性
|
||
if not os.path.exists(VIDEO_PATH):
|
||
print(f"错误：视频文件 {VIDEO_PATH} 不存在！")
|
||
exit()
|
||
|
||
# 步骤2: 提取关键帧
|
||
keyframes, keyframe_times = extract_keyframes_with_time(VIDEO_PATH)
|
||
if not keyframes:
|
||
exit()
|
||
|
||
# 步骤3: 获取视频总时长
|
||
reader = imageio.get_reader(VIDEO_PATH)
|
||
video_duration = reader.get_meta_data()["duration"]
|
||
reader.close()
|
||
|
||
# 步骤4: 语音对齐与聚合
|
||
merged_texts = align_text_with_keyframes(VIDEO_PATH, keyframe_times, video_duration)
|
||
|
||
# 步骤5: 保存结果
|
||
if merged_texts:
|
||
save_summary(merged_texts, keyframes, OUTPUT_DIR)
|
||
print(f"\n处理完成！结果已保存至目录: {os.path.abspath(OUTPUT_DIR)}")
|
||
else:
|
||
print("无有效语音内容可保存") |