-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathframe_stitcher.py
More file actions
152 lines (121 loc) · 6 KB
/
frame_stitcher.py
File metadata and controls
152 lines (121 loc) · 6 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
from PIL import Image, ImageChops
import numpy as np
import glob
import os
class FrameStitcher:
def __init__(self, threshold=0.8, identical_threshold=0.98):
self.threshold = threshold
self.identical_threshold = identical_threshold
def calculate_similarity(self, img1, img2):
"""计算两张图片的相似度"""
# 转换为numpy数组
arr1 = np.array(img1.convert('L'))
arr2 = np.array(img2.convert('L'))
# 确保尺寸相同
if arr1.shape != arr2.shape:
return 0.0
# 计算相关系数
correlation = np.corrcoef(arr1.flatten(), arr2.flatten())[0, 1]
return correlation if not np.isnan(correlation) else 0.0
def find_overlap(self, stitched_img, new_img):
"""找到两张图片的重叠区域"""
# 首先检查是否完全相同
similarity = self.calculate_similarity(stitched_img, new_img)
if similarity > self.identical_threshold:
print(f"检测到相同帧，相似度: {similarity:.3f}")
return True, new_img.height, "identical"
# 转换为灰度图
gray_stitched = stitched_img.convert('L')
gray_new = new_img.convert('L')
h_stitched = gray_stitched.height
w_stitched = gray_stitched.width
h_new = gray_new.height
w_new = gray_new.width
best_overlap = 0
best_similarity = 0
best_type = "none"
# 在stitched的底部区域寻找new_img的顶部
search_height = min(h_stitched // 2, h_new // 2)
for overlap_height in range(50, search_height, 20):
# 取stitched底部区域
stitched_bottom = gray_stitched.crop((0, h_stitched - overlap_height, w_stitched, h_stitched))
# 取new_img顶部区域
new_top = gray_new.crop((0, 0, w_new, overlap_height))
# 调整尺寸以匹配
if stitched_bottom.size == new_top.size:
# 计算相似度
similarity = self.calculate_similarity(stitched_bottom, new_top)
if similarity > best_similarity and similarity > self.threshold:
best_similarity = similarity
best_overlap = overlap_height
best_type = "bottom_top"
if best_similarity > self.threshold:
print(f"找到重叠区域，相似度: {best_similarity:.3f}, 重叠高度: {best_overlap}, 方法: {best_type}")
return True, best_overlap, best_type
return False, 0, "none"
def stitch_frames(self, frames_dir):
"""将滚动视频帧拼接成长图"""
# 获取所有帧文件
frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.png')])
if not frame_files:
print("未找到帧文件")
return None
frame_paths = [os.path.join(frames_dir, f) for f in frame_files]
# 读取第一帧作为基础
stitched = Image.open(frame_paths[0])
print(f"开始拼接，共 {len(frame_paths)} 帧")
print(f"第一帧尺寸: {stitched.size}")
skipped_frames = 0
for i, path in enumerate(frame_paths[1:], 1):
try:
img = Image.open(path)
print(f"\n处理帧 {i}, 尺寸: {img.size}")
# 检查尺寸是否匹配
if img.width != stitched.width:
print(f"帧 {i} 宽度不匹配，跳过")
skipped_frames += 1
continue
# 寻找重叠区域
has_overlap, overlap_height, overlap_type = self.find_overlap(stitched, img)
if has_overlap:
if overlap_type == "identical":
print(f"跳过相同帧 {i}")
skipped_frames += 1
continue
# 计算新的拼接高度
new_height = stitched.height + img.height - overlap_height
# 创建新的拼接图像
new_stitched = Image.new('RGB', (stitched.width, new_height))
# 粘贴原有图像
new_stitched.paste(stitched, (0, 0))
# 粘贴新图像（去除重叠部分）
new_img_crop = img.crop((0, overlap_height, img.width, img.height))
new_stitched.paste(new_img_crop, (0, stitched.height))
stitched = new_stitched
print(f"拼接成功，当前高度: {stitched.height}")
else:
print(f"未找到重叠区域，直接拼接帧 {i}")
# 直接拼接
new_height = stitched.height + img.height
new_stitched = Image.new('RGB', (stitched.width, new_height))
new_stitched.paste(stitched, (0, 0))
new_stitched.paste(img, (0, stitched.height))
stitched = new_stitched
print(f"直接拼接，当前高度: {stitched.height}")
except Exception as e:
print(f"处理帧 {i} 时出错: {e}")
skipped_frames += 1
continue
print(f"\n拼接完成！")
print(f"最终尺寸: {stitched.size}")
print(f"跳过帧数: {skipped_frames}")
return stitched
def save_stitched_image(self, stitched_img, output_path="long_article.png"):
"""保存拼接后的长图"""
try:
stitched_img.save(output_path, 'PNG', optimize=True)
print(f"长图已保存: {output_path}")
return True
except Exception as e:
print(f"保存失败: {e}")
return False