Update app.py

app.py

@@ -17,14 +17,54 @@ 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 = []
+    
+    # 保存当前视频位置
+    current_pos = cap_video.get(cv2.CAP_PROP_POS_FRAMES)
+    
+    # 遍历所有帧计算角度
+    for _ in range(total_frames):
+        ret_video, frame_video = cap_video.read()
+        ret_mask, frame_mask = cap_mask.read()
+        
+        if not ret_video or not ret_mask:
+            break
+            
+        # 处理mask为单通道
+        if len(frame_mask.shape) == 3:
+            frame_mask = cv2.cvtColor(frame_mask, cv2.COLOR_BGR2GRAY)
+        
+        # 找到mask中的白色区域
+        _, 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)
+            angles.append(rect[2])
+            frame_positions.append(cap_video.get(cv2.CAP_PROP_POS_FRAMES))
+    
+    # 计算平均角度
+    mean_angle = np.mean(angles) if angles else 0
+    
+    # 重置视频位置
+    cap_video.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
+    cap_mask.set(cv2.CAP_PROP_POS_FRAMES, current_pos)
     
     # 创建视频写入器
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
     
-    # 初始化固定角度为None
-    fixed_angle = None
-    
+    # 第二阶段：使用平均角度处理视频
     while True:
         ret_video, frame_video = cap_video.read()
         ret_mask, frame_mask = cap_mask.read()
@@ -58,11 +98,6 @@ def process_video_with_mask(video_path, mask_path, text="Joker"):
             # 计算中心点和尺寸
             center = rect[0]
             size = rect[1]
-            current_angle = rect[2]
-            
-            # 如果是第一帧，记录角度
-            if fixed_angle is None:
-                fixed_angle = current_angle
             
             # 创建透明画布绘制文字
             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"):
             # 绘制文字
             text_draw.text((pos_x, pos_y), text, fill=(255, 255, 0, 255), font=font)
             
-            # 使用固定角度旋转文字
-            rotated_text = text_image.rotate(fixed_angle, expand=True, center=(size[0]/2, size[1]/2))
+            # 使用平均角度旋转文字
+            rotated_text = text_image.rotate(mean_angle, expand=True, center=(size[0]/2, size[1]/2))
             
             # 计算粘贴位置
             paste_x = int(center[0] - rotated_text.width / 2)