app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import tempfile
import os

def process_video_with_mask(video_path, mask_path, text="Joker"):
    # 创建临时输出文件
    output_path = os.path.join(tempfile.gettempdir(), "output.mp4")
    
    # 打开视频和mask视频
    cap_video = cv2.VideoCapture(video_path)
    cap_mask = cv2.VideoCapture(mask_path)
    
    # 获取视频参数
    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))
    
    # 第二阶段：使用平均角度处理视频
    while True:
        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)
        
        # 转换为PIL图像处理文字
        pil_image = Image.fromarray(cv2.cvtColor(frame_video, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)
        
        if contours:
            # 找到面积最大的轮廓
            max_contour = max(contours, key=cv2.contourArea)
            
            # 获取旋转矩形
            rect = cv2.minAreaRect(max_contour)
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            
            # 计算中心点和尺寸
            center = rect[0]
            size = rect[1]
            
            # 创建透明画布绘制文字
            text_image = Image.new("RGBA", (int(size[0]), int(size[1])), (0, 0, 0, 0))
            text_draw = ImageDraw.Draw(text_image)
            
            try:
                # 自动调整字体大小
                font_size = int(min(size[0] * 0.8 / max(1, len(text)), size[1] * 0.8))
                font = ImageFont.truetype("华文琥珀.ttf", font_size)
            except:
                font = ImageFont.load_default()
            
            # 计算文字位置
            bbox = text_draw.textbbox((0, 0), text, font=font)
            text_w = bbox[2] - bbox[0]
            text_h = bbox[3] - bbox[1]
            pos_x = (size[0] - text_w) // 2
            pos_y = (size[1] - text_h) // 2
            
            # 绘制文字
            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))
            
            # 计算粘贴位置
            paste_x = int(center[0] - rotated_text.width / 2)
            paste_y = int(center[1] - rotated_text.height / 2)
            
            # 合并到原图
            pil_image.paste(rotated_text, (paste_x, paste_y), rotated_text)
        
        # 转换回OpenCV格式
        frame_processed = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        out.write(frame_processed)
    
    # 释放资源
    cap_video.release()
    cap_mask.release()
    out.release()
    
    return output_path

# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("## 视频文字替换工具")
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="原始视频")
            mask_input = gr.Video(label="Mask视频")
            text_input = gr.Textbox(label="要添加的文字", value="示例文字")
            submit_btn = gr.Button("处理视频")
        with gr.Column():
            video_output = gr.Video(label="处理结果")
    
    submit_btn.click(
        fn=process_video_with_mask,
        inputs=[video_input, mask_input, text_input],
        outputs=video_output,
    )

    gr.Examples(
        [
            ["examples/20250511-013442_3331156457_.mp4", "examples/20250511-013442_3331156457__temp.mp4_alpha_temp.mp4", "今天天气怎么样？"]
        ],
        [video_input, mask_input, text_input]
    )

if __name__ == "__main__":
    demo.launch(share=True)