Video_cutter/utils/caption.py

import os
import json
import subprocess
import time
import ffmpeg
from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
from moviepy.video.tools.subtitles import SubtitlesClip

from .logger_config import setup_logger
from .llm_director import caption_correct

logger = setup_logger(__name__)

def sec_to_srt_time(seconds):
    hours = int(seconds // 3600)
    mins = int((seconds % 3600) // 60)
    secs = seconds % 60
    return f"{hours:02d}:{mins:02d}:{secs:06.3f}".replace('.', ',')

def generate_srt_file(text_sequence, output_path):
    """
    生成SRT字幕文件
    
    参数:
    text_sequence: 字幕列表，格式 [(start_sec, end_sec, text), ...]
    """
    logger.info(f"字幕文件生成。。。")
    with open(output_path, 'w', encoding='utf-8') as f:
        for i, (start, end, text) in enumerate(text_sequence, 1):
            f.write(f"{i}\n")
            f.write(f"{sec_to_srt_time(start)} --> {sec_to_srt_time(end)}\n")
            f.write(f"{caption_correct(text)}\n\n")


def burn_subtitles_bash(input_path, srt_path, output_path):
    cmdLine = f"ffmpeg -i {input_path} -vf subtitles={srt_path} -c:a copy -c:v h264_nvenc -preset fast -cq 23 {output_path}"
    subprocess.call(cmdLine, shell=True)


def burn_subtitles_ffmpeg(input_path, srt_path, output_path, style_config=None):
    """
    通过FFmpeg烧录字幕（硬编码） - 确保音频保留的版本
    """
    logger.info(f"视频加字幕。。。")
    # 默认字幕样式（精简必要参数）
    with open("./config/caption_config.json", 'r', encoding='utf-8') as file:
            caption = json.load(file)

    default_style = caption
    style = style_config or default_style
    
    # 构建字幕过滤器字符串
    subtitle_filter = (
        f"subtitles={srt_path}:force_style='"
        f"Fontname={style['fontname']},"
        f"Fontsize={style['fontsize']},"
        f"PrimaryColour={style['fontcolor']},"
        f"BackColour={style['bordercolor']},"
        f"Outline={style['borderw']},"
        f"Alignment={style['alignment']},"
        f"MarginV={style['margin']}'"
    )
    
    # 使用更简单的命令结构，确保音频流被正确处理
    (
        ffmpeg
        .input(input_path)
        .output(
            output_path,
            vf=subtitle_filter,  # 应用字幕过滤器
            **{
                
                'preset': 'fast',
                'tune': 'fastdecode',
                'c:a': 'copy',  # 尝试直接复制音频
                'threads': '8',
                'y': None
            }
        )
        .global_args('-hwaccel', 'cuda')
        .run()
        
        # 'c:v': 'h264_nvenc',
        # 'cq': 28,
        # 'gpu': '1',
        # 'rc': 'constqp',
    )

if __name__ == "__main__":
    # 生成测试字幕
    text_sequence = [
        (0.0, 10.0, "第一句字幕测试"),
    ]
    generate_srt_file(text_sequence, 'dynamic_subtitles.srt')

    start_time = time.time()
    # 示例用法（硬件加速需配置FFmpeg）
    burn_subtitles_ffmpeg('./output/video_clips/cut_video3.mp4', 'dynamic_subtitles.srt', 'output3.mp4')
    print(f"{time.time() - start_time} secs")



## TODO:
# 1、JSON to SRT
# 2、JSON to CLIPS
# 3、SRT add to CLIPS