Update app.py

app.py

@@ -5,6 +5,11 @@ from PIL import Image, ImageDraw, ImageFont
 import tempfile
 import os
 
+def is_horizontal_angle(angle, threshold=15):
+    """判断角度是否接近水平（0度或180度）"""
+    angle = angle % 180  # 将角度归一化到0-180度范围
+    return abs(angle) < threshold or abs(angle - 180) < threshold
+
 def process_video_with_mask(video_path, mask_path, text="Joker"):
     # 创建临时输出文件
     output_path = os.path.join(tempfile.gettempdir(), "output.mp4")
@@ -17,17 +22,12 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
     fps = cap_video.get(cv2.CAP_PROP_FPS)
     width = int(cap_video.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    total_frames = int(cap_video.get(cv2.CAP_PROP_FRAME_COUNT))
-    
-    # 第一阶段：计算所有帧的平均角度
-    angles = []
-    frame_positions = []
     
-    # 保存当前视频位置
+    # 第一阶段：寻找第一个接近水平的角度
+    fixed_angle = None
     current_pos = cap_video.get(cv2.CAP_PROP_POS_FRAMES)
     
-    # 遍历所有帧计算角度
-    for _ in range(total_frames):
+    while fixed_angle is None:
         ret_video, frame_video = cap_video.read()
         ret_mask, frame_mask = cap_mask.read()
         
@@ -50,21 +50,45 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
             
             # 获取旋转矩形
             rect = cv2.minAreaRect(max_contour)
-            angles.append(rect[2])
-            frame_positions.append(cap_video.get(cv2.CAP_PROP_POS_FRAMES))
+            current_angle = rect[2]
+            
+            # 检查是否接近水平
+            if is_horizontal_angle(current_angle):
+                fixed_angle = current_angle
+    
+    # 如果没有找到接近水平的角度，使用第一个找到的角度
+    if fixed_angle is None:
+        cap_video.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
+        cap_mask.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
+        
+        ret_video, frame_video = cap_video.read()
+        ret_mask, frame_mask = cap_mask.read()
+        
+        if ret_video and ret_mask:
+            if len(frame_mask.shape) == 3:
+                frame_mask = cv2.cvtColor(frame_mask, cv2.COLOR_BGR2GRAY)
+            
+            _, binary = cv2.threshold(frame_mask, 200, 255, cv2.THRESH_BINARY)
+            contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            
+            if contours:
+                max_contour = max(contours, key=cv2.contourArea)
+                rect = cv2.minAreaRect(max_contour)
+                fixed_angle = rect[2]
     
-    # 计算平均角度
-    mean_angle = np.mean(angles) if angles else 0
+    # 如果还是没有角度，默认使用0度
+    if fixed_angle is None:
+        fixed_angle = 0
     
     # 重置视频位置
-    cap_video.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
-    cap_mask.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
+    cap_video.set(cv2.CAP_PROP_POS_FRAMES, 0)
+    cap_mask.set(cv2.CAP_PROP_POS_FRAMES, 0)
     
     # 创建视频写入器
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
     
-    # 第二阶段：使用平均角度处理视频
+    # 第二阶段：使用固定角度处理视频
     while True:
         ret_video, frame_video = cap_video.read()
         ret_mask, frame_mask = cap_mask.read()
@@ -120,8 +144,8 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
             # 绘制文字
             text_draw.text((pos_x, pos_y), text, fill=(255, 255, 0, 255), font=font)
             
-            # 使用平均角度旋转文字
-            rotated_text = text_image.rotate(mean_angle, expand=True, center=(size[0]/2, size[1]/2))
+            # 使用固定角度旋转文字
+            rotated_text = text_image.rotate(fixed_angle, expand=True, center=(size[0]/2, size[1]/2))
             
             # 计算粘贴位置
             paste_x = int(center[0] - rotated_text.width / 2)