-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.py
More file actions
96 lines (75 loc) · 2.91 KB
/
main.py
File metadata and controls
96 lines (75 loc) · 2.91 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
#!/usr/bin/env python3
import argparse
import os
import sys
from pathlib import Path
from frame_extractor import FrameExtractor
from frame_stitcher import FrameStitcher
from ocr_processor import OCRProcessor
def main():
parser = argparse.ArgumentParser(description='视频滚动文章提取工具')
parser.add_argument('input_video', help='输入视频文件路径')
parser.add_argument('-o', '--output', default='output', help='输出目录 (默认: output)')
parser.add_argument('--keep-frames', action='store_true', help='保留中间帧文件')
args = parser.parse_args()
# 检查输入文件
if not os.path.exists(args.input_video):
print(f"错误: 输入视频文件不存在: {args.input_video}")
sys.exit(1)
# 创建输出目录
output_dir = Path(args.output)
output_dir.mkdir(exist_ok=True)
frames_dir = output_dir / "frames"
long_image_path = output_dir / "long_article.png"
text_output_path = output_dir / "article.txt"
print("=== 视频滚动文章提取工具 ===")
print(f"输入视频: {args.input_video}")
print(f"输出目录: {output_dir}")
print()
# 步骤1: 提取帧
print("步骤1: 提取视频帧...")
extractor = FrameExtractor()
# 优先使用ffmpeg，失败则使用OpenCV
success = extractor.extract_frames_ffmpeg(args.input_video, str(frames_dir))
if not success:
print("ffmpeg提取失败，尝试使用OpenCV...")
success = extractor.extract_frames_opencv(args.input_video, str(frames_dir))
if not success:
print("错误: 帧提取失败")
sys.exit(1)
print()
# 步骤2: 拼接长图
print("步骤2: 拼接长图...")
stitcher = FrameStitcher()
stitched_img = stitcher.stitch_frames(str(frames_dir))
if stitched_img is None:
print("错误: 帧拼接失败")
sys.exit(1)
success = stitcher.save_stitched_image(stitched_img, str(long_image_path))
if not success:
print("错误: 长图保存失败")
sys.exit(1)
print()
# 步骤3: OCR识别 (使用Tesseract)
print("步骤3: OCR文字识别 (Tesseract)...")
ocr_processor = OCRProcessor()
texts = ocr_processor.extract_text_tesseract(str(long_image_path))
if not texts:
print("警告: 未识别到文字")
else:
success = ocr_processor.save_text(texts, str(text_output_path))
if not success:
print("错误: 文字保存失败")
sys.exit(1)
print(f"识别了 {len(texts)} 行文字")
print()
# 清理中间文件
if not args.keep_frames:
print("清理中间文件...")
import shutil
shutil.rmtree(frames_dir)
print("=== 处理完成 ===")
print(f"长图保存到: {long_image_path}")
print(f"文字保存到: {text_output_path}")
if __name__ == "__main__":
main()