svjack commited on
Commit
6662c85
·
verified ·
1 Parent(s): b2b3e1d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -10
app.py CHANGED
@@ -17,14 +17,54 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
17
  fps = cap_video.get(cv2.CAP_PROP_FPS)
18
  width = int(cap_video.get(cv2.CAP_PROP_FRAME_WIDTH))
19
  height = int(cap_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
  # 创建视频写入器
22
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
23
  out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
24
 
25
- # 初始化固定角度为None
26
- fixed_angle = None
27
-
28
  while True:
29
  ret_video, frame_video = cap_video.read()
30
  ret_mask, frame_mask = cap_mask.read()
@@ -58,11 +98,6 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
58
  # 计算中心点和尺寸
59
  center = rect[0]
60
  size = rect[1]
61
- current_angle = rect[2]
62
-
63
- # 如果是第一帧,记录角度
64
- if fixed_angle is None:
65
- fixed_angle = current_angle
66
 
67
  # 创建透明画布绘制文字
68
  text_image = Image.new("RGBA", (int(size[0]), int(size[1])), (0, 0, 0, 0))
@@ -85,8 +120,8 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
85
  # 绘制文字
86
  text_draw.text((pos_x, pos_y), text, fill=(255, 255, 0, 255), font=font)
87
 
88
- # 使用固定角度旋转文字
89
- rotated_text = text_image.rotate(fixed_angle, expand=True, center=(size[0]/2, size[1]/2))
90
 
91
  # 计算粘贴位置
92
  paste_x = int(center[0] - rotated_text.width / 2)
 
17
  fps = cap_video.get(cv2.CAP_PROP_FPS)
18
  width = int(cap_video.get(cv2.CAP_PROP_FRAME_WIDTH))
19
  height = int(cap_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
20
+ total_frames = int(cap_video.get(cv2.CAP_PROP_FRAME_COUNT))
21
+
22
+ # 第一阶段:计算所有帧的平均角度
23
+ angles = []
24
+ frame_positions = []
25
+
26
+ # 保存当前视频位置
27
+ current_pos = cap_video.get(cv2.CAP_PROP_POS_FRAMES)
28
+
29
+ # 遍历所有帧计算角度
30
+ for _ in range(total_frames):
31
+ ret_video, frame_video = cap_video.read()
32
+ ret_mask, frame_mask = cap_mask.read()
33
+
34
+ if not ret_video or not ret_mask:
35
+ break
36
+
37
+ # 处理mask为单通道
38
+ if len(frame_mask.shape) == 3:
39
+ frame_mask = cv2.cvtColor(frame_mask, cv2.COLOR_BGR2GRAY)
40
+
41
+ # 找到mask中的白色区域
42
+ _, binary = cv2.threshold(frame_mask, 200, 255, cv2.THRESH_BINARY)
43
+
44
+ # 查找轮廓
45
+ contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
46
+
47
+ if contours:
48
+ # 找到面积最大的轮廓
49
+ max_contour = max(contours, key=cv2.contourArea)
50
+
51
+ # 获取旋转矩形
52
+ rect = cv2.minAreaRect(max_contour)
53
+ angles.append(rect[2])
54
+ frame_positions.append(cap_video.get(cv2.CAP_PROP_POS_FRAMES))
55
+
56
+ # 计算平均角度
57
+ mean_angle = np.mean(angles) if angles else 0
58
+
59
+ # 重置视频位置
60
+ cap_video.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
61
+ cap_mask.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
62
 
63
  # 创建视频写入器
64
  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
65
  out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
66
 
67
+ # 第二阶段:使用平均角度处理视频
 
 
68
  while True:
69
  ret_video, frame_video = cap_video.read()
70
  ret_mask, frame_mask = cap_mask.read()
 
98
  # 计算中心点和尺寸
99
  center = rect[0]
100
  size = rect[1]
 
 
 
 
 
101
 
102
  # 创建透明画布绘制文字
103
  text_image = Image.new("RGBA", (int(size[0]), int(size[1])), (0, 0, 0, 0))
 
120
  # 绘制文字
121
  text_draw.text((pos_x, pos_y), text, fill=(255, 255, 0, 255), font=font)
122
 
123
+ # 使用平均角度旋转文字
124
+ rotated_text = text_image.rotate(mean_angle, expand=True, center=(size[0]/2, size[1]/2))
125
 
126
  # 计算粘贴位置
127
  paste_x = int(center[0] - rotated_text.width / 2)