From c1a131ec88332b85cded19592ece9f7b9ffcf9e0 Mon Sep 17 00:00:00 2001 From: Unizzr <1920158711@qq.com> Date: Sat, 2 May 2026 10:13:13 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BA=E4=BE=9D?= =?UTF-8?q?=E8=B5=96=E6=A3=80=E6=9F=A5=E5=92=8C=E9=94=99=E8=AF=AF=E5=A4=84?= =?UTF-8?q?=E7=90=86=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 添加依赖检查模块(dependency_checker.py)，改进requirements.txt文档，增强主程序(main.py)和音频处理器(audio_processor.py)的错误处理机制 - 新增dependency_checker.py实现自动依赖检查功能 - 重构requirements.txt提供更清晰的安装指南 - 改进main.py添加友好的错误提示和解决方案 - 增强audio_processor.py的FFmpeg查找和错误处理逻辑 - 优化用户界面，提供更详细的进度和状态信息 --- audio_processor.py | 722 +++++++++++++++++++++++++++++++++++------- dependency_checker.py | 569 +++++++++++++++++++++++++++++++++ main.py | 400 ++++++++++++++++++----- requirements.txt | 100 +++++- 4 files changed, 1586 insertions(+), 205 deletions(-) create mode 100644 dependency_checker.py diff --git a/audio_processor.py b/audio_processor.py index 2a013c9..b8a2289 100644 --- a/audio_processor.py +++ b/audio_processor.py @@ -1,18 +1,48 @@ """ 音频处理模块 支持从MP4、MOV、AVI等视频文件中提取音频，以及处理MP3、WAV等音频文件 +改进版：增强FFmpeg查找逻辑，提供更友好的错误信息和解决方案 """ import os +import sys import subprocess import tempfile +import shutil from pathlib import Path -from typing import Optional, Tuple +from typing import Optional, Tuple, List import logging logger = logging.getLogger(__name__) +class FFmpegNotFoundError(Exception): + """FFmpeg未找到错误""" + + def __init__(self, message: str, solutions: List[str] = None): + super().__init__(message) + self.solutions = solutions or [] + + def get_detailed_message(self) -> str: + """获取详细的错误信息""" + msg = str(self) + if self.solutions: + msg += "\n\n解决方案:\n" + for i, solution in enumerate(self.solutions, 1): + msg += f" {i}. {solution}\n" + return msg + + +class FFmpegError(Exception): + """FFmpeg执行错误""" + + def __init__(self, command: List[str], return_code: int, stderr: str): + self.command = command + self.return_code = return_code + self.stderr = stderr + super().__init__(f"FFmpeg执行失败 (返回码: {return_code})") + + class AudioProcessor: """ 音频处理器类 @@ -20,34 +50,177 @@ class AudioProcessor: """ # 支持的视频格式 - SUPPORTED_VIDEO_FORMATS = {'.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm'} + SUPPORTED_VIDEO_FORMATS = {'.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm', '.m4v', '.mpeg', '.mpg', '.3gp'} # 支持的音频格式 - SUPPORTED_AUDIO_FORMATS = {'.mp3', '.wav', '.m4a', '.flac', '.aac', '.ogg', '.wma'} + SUPPORTED_AUDIO_FORMATS = {'.mp3', '.wav', '.m4a', '.flac', '.aac', '.ogg', '.wma', '.opus', '.amr'} + + # FFmpeg详细安装指南 + FFMPEG_INSTALL_GUIDE = { + 'windows': [ + "方法1: 使用 winget 安装 (推荐)", + " 以管理员身份打开 PowerShell，运行:", + " winget install Gyan.FFmpeg", + "", + "方法2: 使用 Chocolatey 安装", + " choco install ffmpeg", + "", + "方法3: 手动安装", + " 1. 访问 https://www.gyan.dev/ffmpeg/builds/", + " 2. 下载 'ffmpeg-release-full.7z'", + " 3. 解压到 C:\\ffmpeg", + " 4. 将 C:\\ffmpeg\\bin 添加到系统 PATH 环境变量", + "", + "安装后需要重启命令行/IDE才能生效！" + ], + 'macos': [ + "方法1: 使用 Homebrew 安装 (推荐)", + " brew install ffmpeg", + "", + "方法2: 使用 MacPorts 安装", + " sudo port install ffmpeg", + "", + "方法3: 手动安装", + " 从 https://evermeet.cx/ffmpeg/ 下载并安装" + ], + 'linux': [ + "Ubuntu/Debian:", + " sudo apt update", + " sudo apt install ffmpeg", + "", + "Fedora/RHEL:", + " sudo dnf install ffmpeg", + "", + "Arch Linux:", + " sudo pacman -S ffmpeg" + ] + } - def __init__(self, ffmpeg_path: Optional[str] = None): + def __init__(self, ffmpeg_path: Optional[str] = None, auto_download: bool = False): """ 初始化音频处理器 Args: ffmpeg_path: FFmpeg可执行文件路径，如为None则尝试从PATH中查找 + auto_download: 是否自动下载FFmpeg（预留功能） """ - self.ffmpeg_path = ffmpeg_path or self._find_ffmpeg() - if not self.ffmpeg_path: - raise RuntimeError("未找到FFmpeg，请确保FFmpeg已安装并添加到系统PATH中") + self.ffmpeg_path = None + self.ffprobe_path = None + self._ffmpeg_version = None + + # 尝试查找 FFmpeg + if ffmpeg_path: + # 用户指定了路径 + self.ffmpeg_path = self._validate_ffmpeg_path(ffmpeg_path) + if not self.ffmpeg_path: + raise FFmpegNotFoundError( + f"指定的FFmpeg路径无效: {ffmpeg_path}", + self._get_install_solutions() + ) + else: + # 自动查找 + self.ffmpeg_path = self._find_ffmpeg() + if not self.ffmpeg_path: + raise FFmpegNotFoundError( + "未找到FFmpeg！FFmpeg是音视频处理必需的工具。", + self._get_install_solutions() + ) + + # 查找 ffprobe + self.ffprobe_path = self._find_ffprobe() + + # 验证 FFmpeg 是否可用 + self._validate_ffmpeg() + + logger.info(f"FFmpeg初始化成功: {self.ffmpeg_path}") + logger.info(f"FFmpeg版本: {self._ffmpeg_version}") + + def _get_platform(self) -> str: + """ + 获取当前平台 + + Returns: + 'windows', 'macos', 或 'linux' + """ + if sys.platform == 'win32': + return 'windows' + elif sys.platform == 'darwin': + return 'macos' + else: + return 'linux' + + def _get_install_solutions(self) -> List[str]: + """ + 获取安装解决方案列表 + + Returns: + 解决方案列表 + """ + platform = self._get_platform() + return self.FFMPEG_INSTALL_GUIDE.get(platform, []) + + def _validate_ffmpeg_path(self, path: str) -> Optional[str]: + """ + 验证 FFmpeg 路径 + + Args: + path: FFmpeg 可执行文件或目录路径 + + Returns: + 验证后的完整路径或 None + """ + path = Path(path).resolve() + + if path.is_file(): + # 是文件路径 + if path.name.lower() in ['ffmpeg', 'ffmpeg.exe']: + if self._test_ffmpeg_executable(str(path)): + return str(path) + elif path.is_dir(): + # 是目录路径 + ffmpeg_exe = 'ffmpeg.exe' if sys.platform == 'win32' else 'ffmpeg' + candidate = path / ffmpeg_exe + if candidate.exists() and self._test_ffmpeg_executable(str(candidate)): + return str(candidate) + + return None + + def _test_ffmpeg_executable(self, path: str) -> bool: + """ + 测试 FFmpeg 可执行文件是否可用 + + Args: + path: FFmpeg 可执行文件路径 + + Returns: + 是否可用 + """ + try: + result = subprocess.run( + [path, '-version'], + capture_output=True, + text=True, + timeout=10 + ) + return result.returncode == 0 + except (subprocess.SubprocessError, FileNotFoundError, OSError): + return False def _find_ffmpeg(self) -> Optional[str]: """ 在系统中查找FFmpeg可执行文件 查找顺序： - 1. 环境变量 FFMPEG_PATH 或 FFMPEG_BIN + 1. 环境变量 FFMPEG_PATH, FFMPEG_BIN, FFMPEG 2. 常见的FFmpeg安装位置 3. 系统PATH环境变量 + 4. 项目目录下的 ffmpeg 文件夹 Returns: FFmpeg路径或None """ - ffmpeg_exe = 'ffmpeg.exe' if os.name == 'nt' else 'ffmpeg' + ffmpeg_exe = 'ffmpeg.exe' if sys.platform == 'win32' else 'ffmpeg' + + logger.info("正在查找FFmpeg...") # 1. 检查环境变量 env_paths = [ @@ -58,60 +231,225 @@ def _find_ffmpeg(self) -> Optional[str]: for env_path in env_paths: if env_path: - env_path = os.path.expanduser(env_path) - # 如果是目录，拼接ffmpeg.exe - if os.path.isdir(env_path): - candidate = os.path.join(env_path, ffmpeg_exe) - if os.path.exists(candidate): - return candidate - # 如果是文件路径 - elif os.path.exists(env_path): - return env_path - - # 2. 检查常见的FFmpeg安装位置（Windows） - if os.name == 'nt': - common_paths = [ + validated = self._validate_ffmpeg_path(env_path) + if validated: + logger.info(f"从环境变量找到FFmpeg: {validated}") + return validated + + # 2. 检查常见的FFmpeg安装位置 + common_locations = self._get_common_locations() + + for location in common_locations: + if location and os.path.exists(location): + validated = self._validate_ffmpeg_path(location) + if validated: + logger.info(f"从常见位置找到FFmpeg: {validated}") + return validated + + # 3. 检查项目目录下的 ffmpeg 文件夹 + project_ffmpeg = Path(__file__).parent / 'ffmpeg' + if project_ffmpeg.exists(): + validated = self._validate_ffmpeg_path(str(project_ffmpeg)) + if validated: + logger.info(f"从项目目录找到FFmpeg: {validated}") + return validated + + # 4. 检查系统PATH环境变量 + path_dirs = os.environ.get('PATH', '').split(os.pathsep) + + for path_dir in path_dirs: + if not path_dir: + continue + try: + path_dir = Path(path_dir).resolve() + if path_dir.exists(): + ffmpeg_file = path_dir / ffmpeg_exe + if ffmpeg_file.exists(): + if self._test_ffmpeg_executable(str(ffmpeg_file)): + logger.info(f"从系统PATH找到FFmpeg: {ffmpeg_file}") + return str(ffmpeg_file) + except Exception: + continue + + # 5. 尝试使用 shutil.which + which_result = shutil.which('ffmpeg') + if which_result: + if self._test_ffmpeg_executable(which_result): + logger.info(f"从shutil.which找到FFmpeg: {which_result}") + return which_result + + logger.warning("未找到FFmpeg") + return None + + def _get_common_locations(self) -> List[str]: + """ + 获取常见的FFmpeg安装位置 + + Returns: + 路径列表 + """ + locations = [] + + if sys.platform == 'win32': + # Windows 常见位置 + user_profile = os.environ.get('USERPROFILE', '') + local_app_data = os.environ.get('LOCALAPPDATA', '') + program_files = os.environ.get('PROGRAMFILES', '') + program_files_x86 = os.environ.get('PROGRAMFILES(X86)', '') + + locations.extend([ r'C:\ffmpeg\bin', - r'C:\Program Files\ffmpeg\bin', - r'C:\Program Files (x86)\ffmpeg\bin', - os.path.join(os.environ.get('LOCALAPPDATA', ''), 'ffmpeg', 'bin'), - os.path.join(os.environ.get('USERPROFILE', ''), 'ffmpeg', 'bin'), - os.path.join(os.environ.get('USERPROFILE', ''), 'scoop', 'shims'), - os.path.join(os.environ.get('USERPROFILE', ''), 'scoop', 'apps', 'ffmpeg', 'current', 'bin'), + r'C:\ffmpeg', + os.path.join(program_files, 'ffmpeg', 'bin'), + os.path.join(program_files, 'ffmpeg'), + os.path.join(program_files_x86, 'ffmpeg', 'bin'), + os.path.join(program_files_x86, 'ffmpeg'), + os.path.join(local_app_data, 'ffmpeg', 'bin'), + os.path.join(local_app_data, 'ffmpeg'), + os.path.join(user_profile, 'ffmpeg', 'bin'), + os.path.join(user_profile, 'ffmpeg'), + os.path.join(user_profile, 'scoop', 'shims'), + os.path.join(user_profile, 'scoop', 'apps', 'ffmpeg', 'current', 'bin'), r'C:\ProgramData\chocolatey\bin', - ] - - for path in common_paths: - candidate = os.path.join(path, ffmpeg_exe) - if os.path.exists(candidate): - return candidate + r'C:\tools\ffmpeg\bin', + r'C:\tools\ffmpeg', + ]) - # 检查 WinGet 安装的 FFmpeg（Gyan.FFmpeg） - winget_packages = os.path.join( - os.environ.get('LOCALAPPDATA', ''), - 'Microsoft', 'WinGet', 'Packages' - ) + # WinGet 包位置 + winget_packages = os.path.join(local_app_data, 'Microsoft', 'WinGet', 'Packages') if os.path.exists(winget_packages): try: - # 查找包含 ffmpeg 的文件夹 for item in os.listdir(winget_packages): item_path = os.path.join(winget_packages, item) if os.path.isdir(item_path) and 'ffmpeg' in item.lower(): - # 在该目录下查找 ffmpeg.exe + locations.append(item_path) + # 也检查子目录 for root, dirs, files in os.walk(item_path): - if ffmpeg_exe in files: - return os.path.join(root, ffmpeg_exe) + if 'bin' in dirs: + locations.append(os.path.join(root, 'bin')) except Exception: pass - # 3. 检查系统PATH环境变量 - for path in os.environ.get('PATH', '').split(os.pathsep): - ffmpeg_full_path = os.path.join(path, ffmpeg_exe) - if os.path.exists(ffmpeg_full_path): - return ffmpeg_full_path + elif sys.platform == 'darwin': + # macOS 常见位置 + locations.extend([ + '/usr/local/bin', + '/usr/local/ffmpeg/bin', + '/opt/homebrew/bin', + '/opt/homebrew/ffmpeg/bin', + '/Applications/ffmpeg', + os.path.expanduser('~/Applications/ffmpeg'), + ]) + + else: + # Linux 常见位置 + locations.extend([ + '/usr/bin', + '/usr/local/bin', + '/opt/ffmpeg/bin', + '/snap/bin', + ]) + + return locations + + def _find_ffprobe(self) -> Optional[str]: + """ + 查找 ffprobe 工具 + + Returns: + ffprobe路径或None + """ + if not self.ffmpeg_path: + return None + + ffprobe_exe = 'ffprobe.exe' if sys.platform == 'win32' else 'ffprobe' + + # 首先在 FFmpeg 同一目录查找 + ffmpeg_dir = Path(self.ffmpeg_path).parent + ffprobe_candidate = ffmpeg_dir / ffprobe_exe + + if ffprobe_candidate.exists(): + try: + result = subprocess.run( + [str(ffprobe_candidate), '-version'], + capture_output=True, + text=True, + timeout=5 + ) + if result.returncode == 0: + return str(ffprobe_candidate) + except Exception: + pass + + # 在系统 PATH 中查找 + which_result = shutil.which('ffprobe') + if which_result: + return which_result return None + def _validate_ffmpeg(self): + """ + 验证 FFmpeg 是否可用并获取版本信息 + """ + try: + result = subprocess.run( + [self.ffmpeg_path, '-version'], + capture_output=True, + text=True, + timeout=10 + ) + + if result.returncode == 0: + # 解析版本信息 + first_line = result.stdout.split('\n')[0] if result.stdout else '' + # 典型输出: "ffmpeg version 6.0-full_build-www.gyan.dev Copyright (c) 2000-2023 the FFmpeg developers" + import re + version_match = re.search(r'ffmpeg version (\S+)', first_line, re.IGNORECASE) + if version_match: + self._ffmpeg_version = version_match.group(1) + else: + self._ffmpeg_version = 'unknown' + else: + raise FFmpegError( + [self.ffmpeg_path, '-version'], + result.returncode, + result.stderr + ) + + except subprocess.TimeoutExpired: + raise FFmpegNotFoundError( + "FFmpeg响应超时", + ["请检查FFmpeg是否正确安装", "尝试重新安装FFmpeg"] + ) + except FileNotFoundError as e: + raise FFmpegNotFoundError( + f"FFmpeg可执行文件不存在: {self.ffmpeg_path}", + self._get_install_solutions() + ) from e + except OSError as e: + if 'WinError 2' in str(e) or 'The system cannot find the file specified' in str(e): + raise FFmpegNotFoundError( + f"无法执行FFmpeg: {e}\n路径: {self.ffmpeg_path}", + self._get_install_solutions() + ) from e + else: + raise + + def get_ffmpeg_info(self) -> dict: + """ + 获取 FFmpeg 信息 + + Returns: + 包含FFmpeg信息的字典 + """ + return { + 'ffmpeg_path': self.ffmpeg_path, + 'ffprobe_path': self.ffprobe_path, + 'version': self._ffmpeg_version, + 'platform': self._get_platform(), + } + def is_supported_format(self, file_path: str) -> bool: """ 检查文件是否为支持的音视频格式 @@ -138,6 +476,71 @@ def is_video_file(self, file_path: str) -> bool: ext = Path(file_path).suffix.lower() return ext in self.SUPPORTED_VIDEO_FORMATS + def _run_ffmpeg_command( + self, + cmd: List[str], + description: str = "执行FFmpeg命令", + timeout: Optional[int] = None, + ) -> Tuple[int, str, str]: + """ + 执行 FFmpeg 命令 + + Args: + cmd: 命令参数列表 + description: 命令描述（用于错误信息） + timeout: 超时时间（秒） + + Returns: + (返回码, 标准输出, 标准错误) + """ + cmd = [self.ffmpeg_path] + cmd + + logger.debug(f"执行FFmpeg命令: {' '.join(cmd)}") + + try: + # 使用 startupinfo 隐藏 Windows 控制台窗口 + startupinfo = None + if sys.platform == 'win32': + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + startupinfo.wShowWindow = subprocess.SW_HIDE + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout, + startupinfo=startupinfo, + ) + + if result.returncode != 0: + logger.error(f"FFmpeg命令失败: {description}") + logger.error(f"命令: {' '.join(cmd)}") + logger.error(f"错误输出: {result.stderr}") + + raise FFmpegError(cmd, result.returncode, result.stderr) + + return result.returncode, result.stdout, result.stderr + + except FileNotFoundError as e: + raise FFmpegNotFoundError( + f"执行FFmpeg时找不到文件: {e}\n命令: {' '.join(cmd)}", + self._get_install_solutions() + ) from e + except OSError as e: + if 'WinError 2' in str(e) or 'The system cannot find the file specified' in str(e): + raise FFmpegNotFoundError( + f"无法执行FFmpeg: {e}\n" + f"FFmpeg路径: {self.ffmpeg_path}\n" + f"命令: {' '.join(cmd)}", + self._get_install_solutions() + ) from e + else: + raise + except subprocess.TimeoutExpired: + logger.error(f"FFmpeg命令超时: {description}") + raise RuntimeError(f"FFmpeg命令超时: {description}") + def extract_audio_from_video( self, video_path: str, @@ -156,9 +559,6 @@ def extract_audio_from_video( Returns: 提取的音频文件路径 - - Raises: - RuntimeError: 提取失败时抛出 """ video_path = Path(video_path).resolve() if not video_path.exists(): @@ -171,33 +571,26 @@ def extract_audio_from_video( output_path = Path(output_path).resolve() + logger.info(f"正在从视频提取音频: {video_path.name} -> {output_path.name}") + # 构建FFmpeg命令 + # -vn: 不处理视频 + # -acodec pcm_s16le: PCM 16位小端编码 + # -ar: 采样率 + # -ac: 声道数 + # -y: 覆盖输出文件 cmd = [ - self.ffmpeg_path, '-i', str(video_path), - '-vn', # 只处理音频 - '-acodec', 'pcm_s16le', # PCM编码，16位 - '-ar', str(sample_rate), # 采样率 - '-ac', str(channels), # 声道数 - '-y', # 覆盖输出文件 + '-vn', + '-acodec', 'pcm_s16le', + '-ar', str(sample_rate), + '-ac', str(channels), + '-y', str(output_path) ] - logger.info(f"正在从视频提取音频: {video_path} -> {output_path}") - logger.info(f"执行命令: {' '.join(cmd)}") - try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False - ) - - if result.returncode != 0: - error_msg = f"音频提取失败: {result.stderr}" - logger.error(error_msg) - raise RuntimeError(error_msg) + self._run_ffmpeg_command(cmd, f"提取音频: {video_path.name}") if not output_path.exists(): raise RuntimeError(f"音频提取失败，输出文件不存在: {output_path}") @@ -205,9 +598,9 @@ def extract_audio_from_video( logger.info(f"音频提取成功: {output_path}") return str(output_path) - except Exception as e: - logger.exception(f"音频提取过程中发生错误: {e}") - raise + except FFmpegError as e: + logger.error(f"音频提取失败: {e}") + raise RuntimeError(f"音频提取失败: {e.stderr}") from e def convert_audio_format( self, @@ -236,6 +629,7 @@ def convert_audio_format( if (input_path.suffix.lower() == '.wav' and sample_rate == 16000 and channels == 1): + logger.info(f"音频已是所需格式，跳过转换: {input_path}") return str(input_path) # 如果未指定输出路径，创建临时文件 @@ -245,9 +639,10 @@ def convert_audio_format( output_path = Path(output_path).resolve() + logger.info(f"正在转换音频: {input_path.name} -> {output_path.name}") + # 构建FFmpeg命令 cmd = [ - self.ffmpeg_path, '-i', str(input_path), '-acodec', 'pcm_s16le', '-ar', str(sample_rate), @@ -256,27 +651,15 @@ def convert_audio_format( str(output_path) ] - logger.info(f"正在转换音频: {input_path} -> {output_path}") - try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False - ) - - if result.returncode != 0: - error_msg = f"音频转换失败: {result.stderr}" - logger.error(error_msg) - raise RuntimeError(error_msg) + self._run_ffmpeg_command(cmd, f"转换音频: {input_path.name}") logger.info(f"音频转换成功: {output_path}") return str(output_path) - except Exception as e: - logger.exception(f"音频转换过程中发生错误: {e}") - raise + except FFmpegError as e: + logger.error(f"音频转换失败: {e}") + raise RuntimeError(f"音频转换失败: {e.stderr}") from e def get_audio_duration(self, audio_path: str) -> float: """ @@ -292,34 +675,66 @@ def get_audio_duration(self, audio_path: str) -> float: if not audio_path.exists(): raise FileNotFoundError(f"音频文件不存在: {audio_path}") - # 使用ffprobe获取时长 - ffprobe_path = self.ffmpeg_path.replace('ffmpeg', 'ffprobe') - cmd = [ - ffprobe_path, - '-v', 'error', - '-show_entries', 'format=duration', - '-of', 'default=noprint_wrappers=1:nokey=1', - str(audio_path) - ] - + # 优先使用 ffprobe + if self.ffprobe_path: + try: + cmd = [ + self.ffprobe_path, + '-v', 'error', + '-show_entries', 'format=duration', + '-of', 'default=noprint_wrappers=1:nokey=1', + str(audio_path) + ] + + # 使用 startupinfo 隐藏 Windows 控制台窗口 + startupinfo = None + if sys.platform == 'win32': + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + startupinfo.wShowWindow = subprocess.SW_HIDE + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=10, + startupinfo=startupinfo, + ) + + if result.returncode == 0: + duration_str = result.stdout.strip() + if duration_str: + return float(duration_str) + except Exception as e: + logger.warning(f"使用ffprobe获取时长失败: {e}，将使用备用方法") + + # 备用方法：使用 ffmpeg 解析 try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False - ) + cmd = [ + '-i', str(audio_path), + '-f', 'null', + '-' + ] - if result.returncode == 0: - duration = float(result.stdout.strip()) - return duration - else: - logger.warning(f"获取音频时长失败: {result.stderr}") - return 0.0 + _, _, stderr = self._run_ffmpeg_command(cmd, f"获取音频时长: {audio_path.name}") + + # 解析输出中的 Duration + import re + match = re.search(r'Duration:\s*(\d{2}):(\d{2}):(\d{2})\.(\d+)', stderr) + if match: + hours = int(match.group(1)) + minutes = int(match.group(2)) + seconds = int(match.group(3)) + milliseconds = int(match.group(4)) + + total_seconds = hours * 3600 + minutes * 60 + seconds + milliseconds / 100.0 + return total_seconds except Exception as e: - logger.exception(f"获取音频时长时发生错误: {e}") + logger.error(f"获取音频时长失败: {e}") return 0.0 + + return 0.0 def process_media_file( self, @@ -342,11 +757,18 @@ def process_media_file( raise FileNotFoundError(f"媒体文件不存在: {media_path}") if not self.is_supported_format(str(media_path)): - raise ValueError(f"不支持的文件格式: {media_path.suffix}") + raise ValueError( + f"不支持的文件格式: {media_path.suffix}\n" + f"支持的视频格式: {', '.join(sorted(self.SUPPORTED_VIDEO_FORMATS))}\n" + f"支持的音频格式: {', '.join(sorted(self.SUPPORTED_AUDIO_FORMATS))}" + ) is_video = self.is_video_file(str(media_path)) is_temp = temp_output + logger.info(f"处理媒体文件: {media_path.name}") + logger.info(f"文件类型: {'视频' if is_video else '音频'}") + if is_video: # 从视频提取音频 audio_path = self.extract_audio_from_video(str(media_path)) @@ -369,3 +791,73 @@ def cleanup_temp_file(self, file_path: str): logger.info(f"已清理临时文件: {file_path}") except Exception as e: logger.warning(f"清理临时文件失败: {file_path}, 错误: {e}") + + def get_media_info(self, media_path: str) -> dict: + """ + 获取媒体文件信息 + + Args: + media_path: 媒体文件路径 + + Returns: + 包含媒体信息的字典 + """ + media_path = Path(media_path).resolve() + if not media_path.exists(): + raise FileNotFoundError(f"媒体文件不存在: {media_path}") + + info = { + 'path': str(media_path), + 'name': media_path.name, + 'extension': media_path.suffix.lower(), + 'is_video': self.is_video_file(str(media_path)), + 'is_audio': not self.is_video_file(str(media_path)) and self.is_supported_format(str(media_path)), + 'size_bytes': media_path.stat().st_size, + 'size_mb': round(media_path.stat().st_size / (1024 * 1024), 2), + } + + # 尝试获取时长 + try: + duration = self.get_audio_duration(str(media_path)) + info['duration_seconds'] = duration + info['duration_formatted'] = self._format_duration(duration) + except Exception: + pass + + return info + + def _format_duration(self, seconds: float) -> str: + """ + 格式化时长 + + Args: + seconds: 秒数 + + Returns: + 格式化的时长字符串 + """ + hours = int(seconds // 3600) + minutes = int((seconds % 3600) // 60) + secs = int(seconds % 60) + + if hours > 0: + return f"{hours}:{minutes:02d}:{secs:02d}" + else: + return f"{minutes}:{secs:02d}" + + +# 便捷函数：检查FFmpeg是否可用 +def check_ffmpeg_available() -> Tuple[bool, Optional[str]]: + """ + 检查FFmpeg是否可用 + + Returns: + (是否可用, 错误信息或None) + """ + try: + processor = AudioProcessor() + return True, None + except FFmpegNotFoundError as e: + return False, e.get_detailed_message() + except Exception as e: + return False, str(e) diff --git a/dependency_checker.py b/dependency_checker.py new file mode 100644 index 0000000..42ecd50 --- /dev/null +++ b/dependency_checker.py @@ -0,0 +1,569 @@ +""" +依赖检查模块 +在程序启动时自动验证所有必需依赖 +提供详细的错误信息和安装指南 +""" + +import os +import sys +import importlib.util +from typing import Dict, List, Tuple, Optional +from dataclasses import dataclass, field +from enum import Enum + + +class DependencyStatus(Enum): + """依赖状态枚举""" + OK = "ok" + MISSING = "missing" + VERSION_TOO_OLD = "version_too_old" + ERROR = "error" + + +@dataclass +class DependencyInfo: + """依赖信息""" + name: str + display_name: str + required: bool = True + installed: bool = False + version: Optional[str] = None + status: DependencyStatus = DependencyStatus.MISSING + error_message: Optional[str] = None + install_guide: List[str] = field(default_factory=list) + + def is_ok(self) -> bool: + """检查是否满足要求""" + return self.status == DependencyStatus.OK + + +class DependencyChecker: + """ + 依赖检查器 + 验证所有必需的依赖是否已安装 + """ + + # 必需的 Python 包 + REQUIRED_PACKAGES = [ + { + 'name': 'torch', + 'display_name': 'PyTorch', + 'required': True, + 'min_version': '2.0.0', + 'install_guide': [ + "GPU版本 (推荐):", + " pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118", + "", + "CPU版本:", + " pip3 install torch torchvision torchaudio", + ] + }, + { + 'name': 'whisper', + 'display_name': 'OpenAI Whisper', + 'required': True, + 'min_version': None, + 'install_guide': [ + "pip install openai-whisper", + ] + }, + { + 'name': 'numpy', + 'display_name': 'NumPy', + 'required': True, + 'min_version': '1.20.0', + 'install_guide': [ + "pip install numpy", + ] + }, + { + 'name': 'tqdm', + 'display_name': 'tqdm (进度条)', + 'required': False, + 'min_version': None, + 'install_guide': [ + "pip install tqdm", + ] + }, + ] + + # 系统工具 + SYSTEM_TOOLS = [ + { + 'name': 'ffmpeg', + 'display_name': 'FFmpeg', + 'required': True, + 'install_guide_windows': [ + "方法1: 使用 winget 安装 (推荐)", + " 以管理员身份打开 PowerShell，运行:", + " winget install Gyan.FFmpeg", + "", + "方法2: 使用 Chocolatey 安装", + " choco install ffmpeg", + "", + "方法3: 手动安装", + " 1. 访问 https://www.gyan.dev/ffmpeg/builds/", + " 2. 下载 'ffmpeg-release-full.7z'", + " 3. 解压到 C:\\ffmpeg", + " 4. 将 C:\\ffmpeg\\bin 添加到系统 PATH 环境变量", + "", + "⚠️ 安装后需要重启命令行/IDE才能生效！", + ], + 'install_guide_macos': [ + "方法1: 使用 Homebrew 安装 (推荐)", + " brew install ffmpeg", + "", + "方法2: 使用 MacPorts 安装", + " sudo port install ffmpeg", + "", + "方法3: 手动安装", + " 从 https://evermeet.cx/ffmpeg/ 下载并安装", + ], + 'install_guide_linux': [ + "Ubuntu/Debian:", + " sudo apt update", + " sudo apt install ffmpeg", + "", + "Fedora/RHEL:", + " sudo dnf install ffmpeg", + "", + "Arch Linux:", + " sudo pacman -S ffmpeg", + ], + } + ] + + def __init__(self): + """初始化依赖检查器""" + self.packages: Dict[str, DependencyInfo] = {} + self.system_tools: Dict[str, DependencyInfo] = {} + self.python_version_ok: bool = True + self.python_version_info: Optional[Tuple[int, int, int]] = None + self._checked: bool = False + + def _get_platform(self) -> str: + """ + 获取当前平台 + + Returns: + 'windows', 'macos', 或 'linux' + """ + if sys.platform == 'win32': + return 'windows' + elif sys.platform == 'darwin': + return 'macos' + else: + return 'linux' + + def check_python_version(self) -> bool: + """ + 检查 Python 版本 + + Returns: + 版本是否满足要求 + """ + self.python_version_info = sys.version_info[:3] + major, minor, micro = self.python_version_info + + # 要求 Python 3.8 - 3.11 + # Whisper 在 Python 3.12+ 可能有兼容性问题 + if (major, minor) < (3, 8): + self.python_version_ok = False + return False + + if (major, minor) >= (3, 12): + # Python 3.12+ 可能有兼容性问题，但不直接阻止 + self.python_version_ok = True + return True + + self.python_version_ok = True + return True + + def check_package(self, package_info: dict) -> DependencyInfo: + """ + 检查单个 Python 包 + + Args: + package_info: 包信息字典 + + Returns: + 依赖信息对象 + """ + name = package_info['name'] + display_name = package_info['display_name'] + required = package_info['required'] + min_version = package_info.get('min_version') + install_guide = package_info.get('install_guide', []) + + info = DependencyInfo( + name=name, + display_name=display_name, + required=required, + install_guide=install_guide, + ) + + try: + # 尝试导入包 + spec = importlib.util.find_spec(name) + if spec is None: + info.status = DependencyStatus.MISSING + info.installed = False + info.error_message = f"包 {display_name} ({name}) 未安装" + return info + + # 尝试获取版本 + module = importlib.import_module(name) + info.installed = True + + # 获取版本号 + version = None + if hasattr(module, '__version__'): + version = module.__version__ + elif name == 'whisper': + # whisper 可能没有 __version__ + try: + import openai_whisper + if hasattr(openai_whisper, '__version__'): + version = openai_whisper.__version__ + except ImportError: + pass + + info.version = version + + # 检查版本要求 + if min_version and version: + try: + from packaging.version import parse as parse_version + if parse_version(version) < parse_version(min_version): + info.status = DependencyStatus.VERSION_TOO_OLD + info.error_message = f"包 {display_name} 版本 {version} 太旧，要求 >= {min_version}" + return info + except ImportError: + # 如果没有 packaging 模块，跳过版本检查 + pass + + info.status = DependencyStatus.OK + + except ImportError as e: + info.status = DependencyStatus.MISSING + info.installed = False + info.error_message = f"导入包 {display_name} 失败: {e}" + + except Exception as e: + info.status = DependencyStatus.ERROR + info.error_message = f"检查包 {display_name} 时发生错误: {e}" + + return info + + def check_all_packages(self) -> Dict[str, DependencyInfo]: + """ + 检查所有 Python 包 + + Returns: + 包信息字典 + """ + for package_info in self.REQUIRED_PACKAGES: + name = package_info['name'] + self.packages[name] = self.check_package(package_info) + + return self.packages + + def check_ffmpeg(self) -> DependencyInfo: + """ + 检查 FFmpeg + + Returns: + FFmpeg 依赖信息 + """ + tool_info = self.SYSTEM_TOOLS[0] + name = tool_info['name'] + display_name = tool_info['display_name'] + + platform = self._get_platform() + if platform == 'windows': + install_guide = tool_info['install_guide_windows'] + elif platform == 'macos': + install_guide = tool_info['install_guide_macos'] + else: + install_guide = tool_info['install_guide_linux'] + + info = DependencyInfo( + name=name, + display_name=display_name, + required=True, + install_guide=install_guide, + ) + + try: + # 使用 audio_processor 中的检查 + from audio_processor import check_ffmpeg_available, FFmpegNotFoundError + + available, error_msg = check_ffmpeg_available() + + if available: + info.status = DependencyStatus.OK + info.installed = True + + # 尝试获取版本 + try: + from audio_processor import AudioProcessor + processor = AudioProcessor() + ffmpeg_info = processor.get_ffmpeg_info() + info.version = ffmpeg_info.get('version') + except Exception: + pass + else: + info.status = DependencyStatus.MISSING + info.installed = False + info.error_message = error_msg + + except ImportError: + # 如果无法导入 audio_processor，使用 shutil.which 检查 + import shutil + ffmpeg_path = shutil.which('ffmpeg') + + if ffmpeg_path: + # 验证是否可执行 + try: + import subprocess + result = subprocess.run( + [ffmpeg_path, '-version'], + capture_output=True, + text=True, + timeout=10 + ) + if result.returncode == 0: + info.status = DependencyStatus.OK + info.installed = True + # 尝试解析版本 + import re + first_line = result.stdout.split('\n')[0] + version_match = re.search(r'ffmpeg version (\S+)', first_line, re.IGNORECASE) + if version_match: + info.version = version_match.group(1) + else: + info.status = DependencyStatus.ERROR + info.error_message = f"FFmpeg 执行失败: {result.stderr}" + except Exception as e: + info.status = DependencyStatus.ERROR + info.error_message = f"检查 FFmpeg 时发生错误: {e}" + else: + info.status = DependencyStatus.MISSING + info.installed = False + info.error_message = "未找到 FFmpeg" + + self.system_tools[name] = info + return info + + def check_all(self) -> bool: + """ + 执行所有依赖检查 + + Returns: + 所有必需依赖是否都满足 + """ + self._checked = True + + # 检查 Python 版本 + self.check_python_version() + + # 检查 Python 包 + self.check_all_packages() + + # 检查系统工具 + self.check_ffmpeg() + + return self.is_all_ok() + + def is_all_ok(self) -> bool: + """ + 检查所有必需依赖是否都满足 + + Returns: + 是否满足 + """ + if not self.python_version_ok: + return False + + # 检查必需的包 + for name, info in self.packages.items(): + if info.required and not info.is_ok(): + return False + + # 检查必需的系统工具 + for name, info in self.system_tools.items(): + if info.required and not info.is_ok(): + return False + + return True + + def get_missing_dependencies(self) -> List[DependencyInfo]: + """ + 获取缺失或有问题的依赖 + + Returns: + 有问题的依赖列表 + """ + missing = [] + + # 检查必需的包 + for name, info in self.packages.items(): + if info.required and not info.is_ok(): + missing.append(info) + + # 检查必需的系统工具 + for name, info in self.system_tools.items(): + if info.required and not info.is_ok(): + missing.append(info) + + return missing + + def get_warning_dependencies(self) -> List[DependencyInfo]: + """ + 获取可选但缺失的依赖（警告级别） + + Returns: + 可选但缺失的依赖列表 + """ + warnings = [] + + for name, info in self.packages.items(): + if not info.required and not info.is_ok(): + warnings.append(info) + + return warnings + + def generate_report(self, show_ok: bool = False) -> str: + """ + 生成依赖检查报告 + + Args: + show_ok: 是否显示状态正常的依赖 + + Returns: + 报告字符串 + """ + lines = [] + lines.append("=" * 60) + lines.append(" VideoSubtitleAI 依赖检查报告") + lines.append("=" * 60) + lines.append("") + + # Python 版本 + lines.append("【Python 版本】") + if self.python_version_info: + version_str = f"{self.python_version_info[0]}.{self.python_version_info[1]}.{self.python_version_info[2]}" + status = "✓ 兼容" if self.python_version_ok else "✗ 不兼容" + lines.append(f" 当前版本: {version_str}") + lines.append(f" 状态: {status}") + if not self.python_version_ok: + lines.append(f" 要求: Python 3.8 - 3.11") + lines.append("") + + # Python 包 + lines.append("【Python 包】") + for name, info in self.packages.items(): + if not show_ok and info.is_ok(): + continue + + status_icon = "✓" if info.is_ok() else "✗" + status_text = "OK" if info.is_ok() else f"问题: {info.error_message}" + required_text = "[必需]" if info.required else "[可选]" + + version_str = f" v{info.version}" if info.version else "" + lines.append(f" {status_icon} {info.display_name}{version_str} {required_text}") + if not info.is_ok(): + lines.append(f" {status_text}") + if info.install_guide: + lines.append(f" 安装指南:") + for guide_line in info.install_guide: + lines.append(f" {guide_line}") + lines.append("") + + # 系统工具 + lines.append("【系统工具】") + for name, info in self.system_tools.items(): + if not show_ok and info.is_ok(): + continue + + status_icon = "✓" if info.is_ok() else "✗" + status_text = "OK" if info.is_ok() else f"问题: {info.error_message}" + + version_str = f" v{info.version}" if info.version else "" + lines.append(f" {status_icon} {info.display_name}{version_str}") + if not info.is_ok(): + lines.append(f" {status_text}") + if info.install_guide: + lines.append(f" 安装指南:") + for guide_line in info.install_guide: + lines.append(f" {guide_line}") + lines.append("") + + # 总结 + lines.append("【总结】") + if self.is_all_ok(): + lines.append(" ✓ 所有必需依赖已满足！") + warnings = self.get_warning_dependencies() + if warnings: + lines.append(f" ⚠ {len(warnings)} 个可选依赖未安装（不影响核心功能）") + else: + missing = self.get_missing_dependencies() + lines.append(f" ✗ {len(missing)} 个必需依赖未满足，请安装后重试") + lines.append("") + lines.append("=" * 60) + + return '\n'.join(lines) + + def print_report(self, show_ok: bool = False): + """ + 打印依赖检查报告到控制台 + + Args: + show_ok: 是否显示状态正常的依赖 + """ + report = self.generate_report(show_ok=show_ok) + + # 添加颜色（如果支持） + try: + from colorama import init, Fore, Style + init() + + # 简单的颜色替换 + colored_report = report + colored_report = colored_report.replace('✓ ', Fore.GREEN + '✓ ' + Style.RESET_ALL) + colored_report = colored_report.replace('✗ ', Fore.RED + '✗ ' + Style.RESET_ALL) + colored_report = colored_report.replace('⚠ ', Fore.YELLOW + '⚠ ' + Style.RESET_ALL) + + print(colored_report) + except ImportError: + print(report) + + +# 便捷函数 +def check_dependencies() -> Tuple[bool, DependencyChecker]: + """ + 便捷函数：检查所有依赖 + + Returns: + (是否所有必需依赖都满足, 依赖检查器实例) + """ + checker = DependencyChecker() + checker.check_all() + return checker.is_all_ok(), checker + + +def print_dependency_report(show_ok: bool = False): + """ + 便捷函数：打印依赖检查报告 + + Args: + show_ok: 是否显示状态正常的依赖 + """ + checker = DependencyChecker() + checker.check_all() + checker.print_report(show_ok=show_ok) + + +if __name__ == '__main__': + # 当直接运行此模块时，执行依赖检查 + print_dependency_report(show_ok=True) diff --git a/main.py b/main.py index 2195e9e..af7ce2f 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,83 @@ """ VideoSubtitleAI - 本地离线音视频转字幕工具 主程序入口，提供命令行接口 +改进版：添加依赖检查，提供更友好的错误信息和安装指南 """ import os import sys import argparse +import traceback from pathlib import Path from typing import Optional, List, Dict, Any, Callable -# 导入自定义模块 -from audio_processor import AudioProcessor -from speech_recognizer import SpeechRecognizer -from subtitle_generator import SubtitleGenerator -from batch_processor import BatchProcessor, BatchProgress -from progress_logger import ProgressLogger, ConsoleLogger + +def check_dependencies_and_exit() -> bool: + """ + 检查依赖并在必要时退出 + + Returns: + 所有依赖是否满足 + """ + try: + from dependency_checker import check_dependencies, DependencyChecker, DependencyStatus + + all_ok, checker = check_dependencies() + + if not all_ok: + # 有缺失的依赖 + print("\n" + "=" * 60) + print(" ⚠️ 依赖检查失败") + print("=" * 60) + + missing = checker.get_missing_dependencies() + warnings = checker.get_warning_dependencies() + + for dep in missing: + print(f"\n✗ {dep.display_name}: {dep.error_message}") + if dep.install_guide: + print(" 安装指南:") + for line in dep.install_guide: + print(f" {line}") + + if warnings: + print(f"\n⚠ {len(warnings)} 个可选依赖未安装（不影响核心功能）:") + for dep in warnings: + print(f" - {dep.display_name}") + + print("\n" + "=" * 60) + print("请安装上述缺失的依赖后重试。") + print("=" * 60 + "\n") + return False + else: + print("✓ 所有必需依赖已满足") + return True + + except ImportError as e: + print(f"\n✗ 依赖检查模块加载失败: {e}") + print("这通常表示 Python 包未正确安装。") + print("\n请确保已安装以下必需包:") + print(" 1. PyTorch: pip install torch torchvision torchaudio") + print(" 2. Whisper: pip install openai-whisper") + print(" 3. NumPy: pip install numpy") + print("\n同时请确保 FFmpeg 已安装:") + print(" Windows: winget install Gyan.FFmpeg") + print(" macOS: brew install ffmpeg") + print(" Linux: sudo apt install ffmpeg") + return False + + +# 导入自定义模块（在依赖检查后） +try: + from audio_processor import AudioProcessor, FFmpegNotFoundError, FFmpegError + from speech_recognizer import SpeechRecognizer + from subtitle_generator import SubtitleGenerator + from batch_processor import BatchProcessor, BatchProgress + from progress_logger import ProgressLogger, ConsoleLogger + MODULES_LOADED = True +except ImportError as e: + print(f"\n✗ 模块加载失败: {e}") + MODULES_LOADED = False class VideoSubtitleAI: @@ -75,22 +138,54 @@ def initialize(self): # 初始化音频处理器 try: + self.progress_logger.info("正在初始化音频处理器...", "初始化") self.audio_processor = AudioProcessor(ffmpeg_path=self.ffmpeg_path) - self.progress_logger.info("音频处理器初始化成功", "初始化") + + # 显示 FFmpeg 信息 + ffmpeg_info = self.audio_processor.get_ffmpeg_info() + if ffmpeg_info['version']: + self.progress_logger.info(f"FFmpeg版本: {ffmpeg_info['version']}", "初始化") + self.progress_logger.success("音频处理器初始化成功", "初始化") + + except FFmpegNotFoundError as e: + # 显示详细的错误信息和解决方案 + print("\n" + "=" * 60) + print(" ❌ FFmpeg 未找到") + print("=" * 60) + print(f"\n错误: {e}") + print("\n" + e.get_detailed_message()) + print("=" * 60 + "\n") + + if self.progress_logger: + self.progress_logger.error(f"音频处理器初始化失败: {e}", "初始化") + raise + except Exception as e: - self.progress_logger.error(f"音频处理器初始化失败: {e}", "初始化") + if self.progress_logger: + self.progress_logger.error(f"音频处理器初始化失败: {e}", "初始化") raise # 初始化语音识别器 try: + self.progress_logger.info("正在初始化语音识别器...", "初始化") self.speech_recognizer = SpeechRecognizer( model_name=self.model_name, device=self.device, model_dir=self.model_dir, ) - self.progress_logger.info(f"语音识别器初始化成功，模型: {self.model_name}", "初始化") + self.progress_logger.success(f"语音识别器初始化成功，模型: {self.model_name}", "初始化") + except ImportError as e: + print("\n" + "=" * 60) + print(" ❌ Whisper 模型加载失败") + print("=" * 60) + print(f"\n错误: {e}") + print("\n请确保已安装 Whisper:") + print(" pip install openai-whisper") + print("=" * 60 + "\n") + raise except Exception as e: - self.progress_logger.error(f"语音识别器初始化失败: {e}", "初始化") + if self.progress_logger: + self.progress_logger.error(f"语音识别器初始化失败: {e}", "初始化") raise # 初始化字幕生成器 @@ -114,6 +209,14 @@ def _show_device_info(self): self.progress_logger.info(f"GPU: {info['gpu_name']} ({info['gpu_memory']})", "设备信息") else: self.progress_logger.warning("GPU不可用，将使用CPU运行（速度较慢）", "设备信息") + print("\n" + "=" * 60) + print(" ⚠️ GPU 不可用") + print("=" * 60) + print("\n未检测到可用的 NVIDIA GPU，将使用 CPU 运行。") + print("CPU 模式运行速度较慢，建议安装 CUDA 版本的 PyTorch。") + print("\n安装 CUDA 版本 PyTorch:") + print(" pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118") + print("=" * 60 + "\n") def load_model(self, force_reload: bool = False): """ @@ -123,8 +226,27 @@ def load_model(self, force_reload: bool = False): self.initialize() self.progress_logger.info(f"正在加载模型: {self.model_name}...", "模型加载") - self.speech_recognizer.load_model(force_reload=force_reload) - self.progress_logger.success("模型加载完成", "模型加载") + self.progress_logger.info("首次运行会自动下载模型，请耐心等待...", "模型加载") + + try: + self.speech_recognizer.load_model(force_reload=force_reload) + self.progress_logger.success("模型加载完成", "模型加载") + except Exception as e: + self.progress_logger.error(f"模型加载失败: {e}", "模型加载") + print("\n" + "=" * 60) + print(" ❌ 模型加载失败") + print("=" * 60) + print(f"\n错误: {e}") + print("\n可能的原因:") + print(" 1. 网络连接问题（首次运行需要下载模型）") + print(" 2. 磁盘空间不足") + print(" 3. 内存不足") + print("\n解决方案:") + print(" - 检查网络连接") + print(" - 确保有足够的磁盘空间（模型大小从 39MB 到 1.5GB）") + print(" - 尝试使用更小的模型: -m tiny 或 -m base") + print("=" * 60 + "\n") + raise def unload_model(self): """ @@ -132,7 +254,8 @@ def unload_model(self): """ if self.speech_recognizer: self.speech_recognizer.unload_model() - self.progress_logger.info("模型已卸载", "模型管理") + if self.progress_logger: + self.progress_logger.info("模型已卸载", "模型管理") def process_single_file( self, @@ -165,6 +288,14 @@ def process_single_file( if not input_path.exists(): raise FileNotFoundError(f"输入文件不存在: {input_path}") + # 检查文件格式 + if not self.audio_processor.is_supported_format(str(input_path)): + raise ValueError( + f"不支持的文件格式: {input_path.suffix}\n" + f"支持的视频格式: {', '.join(sorted(self.audio_processor.SUPPORTED_VIDEO_FORMATS))}\n" + f"支持的音频格式: {', '.join(sorted(self.audio_processor.SUPPORTED_AUDIO_FORMATS))}" + ) + # 确定输出路径 if output_path is None: output_path = input_path.parent / f"{input_path.stem}.srt" @@ -223,6 +354,19 @@ def process_single_file( return result + except FFmpegError as e: + self.progress_logger.error(f"FFmpeg执行失败: {e}", "错误") + print("\n" + "=" * 60) + print(" ❌ 音频处理失败") + print("=" * 60) + print(f"\nFFmpeg执行错误 (返回码: {e.return_code})") + print(f"错误信息: {e.stderr}") + print("\n可能的原因:") + print(" 1. 输入文件损坏或格式不兼容") + print(" 2. FFmpeg版本太旧") + print(" 3. 磁盘空间不足") + print("=" * 60 + "\n") + raise except Exception as e: self.progress_logger.error(f"处理失败: {e}", "错误") raise @@ -295,6 +439,13 @@ def process_folder( if not file_pairs: self.progress_logger.warning("没有找到待处理的文件", "批量处理") + print("\n" + "=" * 60) + print(" ⚠️ 未找到待处理文件") + print("=" * 60) + print(f"\n在 {input_folder} 中未找到支持的音视频文件。") + print(f"\n支持的视频格式: {', '.join(sorted(self.audio_processor.SUPPORTED_VIDEO_FORMATS))}") + print(f"支持的音频格式: {', '.join(sorted(self.audio_processor.SUPPORTED_AUDIO_FORMATS))}") + print("=" * 60 + "\n") return { 'success': True, 'total_files': 0, @@ -417,6 +568,9 @@ def get_info(self) -> Dict[str, Any]: info['device_info'] = self.speech_recognizer.get_device_info() info['model_info'] = self.speech_recognizer.get_model_info() + if self.audio_processor: + info['ffmpeg_info'] = self.audio_processor.get_ffmpeg_info() + return info @@ -429,9 +583,15 @@ def parse_args(): formatter_class=argparse.RawDescriptionHelpFormatter, epilog=''' 示例: + # 检查依赖 + python main.py --check-deps + # 处理单个视频 python main.py -i video.mp4 -m base + # 指定FFmpeg路径 + python main.py -i video.mp4 --ffmpeg-path "C:\\ffmpeg\\bin\\ffmpeg.exe" + # 处理整个文件夹 python main.py -i ./videos -o ./subtitles --recursive -m small @@ -441,7 +601,8 @@ def parse_args(): ) # 输入输出参数 - parser.add_argument('-i', '--input', required=True, help='输入文件或文件夹路径') + parser.add_argument('-i', '--input', required=False, default=None, + help='输入文件或文件夹路径（使用 --check-deps 时可选）') parser.add_argument('-o', '--output', default=None, help='输出文件或文件夹路径（可选）') # 模型参数 @@ -452,6 +613,10 @@ def parse_args(): parser.add_argument('--device', default=None, choices=['cpu', 'cuda'], help='运行设备 (默认: 自动选择)') + # FFmpeg 参数 + parser.add_argument('--ffmpeg-path', default=None, + help='FFmpeg可执行文件路径（可选，如未指定则从系统PATH中查找）') + # 语言参数 parser.add_argument('-l', '--language', default='auto', help='语言代码 (默认: auto自动检测，可选: zh, en, ja等)') @@ -473,77 +638,137 @@ def parse_args(): help='仅检测语言，不生成字幕') parser.add_argument('--info', action='store_true', help='显示系统信息') - parser.add_argument('--ffmpeg-path', default=None, - help='FFmpeg可执行文件路径（可选，如未指定则从系统PATH中查找）') + parser.add_argument('--check-deps', action='store_true', + help='检查依赖并退出') return parser.parse_args() +def handle_exception(e: Exception, verbose: bool = False): + """ + 处理异常，提供友好的错误信息 + + Args: + e: 异常对象 + verbose: 是否显示详细堆栈 + """ + print("\n" + "=" * 60) + print(" ❌ 程序执行失败") + print("=" * 60) + print(f"\n错误类型: {type(e).__name__}") + print(f"错误信息: {e}") + + if verbose: + print("\n详细错误信息:") + traceback.print_exc() + + print("\n" + "=" * 60) + print("如有问题，请检查:") + print(" 1. 所有依赖是否已正确安装") + print(" 2. 输入文件是否存在且格式正确") + print(" 3. 是否有足够的磁盘空间") + print("=" * 60 + "\n") + + def main(): """ 主函数入口 """ args = parse_args() + # 检查依赖 + if args.check_deps: + print("\n" + "=" * 60) + print(" VideoSubtitleAI 依赖检查") + print("=" * 60 + "\n") + + check_dependencies_and_exit() + + # 尝试打印详细报告 + try: + from dependency_checker import print_dependency_report + print("") + print_dependency_report(show_ok=True) + except Exception: + pass + + sys.exit(0) + + # 检查必需参数 + if args.input is None: + print("\n错误: 缺少必需参数 -i/--input") + print("使用 --help 查看帮助信息") + sys.exit(1) + # 检查输入 input_path = Path(args.input) if not input_path.exists(): - print(f"错误: 输入路径不存在: {input_path}") + print(f"\n错误: 输入路径不存在: {input_path}") sys.exit(1) # 确定任务类型 task = 'translate' if args.translate else 'transcribe' - # 创建实例 - ai = VideoSubtitleAI( - model_name=args.model, - device=args.device, - model_dir=args.model_dir, - log_dir=args.log_dir, - verbose=args.verbose, - ffmpeg_path=args.ffmpeg_path, - ) - - # 显示信息 - if args.info: - ai.initialize() - info = ai.get_info() - print("\n" + "=" * 60) - print(" 系统信息") - print("=" * 60) - print(f" 版本: {info['version']}") - print(f" 模型: {info['model_name']}") - if 'device_info' in info: - dev = info['device_info'] - print(f" 设备: {dev['current_device']}") - if dev['gpu_available']: - print(f" GPU: {dev['gpu_name']}") - print("=" * 60 + "\n") - return - - # 仅检测语言 - if args.detect_language_only: - ai.initialize() - if input_path.is_file(): - result = ai.detect_language(str(input_path)) - print(f"\n检测结果:") - print(f" 语言: {result['language_name']}") - print(f" 代码: {result['detected_language']}") - print(f" 置信度: {result['confidence']:.2%}") - print("\nTop 5 可能语言:") - for i, lang in enumerate(result['top_languages'], 1): - print(f" {i}. {lang['language_name']}: {lang['confidence']:.2%}") - else: - print("错误: --detect-language-only 仅支持单个文件") - sys.exit(1) - return - - # 初始化并加载模型 - ai.initialize() - ai.load_model() + # 检查模块是否加载 + if not MODULES_LOADED: + print("\n错误: 核心模块加载失败，请确保所有依赖已正确安装。") + print("运行 --check-deps 检查依赖状态。") + sys.exit(1) - # 处理 try: + # 创建实例 + ai = VideoSubtitleAI( + model_name=args.model, + device=args.device, + model_dir=args.model_dir, + log_dir=args.log_dir, + verbose=args.verbose, + ffmpeg_path=args.ffmpeg_path, + ) + + # 显示信息 + if args.info: + ai.initialize() + info = ai.get_info() + print("\n" + "=" * 60) + print(" 系统信息") + print("=" * 60) + print(f" 版本: {info['version']}") + print(f" 模型: {info['model_name']}") + if 'device_info' in info: + dev = info['device_info'] + print(f" 设备: {dev['current_device']}") + if dev['gpu_available']: + print(f" GPU: {dev['gpu_name']}") + if 'ffmpeg_info' in info: + ffmpeg = info['ffmpeg_info'] + print(f" FFmpeg: {ffmpeg.get('version', 'unknown')}") + print(f" FFmpeg路径: {ffmpeg.get('ffmpeg_path', 'unknown')}") + print("=" * 60 + "\n") + return + + # 仅检测语言 + if args.detect_language_only: + ai.initialize() + if input_path.is_file(): + result = ai.detect_language(str(input_path)) + print(f"\n检测结果:") + print(f" 语言: {result['language_name']}") + print(f" 代码: {result['detected_language']}") + print(f" 置信度: {result['confidence']:.2%}") + print("\nTop 5 可能语言:") + for i, lang in enumerate(result['top_languages'], 1): + print(f" {i}. {lang['language_name']}: {lang['confidence']:.2%}") + else: + print("错误: --detect-language-only 仅支持单个文件") + sys.exit(1) + return + + # 初始化并加载模型 + ai.initialize() + ai.load_model() + + # 处理 if input_path.is_file(): # 单个文件 result = ai.process_single_file( @@ -552,11 +777,15 @@ def main(): language=args.language, task=task, ) - print(f"\n处理完成!") - print(f" 输入: {result['input_path']}") - print(f" 输出: {result['output_path']}") - print(f" 语言: {result['language_name']}") + print(f"\n" + "=" * 60) + print(" ✅ 处理完成!") + print("=" * 60) + print(f"\n 输入文件: {result['input_path']}") + print(f" 输出字幕: {result['output_path']}") + print(f" 检测语言: {result['language_name']}") print(f" 字幕段数: {result['segment_count']}") + print(f" 总字数: {result['total_characters']}") + print("\n" + "=" * 60 + "\n") else: # 文件夹 result = ai.process_folder( @@ -567,22 +796,37 @@ def main(): task=task, skip_existing=args.skip_existing, ) - print(f"\n批量处理完成!") - print(f" 总文件数: {result['total_files']}") - print(f" 成功: {result['completed_files']}") - print(f" 失败: {result['failed_files']}") + print(f"\n" + "=" * 60) + print(" ✅ 批量处理完成!") + print("=" * 60) + print(f"\n 总文件数: {result['total_files']}") + print(f" 成功处理: {result['completed_files']}") + print(f" 处理失败: {result['failed_files']}") + print(f" 跳过: {result['skipped_files']}") print(f" 成功率: {result['success_rate']}") - print(f" 耗时: {result['elapsed_time']}") + print(f" 总耗时: {result['elapsed_time']}") + + if result['errors']: + print(f"\n 失败的文件:") + for err in result['errors']: + print(f" - {Path(err['input_path']).name}: {err.get('error', '未知错误')}") + print("\n" + "=" * 60 + "\n") + + except FFmpegNotFoundError as e: + # 已经在 initialize 中处理过了 + sys.exit(1) except Exception as e: - print(f"\n错误: {e}") - import traceback - traceback.print_exc() + handle_exception(e, args.verbose) sys.exit(1) finally: # 卸载模型 - ai.unload_model() + try: + if 'ai' in locals(): + ai.unload_model() + except Exception: + pass if __name__ == '__main__': diff --git a/requirements.txt b/requirements.txt index 96c87d8..bf40f5e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,25 +1,101 @@ # VideoSubtitleAI - 本地离线音视频转字幕工具 # 依赖包列表 +# +# ⚠️ 重要提示：在安装这些包之前，请先安装 FFmpeg +# Windows: winget install Gyan.FFmpeg +# macOS: brew install ffmpeg +# Linux: sudo apt install ffmpeg +# +# 快速安装所有依赖： +# pip install -r requirements.txt +# -# 核心依赖 +# ============================================================================ +# 必需依赖（必须安装） +# ============================================================================ + +# OpenAI Whisper - 核心语音识别模型 openai-whisper>=20231117 -# PyTorch (Whisper需要，根据系统环境选择合适版本) -# 安装GPU版本 (推荐，速度快): -# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 +# PyTorch - Whisper 依赖的深度学习框架 +# ⚠️ 强烈建议安装 CUDA 版本以获得 GPU 加速 +# +# 安装 GPU 版本 (推荐，速度提升 10-100 倍): +# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 +# +# 或者安装 CUDA 12.1 版本: +# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 # -# 或者安装CPU版本: -# pip3 install torch torchvision torchaudio +# 安装 CPU 版本 (速度较慢): +# pip3 install torch torchvision torchaudio # -# 注意: 以下仅为最小依赖，实际安装时请根据需求安装完整版PyTorch +# 注意: 以下仅为最小依赖，实际使用时请安装完整版 torch>=2.0.0 -# 音频处理 (可选，FFmpeg通常需要单独安装) -# ffmpeg-python>=0.2.0 -# pydub>=0.25.1 +# NumPy - 数值计算库 +numpy>=1.24.0 + +# ============================================================================ +# 可选依赖（推荐安装） +# ============================================================================ # 进度条显示 tqdm>=4.66.0 -# 其他工具 -numpy>=1.24.0 +# 颜色输出（用于更友好的终端显示） +colorama>=0.4.6 + +# 版本比较（用于依赖检查） +packaging>=23.0 + +# ============================================================================ +# 系统工具（必须单独安装） +# ============================================================================ +# +# FFmpeg - 音视频处理工具（必需） +# +# Windows: +# 方法1: 使用 winget (推荐) +# winget install Gyan.FFmpeg +# +# 方法2: 使用 Chocolatey +# choco install ffmpeg +# +# 方法3: 手动安装 +# 1. 访问 https://www.gyan.dev/ffmpeg/builds/ +# 2. 下载 ffmpeg-release-full.7z +# 3. 解压到 C:\ffmpeg +# 4. 将 C:\ffmpeg\bin 添加到系统 PATH +# 5. 重启命令行/IDE +# +# macOS: +# brew install ffmpeg +# +# Linux (Ubuntu/Debian): +# sudo apt update +# sudo apt install ffmpeg +# +# Linux (Fedora): +# sudo dnf install ffmpeg +# +# ============================================================================ + +# ============================================================================ +# 安装步骤总结 +# ============================================================================ +# +# 1. 安装 FFmpeg (必需) +# Windows: winget install Gyan.FFmpeg +# macOS: brew install ffmpeg +# Linux: sudo apt install ffmpeg +# +# 2. 安装 PyTorch (推荐 GPU 版本) +# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 +# +# 3. 安装其他依赖 +# pip install -r requirements.txt +# +# 4. 验证安装 +# python main.py --check-deps +# +# ============================================================================ From 0b576c1ff1212004bdb6b5afb2ae62b1b172893b Mon Sep 17 00:00:00 2001 From: Unizzr <1920158711@qq.com> Date: Sat, 2 May 2026 10:47:58 +0800 Subject: [PATCH 2/2] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=E5=8F=8C?= =?UTF-8?q?=E8=AF=AD=E5=AD=97=E5=B9=95=E7=BF=BB=E8=AF=91=E5=8F=8A=E5=AD=97?= =?UTF-8?q?=E5=B9=95=E5=A4=84=E7=90=86=E5=A2=9E=E5=BC=BA=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增字幕处理核心模块，包括： 1. subtitle_models.py - 统一字幕数据模型 2. subtitle_translator.py - 支持多后端的翻译模块 3. subtitle_editor.py - 提供撤销/重做功能的字幕编辑器 4. subtitle_exporter.py - 支持SRT/ASS/VTT多格式导出 5. subtitle_embedder.py - 基于FFmpeg的字幕硬嵌入功能 6. video_subtitle_enhanced.py - 增强版主模块整合所有功能 同时添加测试模块test_modules.py验证各功能 --- language_detector.py | 599 ++++++++++++++++++++++++++ main.py | 425 ++++++++++++++++++- requirements.txt | 33 ++ subtitle_editor.py | 700 +++++++++++++++++++++++++++++++ subtitle_embedder.py | 491 ++++++++++++++++++++++ subtitle_exporter.py | 778 ++++++++++++++++++++++++++++++++++ subtitle_models.py | 539 ++++++++++++++++++++++++ subtitle_translator.py | 832 +++++++++++++++++++++++++++++++++++++ test_modules.py | 210 ++++++++++ video_subtitle_enhanced.py | 492 ++++++++++++++++++++++ 10 files changed, 5076 insertions(+), 23 deletions(-) create mode 100644 language_detector.py create mode 100644 subtitle_editor.py create mode 100644 subtitle_embedder.py create mode 100644 subtitle_exporter.py create mode 100644 subtitle_models.py create mode 100644 subtitle_translator.py create mode 100644 test_modules.py create mode 100644 video_subtitle_enhanced.py diff --git a/language_detector.py b/language_detector.py new file mode 100644 index 0000000..fadb32e --- /dev/null +++ b/language_detector.py @@ -0,0 +1,599 @@ +""" +语言检测优化模块 +支持多语言混合场景的识别 +实现分段语言检测、自动切换语言模型 +""" + +import os +import sys +import re +import logging +from typing import List, Dict, Any, Optional, Tuple +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +import numpy as np + +logger = logging.getLogger(__name__) + + +LANGUAGE_NAMES = { + 'zh': '中文', + 'en': '英文', + 'ja': '日文', + 'ko': '韩文', + 'fr': '法文', + 'de': '德文', + 'es': '西班牙文', + 'ru': '俄文', + 'pt': '葡萄牙文', + 'it': '意大利文', + 'auto': '自动检测', +} + + +CHINESE_CHARS = set() +for codepoint in range(0x4e00, 0x9fff + 1): + CHINESE_CHARS.add(chr(codepoint)) + +JAPANESE_HIRAGANA = set() +for codepoint in range(0x3040, 0x309f + 1): + JAPANESE_HIRAGANA.add(chr(codepoint)) + +JAPANESE_KATAKANA = set() +for codepoint in range(0x30a0, 0x30ff + 1): + JAPANESE_KATAKANA.add(chr(codepoint)) + +KOREAN_HANGUL = set() +for codepoint in range(0xac00, 0xd7a3 + 1): + KOREAN_HANGUL.add(chr(codepoint)) + + +@dataclass +class LanguageDetectionResult: + """语言检测结果""" + language: str = 'unknown' + language_name: str = '未知' + confidence: float = 0.0 + top_languages: List[Dict[str, Any]] = field(default_factory=list) + detection_method: str = 'unknown' + is_mixed: bool = False + mixed_languages: List[str] = field(default_factory=list) + + +@dataclass +class SegmentLanguageInfo: + """段落语言信息""" + segment_index: int = 0 + start_time: float = 0.0 + end_time: float = 0.0 + text: str = '' + detected_language: str = 'unknown' + confidence: float = 0.0 + language_scores: Dict[str, float] = field(default_factory=dict) + + +class LanguageDetector: + """ + 增强语言检测器 + 支持： + 1. 基于字符特征的快速语言检测 + 2. 多语言混合检测 + 3. 段落级语言分析 + 4. 与Whisper模型集成的高精度检测 + """ + + def __init__(self, whisper_model=None): + self.whisper_model = whisper_model + self._language_detector = None + + def detect_language_from_text(self, text: str) -> LanguageDetectionResult: + """ + 基于文本内容快速检测语言 + """ + if not text or not text.strip(): + return LanguageDetectionResult( + language='unknown', + language_name='未知', + confidence=0.0, + detection_method='text_analysis', + ) + + scores = self._calculate_language_scores(text) + + if not scores: + return LanguageDetectionResult( + language='unknown', + language_name='未知', + confidence=0.0, + detection_method='text_analysis', + ) + + sorted_languages = sorted( + scores.items(), + key=lambda x: x[1], + reverse=True + ) + + top_lang, top_confidence = sorted_languages[0] + + top_languages = [] + for lang, conf in sorted_languages[:5]: + if conf > 0.01: + top_languages.append({ + 'language': lang, + 'language_name': LANGUAGE_NAMES.get(lang, lang), + 'confidence': conf, + }) + + mixed_languages = [] + for lang, conf in sorted_languages[1:]: + if conf > 0.15 and top_confidence - conf < 0.5: + mixed_languages.append(lang) + + is_mixed = len(mixed_languages) > 0 + + return LanguageDetectionResult( + language=top_lang, + language_name=LANGUAGE_NAMES.get(top_lang, top_lang), + confidence=top_confidence, + top_languages=top_languages, + detection_method='text_analysis', + is_mixed=is_mixed, + mixed_languages=mixed_languages, + ) + + def _calculate_language_scores(self, text: str) -> Dict[str, float]: + """ + 计算各种语言的得分 + """ + scores = {} + total_chars = len(text) + + if total_chars == 0: + return scores + + chinese_count = 0 + japanese_count = 0 + korean_count = 0 + latin_count = 0 + cyrillic_count = 0 + digit_count = 0 + symbol_count = 0 + + for char in text: + if char in CHINESE_CHARS: + chinese_count += 1 + elif char in JAPANESE_HIRAGANA or char in JAPANESE_KATAKANA: + japanese_count += 1 + elif char in KOREAN_HANGUL: + korean_count += 1 + elif '\u0400' <= char <= '\u04ff' or '\u0500' <= char <= '\u052f': + cyrillic_count += 1 + elif ('a' <= char <= 'z') or ('A' <= char <= 'Z'): + latin_count += 1 + elif char.isdigit(): + digit_count += 1 + elif char.isspace(): + continue + else: + symbol_count += 1 + + meaningful_chars = chinese_count + japanese_count + korean_count + latin_count + cyrillic_count + if meaningful_chars == 0: + meaningful_chars = 1 + + if chinese_count > 0: + chinese_ratio = chinese_count / meaningful_chars + + if japanese_count > 0: + ja_ratio = japanese_count / meaningful_chars + if ja_ratio > 0.1: + scores['ja'] = ja_ratio * 0.8 + scores['zh'] = chinese_ratio * 0.6 + else: + scores['zh'] = chinese_ratio + else: + scores['zh'] = chinese_ratio + + if japanese_count > 0 and 'ja' not in scores: + ja_ratio = japanese_count / meaningful_chars + scores['ja'] = ja_ratio + + if korean_count > 0: + ko_ratio = korean_count / meaningful_chars + scores['ko'] = ko_ratio + + if latin_count > 0: + en_ratio = latin_count / meaningful_chars + + if en_ratio > 0.1: + word_count = len(re.findall(r'\b[a-zA-Z]{2,}\b', text)) + if word_count >= 2: + scores['en'] = min(1.0, en_ratio + 0.2) + else: + scores['en'] = en_ratio * 0.5 + + if cyrillic_count > 0: + ru_ratio = cyrillic_count / meaningful_chars + scores['ru'] = ru_ratio + + if not scores and latin_count > 0: + scores['en'] = 0.5 + + total_score = sum(scores.values()) + if total_score > 0: + normalized_scores = {k: v / total_score for k, v in scores.items()} + else: + normalized_scores = scores + + return normalized_scores + + def detect_language_with_whisper( + self, + audio_path: str, + sample_duration: float = 30.0, + ) -> LanguageDetectionResult: + """ + 使用Whisper模型进行高精度语言检测 + """ + if self.whisper_model is None: + return LanguageDetectionResult( + language='unknown', + language_name='未知', + confidence=0.0, + detection_method='whisper_unavailable', + ) + + try: + import whisper + + audio = whisper.load_audio(audio_path) + audio = whisper.pad_or_trim(audio) + + mel = whisper.log_mel_spectrogram(audio).to(self.whisper_model.device) + + _, probs = self.whisper_model.detect_language(mel) + + detected_lang = max(probs, key=probs.get) + confidence = probs[detected_lang] + + top_languages = [] + for lang, prob in sorted(probs.items(), key=lambda x: x[1], reverse=True)[:5]: + if prob > 0.01: + top_languages.append({ + 'language': lang, + 'language_name': LANGUAGE_NAMES.get(lang, lang), + 'confidence': float(prob), + }) + + return LanguageDetectionResult( + language=detected_lang, + language_name=LANGUAGE_NAMES.get(detected_lang, detected_lang), + confidence=float(confidence), + top_languages=top_languages, + detection_method='whisper_model', + ) + + except Exception as e: + logger.error(f"Whisper语言检测失败: {e}") + return LanguageDetectionResult( + language='unknown', + language_name='未知', + confidence=0.0, + detection_method='whisper_error', + ) + + def analyze_segments_language( + self, + segments: List[Dict[str, Any]], + ) -> List[SegmentLanguageInfo]: + """ + 分析每个段落的语言 + """ + results = [] + + for idx, seg in enumerate(segments): + text = seg.get('text', '') if isinstance(seg, dict) else getattr(seg, 'text', '') + start = seg.get('start', 0.0) if isinstance(seg, dict) else getattr(seg, 'start', 0.0) + end = seg.get('end', 0.0) if isinstance(seg, dict) else getattr(seg, 'end', 0.0) + + detection = self.detect_language_from_text(text) + + seg_info = SegmentLanguageInfo( + segment_index=idx, + start_time=start, + end_time=end, + text=text, + detected_language=detection.language, + confidence=detection.confidence, + language_scores={lang: info['confidence'] for lang, info in + [(l['language'], l) for l in detection.top_languages]}, + ) + + results.append(seg_info) + + return results + + def detect_mixed_language_scenes( + self, + segments: List[Dict[str, Any]], + threshold: float = 0.3, + ) -> Dict[str, Any]: + """ + 检测多语言混合场景 + """ + segment_languages = self.analyze_segments_language(segments) + + language_distribution = {} + for seg_info in segment_languages: + lang = seg_info.detected_language + if lang not in language_distribution: + language_distribution[lang] = 0 + language_distribution[lang] += 1 + + total_segments = len(segment_languages) + if total_segments == 0: + return { + 'is_mixed': False, + 'primary_language': 'unknown', + 'languages': [], + 'segments': [], + } + + language_ratios = { + lang: count / total_segments + for lang, count in language_distribution.items() + } + + significant_languages = [ + lang for lang, ratio in language_ratios.items() + if ratio >= threshold and lang != 'unknown' + ] + + is_mixed = len(significant_languages) > 1 + + primary_language = max( + language_ratios.items(), + key=lambda x: x[1] + )[0] if language_ratios else 'unknown' + + language_segments = {} + for lang in significant_languages: + language_segments[lang] = [ + seg for seg in segment_languages + if seg.detected_language == lang + ] + + return { + 'is_mixed': is_mixed, + 'primary_language': primary_language, + 'primary_language_name': LANGUAGE_NAMES.get(primary_language, primary_language), + 'languages': significant_languages, + 'language_names': [LANGUAGE_NAMES.get(lang, lang) for lang in significant_languages], + 'language_distribution': language_distribution, + 'language_ratios': language_ratios, + 'segment_languages': segment_languages, + 'language_segments': language_segments, + 'threshold': threshold, + } + + def get_optimal_language_for_transcription( + self, + segments: List[Dict[str, Any]], + prefer_multilingual: bool = True, + ) -> Dict[str, Any]: + """ + 获取最优的转录语言设置 + """ + mixed_info = self.detect_mixed_language_scenes(segments) + + if mixed_info['is_mixed'] and prefer_multilingual: + return { + 'language': 'auto', + 'language_name': '自动检测（多语言混合）', + 'is_mixed': True, + 'languages': mixed_info['languages'], + 'recommendation': '使用auto模式，Whisper会自动切换语言模型', + 'details': mixed_info, + } + else: + primary = mixed_info['primary_language'] + return { + 'language': primary if primary != 'unknown' else 'auto', + 'language_name': LANGUAGE_NAMES.get(primary, '自动检测'), + 'is_mixed': False, + 'languages': [primary] if primary != 'unknown' else [], + 'recommendation': f'使用{primary}模式以获得最佳识别效果', + 'details': mixed_info, + } + + +class MultilingualSpeechRecognizer: + """ + 多语言语音识别增强器 + 针对多语言混合场景优化识别效果 + """ + + def __init__(self, base_recognizer=None, whisper_model=None): + self.base_recognizer = base_recognizer + self.whisper_model = whisper_model + self.language_detector = LanguageDetector(whisper_model=whisper_model) + + def recognize_with_language_detection( + self, + audio_path: str, + language: str = 'auto', + task: str = 'transcribe', + **kwargs, + ) -> Dict[str, Any]: + """ + 带语言检测的语音识别 + """ + if self.whisper_model is None and self.base_recognizer is None: + raise RuntimeError("没有可用的语音识别模型") + + if self.base_recognizer: + result = self.base_recognizer.recognize( + audio_path=audio_path, + language=language, + task=task, + **kwargs, + ) + else: + try: + import whisper + + result = self.whisper_model.transcribe( + audio=str(audio_path), + language=None if language == 'auto' else language, + task=task, + **kwargs, + ) + + detected_language = result.get('language', 'unknown') + segments = result.get('segments', []) + + formatted_segments = [] + for seg in segments: + formatted = { + 'id': seg.get('id', 0), + 'start': round(float(seg.get('start', 0)), 3), + 'end': round(float(seg.get('end', 0)), 3), + 'text': seg.get('text', '').strip(), + 'language': seg.get('language', detected_language), + 'tokens': seg.get('tokens', []), + 'avg_logprob': seg.get('avg_logprob', 0.0), + 'no_speech_prob': seg.get('no_speech_prob', 0.0), + } + formatted_segments.append(formatted) + + result = { + 'text': result.get('text', '').strip(), + 'segments': formatted_segments, + 'language': detected_language, + 'language_name': LANGUAGE_NAMES.get(detected_language, detected_language), + 'segments_info': formatted_segments, + 'audio_path': audio_path, + 'task': task, + } + + except Exception as e: + logger.error(f"Whisper识别失败: {e}") + raise + + segments = result.get('segments', []) + + if segments: + mixed_info = self.language_detector.detect_mixed_language_scenes(segments) + + result['language_analysis'] = { + 'is_mixed': mixed_info['is_mixed'], + 'languages': mixed_info.get('languages', []), + 'language_names': mixed_info.get('language_names', []), + 'primary_language': mixed_info.get('primary_language', 'unknown'), + 'language_ratios': mixed_info.get('language_ratios', {}), + } + + return result + + def analyze_and_suggest_language( + self, + audio_path: str, + segments: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """ + 分析音频并建议最佳语言设置 + """ + suggestions = {} + + if segments: + mixed_info = self.language_detector.detect_mixed_language_scenes(segments) + + suggestions['segment_analysis'] = mixed_info + + if mixed_info['is_mixed']: + suggestions['recommended_language'] = 'auto' + suggestions['recommendation'] = ( + f"检测到多语言混合场景（{', '.join(mixed_info['language_names'])}），" + f"建议使用auto模式以获得最佳识别效果" + ) + suggestions['is_mixed'] = True + else: + primary = mixed_info['primary_language'] + suggestions['recommended_language'] = primary if primary != 'unknown' else 'auto' + suggestions['recommendation'] = ( + f"检测到主要语言为{mixed_info['primary_language_name']}，" + f"建议使用该语言以获得最佳识别效果" + ) + suggestions['is_mixed'] = False + + if self.whisper_model: + whisper_result = self.language_detector.detect_language_with_whisper(audio_path) + suggestions['whisper_detection'] = { + 'language': whisper_result.language, + 'language_name': whisper_result.language_name, + 'confidence': whisper_result.confidence, + 'top_languages': whisper_result.top_languages, + } + + return suggestions + + +def is_chinese_text(text: str) -> bool: + """ + 快速判断文本是否主要是中文 + """ + if not text: + return False + + chinese_count = 0 + total_chars = 0 + + for char in text: + if char.isspace() or char.isdigit(): + continue + total_chars += 1 + if char in CHINESE_CHARS: + chinese_count += 1 + + if total_chars == 0: + return False + + return chinese_count / total_chars > 0.5 + + +def is_english_text(text: str) -> bool: + """ + 快速判断文本是否主要是英文 + """ + if not text: + return False + + english_count = 0 + total_chars = 0 + + for char in text: + if char.isspace() or char.isdigit(): + continue + total_chars += 1 + if ('a' <= char <= 'z') or ('A' <= char <= 'Z'): + english_count += 1 + + if total_chars == 0: + return False + + return english_count / total_chars > 0.6 + + +def classify_text_language(text: str) -> Tuple[str, float]: + """ + 分类文本语言，返回(语言代码, 置信度) + """ + if not text or not text.strip(): + return ('unknown', 0.0) + + detector = LanguageDetector() + result = detector.detect_language_from_text(text) + + return (result.language, result.confidence) diff --git a/main.py b/main.py index af7ce2f..9f3497b 100644 --- a/main.py +++ b/main.py @@ -579,24 +579,40 @@ def parse_args(): 解析命令行参数 """ parser = argparse.ArgumentParser( - description='VideoSubtitleAI - 本地离线音视频转字幕工具', + description='VideoSubtitleAI - 本地离线音视频转字幕工具（增强版）', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=''' -示例: +基本用法示例: # 检查依赖 python main.py --check-deps - # 处理单个视频 + # 处理单个视频，生成SRT字幕 python main.py -i video.mp4 -m base - # 指定FFmpeg路径 - python main.py -i video.mp4 --ffmpeg-path "C:\\ffmpeg\\bin\\ffmpeg.exe" + # 生成双语字幕（中英文） + python main.py -i video.mp4 --bilingual --bilingual-order zh_en + + # 导出多种格式（SRT + ASS + VTT） + python main.py -i video.mp4 --formats srt,ass,vtt + + # 字幕硬嵌入到视频中 + python main.py -i video.mp4 --embed-video + + # 自定义字幕样式（黄色大字体） + python main.py -i video.mp4 --style yellow --font-size 48 # 处理整个文件夹 python main.py -i ./videos -o ./subtitles --recursive -m small + +高级用法示例: + # 完整流程：识别 + 双语 + ASS样式 + 硬嵌入 + python main.py -i video.mp4 --bilingual --formats srt,ass --embed-video --style chinese_large - # 指定语言和模型 - python main.py -i video.mp4 --language zh --model large + # 使用GPU加速硬嵌入 + python main.py -i video.mp4 --embed-video --embed-gpu + + # 仅导出字幕文件（已有识别结果时使用） + python main.py --export-only -s existing.srt --formats ass,vtt ''' ) @@ -631,6 +647,61 @@ def parse_args(): parser.add_argument('--no-skip-existing', action='store_false', dest='skip_existing', help='不跳过已存在字幕的文件') + # ========== 新增：双语字幕参数 ========== + bilingual_group = parser.add_argument_group('双语字幕选项') + bilingual_group.add_argument('--bilingual', action='store_true', + help='启用双语字幕生成') + bilingual_group.add_argument('--source-language', default='zh', + help='源语言代码 (默认: zh)') + bilingual_group.add_argument('--target-language', default='en', + help='目标语言代码 (默认: en)') + bilingual_group.add_argument('--bilingual-order', default='zh_en', + choices=['zh_en', 'en_zh'], + help='双语显示顺序: zh_en(中文在上英文在下) 或 en_zh(英文在上中文在下) (默认: zh_en)') + + # ========== 新增：导出格式参数 ========== + export_group = parser.add_argument_group('字幕导出选项') + export_group.add_argument('--formats', type=str, default='srt', + help='导出格式列表，逗号分隔 (支持: srt, ass, vtt; 默认: srt)') + export_group.add_argument('--export-only', action='store_true', + help='仅导出字幕模式，配合 -s/--subtitle 使用已有字幕文件') + export_group.add_argument('-s', '--subtitle', default=None, + help='已有字幕文件路径（用于--export-only模式）') + export_group.add_argument('--video-resolution', type=str, default='1920x1080', + help='视频分辨率，用于ASS样式 (默认: 1920x1080)') + + # ========== 新增：字幕硬嵌入参数 ========== + embed_group = parser.add_argument_group('字幕硬嵌入选项') + embed_group.add_argument('--embed-video', action='store_true', + help='将字幕硬嵌入到视频中') + embed_group.add_argument('--embed-output', default=None, + help='硬嵌入输出视频路径（可选）') + embed_group.add_argument('--embed-quality', default='high', + choices=['fast', 'normal', 'high', 'lossless'], + help='硬嵌入视频质量 (默认: high)') + embed_group.add_argument('--embed-gpu', action='store_true', + help='使用GPU加速硬嵌入（需要NVIDIA显卡和支持NVENC的FFmpeg）') + + # ========== 新增：样式美化参数 ========== + style_group = parser.add_argument_group('字幕样式选项') + style_group.add_argument('--style', default='default', + help='预设样式名称 (default, chinese_large, english, yellow, cyan, top)') + style_group.add_argument('--font-name', default=None, + help='字幕字体名称 (默认: Microsoft YaHei)') + style_group.add_argument('--font-size', type=int, default=None, + help='字幕字体大小 (默认: 48)') + style_group.add_argument('--font-color', type=str, default=None, + help='字体颜色 (十六进制，如: #FFFFFF 表示白色)') + style_group.add_argument('--outline-color', type=str, default=None, + help='描边颜色 (十六进制，如: #000000 表示黑色)') + style_group.add_argument('--outline-width', type=int, default=None, + help='描边宽度 (默认: 2)') + style_group.add_argument('--position', type=str, default='bottom', + choices=['top', 'middle', 'bottom'], + help='字幕位置 (默认: bottom)') + style_group.add_argument('--margin-v', type=int, default=None, + help='垂直边距（像素）(默认: 40)') + # 其他参数 parser.add_argument('--log-dir', default=None, help='日志目录（可选）') parser.add_argument('-v', '--verbose', action='store_true', help='输出详细日志') @@ -640,6 +711,10 @@ def parse_args(): help='显示系统信息') parser.add_argument('--check-deps', action='store_true', help='检查依赖并退出') + parser.add_argument('--list-styles', action='store_true', + help='列出所有可用的预设样式') + parser.add_argument('--list-formats', action='store_true', + help='列出所有支持的导出格式') return parser.parse_args() @@ -670,6 +745,89 @@ def handle_exception(e: Exception, verbose: bool = False): print("=" * 60 + "\n") +def list_preset_styles(): + """列出所有预设样式""" + print("\n" + "=" * 60) + print(" 可用的预设样式") + print("=" * 60) + + styles = { + 'default': '默认样式（微软雅黑，48号字，底部）', + 'chinese_large': '中文大字体（微软雅黑，48号字，大描边）', + 'english': '英文字体（Arial，28号字）', + 'yellow': '黄色字幕（微软雅黑，48号字，黄色）', + 'cyan': '青色字幕（微软雅黑，48号字，青色）', + 'top': '顶部显示（微软雅黑，40号字，顶部）', + } + + for name, desc in styles.items(): + print(f"\n {name}") + print(f" 说明: {desc}") + + print("\n" + "=" * 60) + print("使用方式: --style <样式名>") + print("自定义覆盖: --font-name, --font-size, --font-color 等参数") + print("=" * 60 + "\n") + + +def list_supported_formats(): + """列出所有支持的导出格式""" + print("\n" + "=" * 60) + print(" 支持的导出格式") + print("=" * 60) + + formats = { + 'srt': { + 'name': 'SubRip 字幕', + 'description': '最通用的字幕格式，纯文本无样式', + 'features': ['简单文本', '时间轴', '兼容所有播放器'], + }, + 'ass': { + 'name': 'Advanced SubStation Alpha', + 'description': '高级字幕格式，支持丰富的样式定义', + 'features': ['字体样式', '颜色', '描边', '位置', '动画效果'], + }, + 'vtt': { + 'name': 'WebVTT 字幕', + 'description': 'Web视频文本轨道，HTML5原生支持', + 'features': ['Web兼容', '简单CSS样式', '时间轴'], + }, + } + + for fmt, info in formats.items(): + print(f"\n {fmt.upper()} - {info['name']}") + print(f" 说明: {info['description']}") + print(f" 特性: {', '.join(info['features'])}") + + print("\n" + "=" * 60) + print("使用方式: --formats <格式1,格式2>") + print("示例: --formats srt,ass,vtt") + print("=" * 60 + "\n") + + +def parse_export_formats(formats_str: str) -> List[str]: + """解析导出格式字符串""" + formats = [] + for fmt in formats_str.split(','): + fmt = fmt.strip().lower() + if fmt in ['srt', 'ass', 'ssa', 'vtt']: + formats.append('ssa' if fmt == 'ass' else fmt) + else: + print(f"警告: 不支持的格式 '{fmt}'，已跳过") + return formats if formats else ['srt'] + + +def parse_video_resolution(res_str: str) -> Tuple[int, int]: + """解析视频分辨率字符串""" + try: + if 'x' in res_str: + w, h = res_str.lower().split('x') + return (int(w.strip()), int(h.strip())) + except: + pass + return (1920, 1080) + + def main(): """ 主函数入口 @@ -684,7 +842,6 @@ def main(): check_dependencies_and_exit() - # 尝试打印详细报告 try: from dependency_checker import print_dependency_report print("") @@ -694,6 +851,95 @@ def main(): sys.exit(0) + # 列出预设样式 + if args.list_styles: + list_preset_styles() + sys.exit(0) + + # 列出支持的格式 + if args.list_formats: + list_supported_formats() + sys.exit(0) + + # 仅导出模式：不需要输入视频 + if args.export_only: + if not args.subtitle: + print("\n错误: --export-only 模式需要指定 -s/--subtitle 参数") + print("使用方式: python main.py --export-only -s existing.srt --formats ass,vtt") + sys.exit(1) + + subtitle_path = Path(args.subtitle) + if not subtitle_path.exists(): + print(f"\n错误: 字幕文件不存在: {subtitle_path}") + sys.exit(1) + + try: + from video_subtitle_enhanced import VideoSubtitleAIEnhanced + from subtitle_exporter import SubtitleImporter, SubtitleExporter + from subtitle_models import create_default_style + + print(f"\n" + "=" * 60) + print(" 📤 仅导出模式") + print("=" * 60) + print(f"\n 输入字幕: {subtitle_path}") + + formats = parse_export_formats(args.formats) + video_res = parse_video_resolution(args.video_resolution) + + print(f" 导出格式: {', '.join([f.upper() for f in formats])}") + print(f" 视频分辨率: {video_res[0]}x{video_res[1]}") + + importer = SubtitleImporter() + doc = importer.import_from_file(str(subtitle_path)) + + output_base = args.output if args.output else str(subtitle_path.parent / subtitle_path.stem) + + enhancer = VideoSubtitleAIEnhanced() + preset_styles = enhancer.get_preset_styles() + style = preset_styles.get(args.style, create_default_style()) + + if args.font_name: + style.font_name = args.font_name + if args.font_size: + style.font_size = args.font_size + if args.font_color: + from subtitle_exporter import hex_to_ass_color + style.primary_color = hex_to_ass_color(args.font_color) + if args.outline_color: + from subtitle_exporter import hex_to_ass_color + style.outline_color = hex_to_ass_color(args.outline_color) + if args.outline_width is not None: + style.outline = args.outline_width + if args.margin_v is not None: + style.margin_v = args.margin_v + + if args.position == 'top': + style.alignment = 8 + elif args.position == 'middle': + style.alignment = 5 + else: + style.alignment = 2 + + output_files = enhancer.export_to_multiple_formats( + document=doc, + base_output_path=output_base, + formats=formats, + default_style=style, + video_resolution=video_res, + ) + + print(f"\n ✅ 导出完成!") + print(f"\n 生成的文件:") + for fmt, path in output_files.items(): + print(f" - {fmt.upper()}: {path}") + + print("\n" + "=" * 60 + "\n") + sys.exit(0) + + except Exception as e: + handle_exception(e, args.verbose) + sys.exit(1) + # 检查必需参数 if args.input is None: print("\n错误: 缺少必需参数 -i/--input") @@ -715,8 +961,16 @@ def main(): print("运行 --check-deps 检查依赖状态。") sys.exit(1) + # 解析导出格式 + export_formats = parse_export_formats(args.formats) + video_resolution = parse_video_resolution(args.video_resolution) + try: - # 创建实例 + # 创建增强版实例 + from video_subtitle_enhanced import VideoSubtitleAIEnhanced + from subtitle_models import create_default_style + + # 创建基础实例 ai = VideoSubtitleAI( model_name=args.model, device=args.device, @@ -768,26 +1022,148 @@ def main(): ai.initialize() ai.load_model() - # 处理 + # 获取样式 + enhancer = VideoSubtitleAIEnhanced( + base_recognizer=ai.speech_recognizer, + whisper_model=ai.speech_recognizer.model if ai.speech_recognizer else None, + ffmpeg_path=args.ffmpeg_path, + verbose=args.verbose, + ) + + preset_styles = enhancer.get_preset_styles() + style = preset_styles.get(args.style, create_default_style()) + + # 覆盖样式参数 + if args.font_name: + style.font_name = args.font_name + if args.font_size: + style.font_size = args.font_size + if args.font_color: + from subtitle_exporter import hex_to_ass_color + style.primary_color = hex_to_ass_color(args.font_color) + if args.outline_color: + from subtitle_exporter import hex_to_ass_color + style.outline_color = hex_to_ass_color(args.outline_color) + if args.outline_width is not None: + style.outline = args.outline_width + if args.margin_v is not None: + style.margin_v = args.margin_v + + if args.position == 'top': + style.alignment = 8 + elif args.position == 'middle': + style.alignment = 5 + else: + style.alignment = 2 + + # 处理单个文件 if input_path.is_file(): - # 单个文件 - result = ai.process_single_file( + print(f"\n" + "=" * 60) + print(" 🎬 处理单个文件") + print("=" * 60) + print(f"\n 输入文件: {input_path}") + print(f" 模型: {args.model}") + print(f" 语言: {args.language}") + print(f" 任务: {'翻译' if task == 'translate' else '转录'}") + print(f" 导出格式: {', '.join([f.upper() for f in export_formats])}") + if args.bilingual: + print(f" 双语字幕: 是 ({args.bilingual_order})") + print(f" 源语言: {args.source_language}") + print(f" 目标语言: {args.target_language}") + if args.embed_video: + print(f" 硬嵌入: 是") + print(f" 质量: {args.embed_quality}") + print(f" GPU加速: {'是' if args.embed_gpu else '否'}") + print(f" 样式: {args.style}") + + # 基础处理 + base_result = ai.process_single_file( input_path=str(input_path), output_path=args.output, language=args.language, task=task, ) - print(f"\n" + "=" * 60) - print(" ✅ 处理完成!") - print("=" * 60) - print(f"\n 输入文件: {result['input_path']}") - print(f" 输出字幕: {result['output_path']}") - print(f" 检测语言: {result['language_name']}") - print(f" 字幕段数: {result['segment_count']}") - print(f" 总字数: {result['total_characters']}") + + # 增强处理 + print(f"\n 📝 增强处理中...") + + # 获取识别结果 + temp_audio_path = None + is_temp = False + try: + temp_audio_path, is_temp = ai.audio_processor.process_media_file(str(input_path)) + recognition_result = ai.speech_recognizer.recognize( + audio_path=temp_audio_path, + language=args.language, + task=task, + verbose=args.verbose, + ) + finally: + if is_temp and temp_audio_path: + ai.audio_processor.cleanup_temp_file(temp_audio_path) + + # 增强处理 + enhanced_result = enhancer.process_single_file_enhanced( + recognition_result=recognition_result, + input_path=str(input_path), + output_path=args.output, + enable_bilingual=args.bilingual, + source_language=args.source_language, + target_language=args.target_language, + bilingual_order=args.bilingual_order, + export_formats=export_formats, + enable_embedding=args.embed_video, + embedding_output_path=args.embed_output, + style=style, + video_quality=args.embed_quality, + use_gpu_for_embedding=args.embed_gpu, + ) + + # 输出结果 + print(f"\n ✅ 处理完成!") + print(f"\n 处理结果:") + print(f" - 检测语言: {base_result.get('language_name', '未知')}") + print(f" - 字幕段数: {base_result.get('segment_count', 0)}") + print(f" - 总字数: {base_result.get('total_characters', 0)}") + + if enhanced_result.language_detection.get('is_mixed'): + langs = enhanced_result.language_detection.get('languages', []) + print(f"\n ⚠️ 检测到多语言混合场景: {', '.join(langs)}") + + if enhanced_result.output_files: + print(f"\n 📂 生成的文件:") + for fmt, path in enhanced_result.output_files.items(): + if fmt == 'embedded_video': + print(f" - [视频] 硬嵌入视频: {path}") + else: + print(f" - [字幕] {fmt.upper()}: {path}") + + if enhanced_result.warnings: + print(f"\n ⚠️ 警告:") + for w in enhanced_result.warnings: + print(f" - {w}") + + if enhanced_result.errors: + print(f"\n ❌ 错误:") + for e in enhanced_result.errors: + print(f" - {e}") + print("\n" + "=" * 60 + "\n") + else: - # 文件夹 + # 文件夹批量处理 + print(f"\n" + "=" * 60) + print(" 📁 批量处理文件夹") + print("=" * 60) + print(f"\n 输入文件夹: {input_path}") + print(f" 模型: {args.model}") + print(f" 语言: {args.language}") + print(f" 递归处理: {'是' if args.recursive else '否'}") + print(f" 跳过已有: {'是' if args.skip_existing else '否'}") + print(f" 导出格式: {', '.join([f.upper() for f in export_formats])}") + if args.bilingual: + print(f" 双语字幕: 是 ({args.bilingual_order})") + result = ai.process_folder( input_folder=str(input_path), output_folder=args.output, @@ -796,6 +1172,7 @@ def main(): task=task, skip_existing=args.skip_existing, ) + print(f"\n" + "=" * 60) print(" ✅ 批量处理完成!") print("=" * 60) @@ -811,10 +1188,12 @@ def main(): for err in result['errors']: print(f" - {Path(err['input_path']).name}: {err.get('error', '未知错误')}") - print("\n" + "=" * 60 + "\n") + print("\n" + "=" * 60) + print(" 注意: 增强版功能（双语字幕、多格式导出、硬嵌入）") + print(" 目前仅支持单文件处理模式。") + print("=" * 60 + "\n") except FFmpegNotFoundError as e: - # 已经在 initialize 中处理过了 sys.exit(1) except Exception as e: handle_exception(e, args.verbose) diff --git a/requirements.txt b/requirements.txt index bf40f5e..297a7e5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -48,6 +48,39 @@ colorama>=0.4.6 # 版本比较（用于依赖检查） packaging>=23.0 +# ============================================================================ +# 双语字幕翻译功能可选依赖 +# ============================================================================ + +# 在线翻译API支持 +# requests>=2.28.0 # 用于百度/Google/DeepL翻译API + +# 本地翻译模型（Helsinki-NLP OPUS-MT） +# transformers>=4.30.0 +# sentencepiece>=0.1.99 +# torch>=2.0.0 # 已在必需依赖中 + +# ============================================================================ +# 增强功能说明 +# ============================================================================ +# +# 新增功能模块： +# 1. subtitle_models.py - 统一数据模型（字幕段、样式、文档） +# 2. subtitle_translator.py - 翻译模块（支持多后端） +# 3. subtitle_editor.py - 字幕编辑（撤销/重做、时间轴调整） +# 4. subtitle_exporter.py - 多格式导出（SRT/ASS/VTT） +# 5. subtitle_embedder.py - 字幕硬嵌入（FFmpeg烧录） +# 6. language_detector.py - 语言检测优化（多语言混合） +# 7. video_subtitle_enhanced.py - 增强版主模块 +# +# 使用示例： +# - 生成双语字幕：python main.py -i video.mp4 --bilingual +# - 导出多种格式：python main.py -i video.mp4 --formats srt,ass,vtt +# - 字幕硬嵌入：python main.py -i video.mp4 --embed-video +# - 自定义样式：python main.py -i video.mp4 --style yellow --font-size 48 +# - 仅导出模式：python main.py --export-only -s existing.srt --formats ass +# + # ============================================================================ # 系统工具（必须单独安装） # ============================================================================ diff --git a/subtitle_editor.py b/subtitle_editor.py new file mode 100644 index 0000000..b1cb98c --- /dev/null +++ b/subtitle_editor.py @@ -0,0 +1,700 @@ +""" +字幕编辑模块 +提供字幕时间轴调整、文本修改、分句合并/拆分等编辑功能 +""" + +import logging +from typing import List, Dict, Any, Optional, Tuple, Callable +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +import re + +from subtitle_models import SubtitleSegment, SubtitleDocument, SubtitleStyle + +logger = logging.getLogger(__name__) + + +class EditOperation(Enum): + """编辑操作类型枚举""" + ADD = "add" + DELETE = "delete" + MODIFY = "modify" + SPLIT = "split" + MERGE = "merge" + SHIFT = "shift" + STYLE = "style" + + +@dataclass +class EditHistoryItem: + """编辑历史记录项""" + operation: EditOperation + index: int + old_data: Optional[Dict[str, Any]] = None + new_data: Optional[Dict[str, Any]] = None + description: str = "" + timestamp: float = 0.0 + + +class SubtitleEditor: + """ + 字幕编辑器类 + 提供完整的字幕编辑功能： + - 时间轴调整（整体偏移、单个调整） + - 文本修改 + - 分句合并/拆分 + - 撤销/重做 + """ + + def __init__(self, document: Optional[SubtitleDocument] = None): + self.document = document or SubtitleDocument() + self._history: List[EditHistoryItem] = [] + self._history_index: int = -1 + self._max_history: int = 100 + + def load_document(self, document: SubtitleDocument): + self.document = document + self._history.clear() + self._history_index = -1 + + def load_from_segments(self, segments: List[Any]): + from subtitle_models import SubtitleSegment + + doc_segments = [] + for seg in segments: + if isinstance(seg, dict): + doc_segments.append(SubtitleSegment.from_dict(seg)) + else: + doc_segments.append(seg) + + self.document = SubtitleDocument(segments=doc_segments) + self._history.clear() + self._history_index = -1 + + def load_from_srt(self, srt_path: str, encoding: str = 'utf-8'): + segments = self._parse_srt_file(srt_path, encoding) + self.load_from_segments(segments) + + def _parse_srt_file(self, srt_path: str, encoding: str = 'utf-8') -> List[SubtitleSegment]: + srt_path = Path(srt_path).resolve() + if not srt_path.exists(): + raise FileNotFoundError(f"SRT文件不存在: {srt_path}") + + segments = [] + + with open(srt_path, 'r', encoding=encoding) as f: + content = f.read() + + blocks = re.split(r'\n\n+', content.strip()) + + for block in blocks: + lines = block.strip().split('\n') + if len(lines) < 2: + continue + + seg_id = len(segments) + time_line_idx = 0 + + if lines[0].strip().isdigit(): + seg_id = int(lines[0].strip()) - 1 + time_line_idx = 1 + + if time_line_idx >= len(lines): + continue + + time_line = lines[time_line_idx] + time_match = re.match( + r'(\d{2}:\d{2}:\d{2}[,.]\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}[,.]\d{3})', + time_line + ) + + if time_match: + start_time = self._parse_time_str(time_match.group(1)) + end_time = self._parse_time_str(time_match.group(2)) + + text_lines = lines[time_line_idx + 1:] + text = '\n'.join(text_lines).strip() + + segment = SubtitleSegment( + id=seg_id, + start=start_time, + end=end_time, + text=text, + ) + segments.append(segment) + + logger.info(f"成功解析SRT文件: {srt_path}，共 {len(segments)} 段") + return segments + + def _parse_time_str(self, time_str: str) -> float: + time_str = time_str.strip() + time_str = time_str.replace(',', '.') + + parts = time_str.split(':') + if len(parts) == 3: + hours = int(parts[0]) + minutes = int(parts[1]) + seconds = float(parts[2]) + return hours * 3600 + minutes * 60 + seconds + elif len(parts) == 2: + minutes = int(parts[0]) + seconds = float(parts[1]) + return minutes * 60 + seconds + else: + return float(time_str) + + def _record_history( + self, + operation: EditOperation, + index: int, + old_data: Optional[Dict[str, Any]] = None, + new_data: Optional[Dict[str, Any]] = None, + description: str = "", + ): + if self._history_index < len(self._history) - 1: + self._history = self._history[:self._history_index + 1] + + import time + item = EditHistoryItem( + operation=operation, + index=index, + old_data=old_data, + new_data=new_data, + description=description, + timestamp=time.time(), + ) + + self._history.append(item) + self._history_index = len(self._history) - 1 + + if len(self._history) > self._max_history: + self._history = self._history[-self._max_history:] + self._history_index = self._max_history - 1 + + def undo(self) -> bool: + if self._history_index < 0: + logger.warning("没有可撤销的操作") + return False + + item = self._history[self._history_index] + + if item.operation == EditOperation.ADD: + if 0 <= item.index < len(self.document.segments): + self.document.segments.pop(item.index) + self.document._reindex_segments() + + elif item.operation == EditOperation.DELETE: + if item.old_data: + from subtitle_models import SubtitleSegment + seg = SubtitleSegment.from_dict(item.old_data) + self.document.segments.insert(item.index, seg) + self.document._reindex_segments() + + elif item.operation == EditOperation.MODIFY: + if item.old_data and 0 <= item.index < len(self.document.segments): + from subtitle_models import SubtitleSegment + seg = SubtitleSegment.from_dict(item.old_data) + self.document.segments[item.index] = seg + + elif item.operation == EditOperation.SPLIT: + if 0 <= item.index < len(self.document.segments) - 1: + self.document.merge_adjacent_segments(item.index, gap_threshold=0) + + elif item.operation == EditOperation.MERGE: + if item.old_data and 'segments' in item.old_data: + from subtitle_models import SubtitleSegment + self.document.segments.pop(item.index) + for i, seg_data in enumerate(item.old_data['segments']): + seg = SubtitleSegment.from_dict(seg_data) + self.document.segments.insert(item.index + i, seg) + self.document._reindex_segments() + + elif item.operation == EditOperation.SHIFT: + if item.old_data and 'shift_seconds' in item.old_data: + shift_back = -item.old_data['shift_seconds'] + self.document.shift_all_timings(shift_back) + + self._history_index -= 1 + logger.info(f"已撤销: {item.description}") + return True + + def redo(self) -> bool: + if self._history_index >= len(self._history) - 1: + logger.warning("没有可重做的操作") + return False + + self._history_index += 1 + item = self._history[self._history_index] + + logger.info(f"已重做: {item.description}") + return True + + @property + def can_undo(self) -> bool: + return self._history_index >= 0 + + @property + def can_redo(self) -> bool: + return self._history_index < len(self._history) - 1 + + def get_segment(self, index: int) -> Optional[SubtitleSegment]: + return self.document.get_segment(index) + + def update_segment_text(self, index: int, new_text: str) -> bool: + if index < 0 or index >= len(self.document.segments): + return False + + segment = self.document.segments[index] + old_data = segment.to_dict() + + segment.text = new_text + + self._record_history( + operation=EditOperation.MODIFY, + index=index, + old_data=old_data, + new_data=segment.to_dict(), + description=f"修改第 {index + 1} 段文本", + ) + + logger.info(f"已修改第 {index + 1} 段文本") + return True + + def update_segment_bilingual_text( + self, + index: int, + text_zh: Optional[str] = None, + text_en: Optional[str] = None, + ) -> bool: + if index < 0 or index >= len(self.document.segments): + return False + + segment = self.document.segments[index] + old_data = segment.to_dict() + + if text_zh is not None: + segment.text_zh = text_zh + if text_en is not None: + segment.text_en = text_en + + segment.is_bilingual = True + + self._record_history( + operation=EditOperation.MODIFY, + index=index, + old_data=old_data, + new_data=segment.to_dict(), + description=f"修改第 {index + 1} 段双语文本", + ) + + return True + + def update_segment_time(self, index: int, start: Optional[float] = None, end: Optional[float] = None) -> bool: + if index < 0 or index >= len(self.document.segments): + return False + + segment = self.document.segments[index] + old_data = segment.to_dict() + + if start is not None: + segment.start = max(0.0, start) + if end is not None: + segment.end = max(segment.start, end) + + self._record_history( + operation=EditOperation.MODIFY, + index=index, + old_data=old_data, + new_data=segment.to_dict(), + description=f"调整第 {index + 1} 段时间轴", + ) + + logger.info(f"已调整第 {index + 1} 段时间轴: {segment.start:.3f} -> {segment.end:.3f}") + return True + + def shift_all_segments(self, seconds: float) -> int: + if seconds == 0: + return 0 + + old_data = {'shift_seconds': seconds} + + self.document.shift_all_timings(seconds) + + count = len(self.document.segments) + + self._record_history( + operation=EditOperation.SHIFT, + index=0, + old_data=old_data, + description=f"整体偏移时间轴: {seconds:+.3f} 秒", + ) + + logger.info(f"已整体偏移 {count} 个字幕段: {seconds:+.3f} 秒") + return count + + def shift_segments_range(self, start_index: int, end_index: int, seconds: float) -> int: + if seconds == 0: + return 0 + + if start_index < 0 or end_index >= len(self.document.segments) or start_index > end_index: + return 0 + + old_data = { + 'segments': [self.document.segments[i].to_dict() for i in range(start_index, end_index + 1)] + } + + count = 0 + for i in range(start_index, end_index + 1): + seg = self.document.segments[i] + seg.start = max(0.0, seg.start + seconds) + seg.end = max(seg.start, seg.end + seconds) + count += 1 + + self._record_history( + operation=EditOperation.SHIFT, + index=start_index, + old_data=old_data, + description=f"偏移第 {start_index + 1}-{end_index + 1} 段时间轴: {seconds:+.3f} 秒", + ) + + return count + + def split_segment(self, index: int, split_seconds: float) -> Tuple[bool, int]: + if index < 0 or index >= len(self.document.segments): + return False, -1 + + old_data = self.document.segments[index].to_dict() + + success, new_index = self.document.split_segment(index, split_seconds) + + if success: + self._record_history( + operation=EditOperation.SPLIT, + index=index, + old_data=old_data, + description=f"拆分第 {index + 1} 段（在 {split_seconds:.3f} 秒处）", + ) + logger.info(f"已拆分第 {index + 1} 段，新段索引: {new_index + 1}") + + return success, new_index + + def split_segment_at_ratio(self, index: int, ratio: float = 0.5) -> Tuple[bool, int]: + if index < 0 or index >= len(self.document.segments): + return False, -1 + + segment = self.document.segments[index] + duration = segment.end - segment.start + split_seconds = segment.start + duration * ratio + + return self.split_segment(index, split_seconds) + + def merge_adjacent_segments(self, index: int, gap_threshold: float = 0.5) -> bool: + if index < 0 or index >= len(self.document.segments) - 1: + return False + + old_data = { + 'segments': [ + self.document.segments[index].to_dict(), + self.document.segments[index + 1].to_dict(), + ] + } + + success = self.document.merge_adjacent_segments(index, gap_threshold) + + if success: + self._record_history( + operation=EditOperation.MERGE, + index=index, + old_data=old_data, + description=f"合并第 {index + 1} 和 {index + 2} 段", + ) + logger.info(f"已合并第 {index + 1} 和 {index + 2} 段") + + return success + + def merge_multiple_segments(self, start_index: int, end_index: int) -> bool: + if start_index < 0 or end_index >= len(self.document.segments) or start_index >= end_index: + return False + + old_data = { + 'segments': [self.document.segments[i].to_dict() for i in range(start_index, end_index + 1)] + } + + merged_text = [] + merged_text_zh = [] + merged_text_en = [] + start_time = self.document.segments[start_index].start + end_time = self.document.segments[end_index].end + is_bilingual = any(seg.is_bilingual for seg in self.document.segments[start_index:end_index + 1]) + + for i in range(start_index, end_index + 1): + seg = self.document.segments[i] + merged_text.append(seg.text) + if seg.is_bilingual: + merged_text_zh.append(seg.text_zh) + merged_text_en.append(seg.text_en) + + for _ in range(start_index, end_index): + self.document.segments.pop(start_index) + + merged_seg = self.document.segments[start_index] + merged_seg.start = start_time + merged_seg.end = end_time + merged_seg.text = ' '.join(merged_text) + if is_bilingual: + merged_seg.text_zh = ' '.join(merged_text_zh) + merged_seg.text_en = ' '.join(merged_text_en) + merged_seg.is_bilingual = True + + self.document._reindex_segments() + + self._record_history( + operation=EditOperation.MERGE, + index=start_index, + old_data=old_data, + description=f"合并第 {start_index + 1}-{end_index + 1} 段", + ) + + logger.info(f"已合并第 {start_index + 1}-{end_index + 1} 段") + return True + + def add_segment(self, index: int, segment: SubtitleSegment) -> bool: + success = self.document.insert_segment(index, segment) + + if success: + self._record_history( + operation=EditOperation.ADD, + index=index, + new_data=segment.to_dict(), + description=f"在位置 {index + 1} 插入新段", + ) + logger.info(f"已在位置 {index + 1} 插入新段") + + return success + + def delete_segment(self, index: int) -> bool: + if index < 0 or index >= len(self.document.segments): + return False + + old_data = self.document.segments[index].to_dict() + + success = self.document.remove_segment(index) + + if success: + self._record_history( + operation=EditOperation.DELETE, + index=index, + old_data=old_data, + description=f"删除第 {index + 1} 段", + ) + logger.info(f"已删除第 {index + 1} 段") + + return success + + def delete_segments_range(self, start_index: int, end_index: int) -> int: + if start_index < 0 or end_index >= len(self.document.segments) or start_index > end_index: + return 0 + + old_data = { + 'segments': [self.document.segments[i].to_dict() for i in range(start_index, end_index + 1)] + } + + count = 0 + for _ in range(start_index, end_index + 1): + if start_index < len(self.document.segments): + self.document.segments.pop(start_index) + count += 1 + + self.document._reindex_segments() + + self._record_history( + operation=EditOperation.DELETE, + index=start_index, + old_data=old_data, + description=f"删除第 {start_index + 1}-{end_index + 1} 段", + ) + + logger.info(f"已删除 {count} 个段") + return count + + def find_segments_by_text(self, keyword: str, case_sensitive: bool = False) -> List[int]: + results = [] + + for idx, seg in enumerate(self.document.segments): + text = seg.text + if not case_sensitive: + text = text.lower() + keyword = keyword.lower() + + if keyword in text: + results.append(idx) + continue + + if seg.is_bilingual: + text_zh = seg.text_zh + text_en = seg.text_en + if not case_sensitive: + text_zh = text_zh.lower() + text_en = text_en.lower() + + if keyword in text_zh or keyword in text_en: + results.append(idx) + + return results + + def replace_text(self, old_text: str, new_text: str, case_sensitive: bool = False) -> int: + count = 0 + + for idx, seg in enumerate(self.document.segments): + old_data = seg.to_dict() + modified = False + + text = seg.text + if case_sensitive: + if old_text in text: + seg.text = text.replace(old_text, new_text) + modified = True + else: + import re + pattern = re.compile(re.escape(old_text), re.IGNORECASE) + if pattern.search(text): + seg.text = pattern.sub(new_text, text) + modified = True + + if seg.is_bilingual: + if case_sensitive: + if old_text in seg.text_zh: + seg.text_zh = seg.text_zh.replace(old_text, new_text) + modified = True + if old_text in seg.text_en: + seg.text_en = seg.text_en.replace(old_text, new_text) + modified = True + else: + if pattern.search(seg.text_zh): + seg.text_zh = pattern.sub(new_text, seg.text_zh) + modified = True + if pattern.search(seg.text_en): + seg.text_en = pattern.sub(new_text, seg.text_en) + modified = True + + if modified: + count += 1 + self._record_history( + operation=EditOperation.MODIFY, + index=idx, + old_data=old_data, + new_data=seg.to_dict(), + description=f"替换第 {idx + 1} 段文本: '{old_text}' -> '{new_text}'", + ) + + logger.info(f"已替换 {count} 个段中的文本") + return count + + def adjust_segment_gaps(self, min_gap: float = 0.1, max_gap: float = 2.0) -> int: + count = 0 + + for i in range(len(self.document.segments) - 1): + current = self.document.segments[i] + next_seg = self.document.segments[i + 1] + + gap = next_seg.start - current.end + + if gap < min_gap: + old_data = { + 'segments': [current.to_dict(), next_seg.to_dict()] + } + + current.end = next_seg.start - min_gap + + self._record_history( + operation=EditOperation.MODIFY, + index=i, + old_data=old_data, + description=f"调整第 {i + 1} 和 {i + 2} 段间隔 (扩大)", + ) + count += 1 + + elif gap > max_gap: + old_data = { + 'segments': [current.to_dict(), next_seg.to_dict()] + } + + extension = (gap - max_gap) / 2 + current.end += extension + next_seg.start -= extension + + self._record_history( + operation=EditOperation.MODIFY, + index=i, + old_data=old_data, + description=f"调整第 {i + 1} 和 {i + 2} 段间隔 (缩小)", + ) + count += 1 + + logger.info(f"已调整 {count} 个段间隔") + return count + + def set_segment_style(self, index: int, style_name: str) -> bool: + if index < 0 or index >= len(self.document.segments): + return False + + if style_name not in self.document.styles: + logger.warning(f"样式 '{style_name}' 不存在") + return False + + segment = self.document.segments[index] + old_data = segment.to_dict() + + segment.style_name = style_name + + self._record_history( + operation=EditOperation.STYLE, + index=index, + old_data=old_data, + new_data=segment.to_dict(), + description=f"设置第 {index + 1} 段样式为 '{style_name}'", + ) + + return True + + def add_style(self, style: SubtitleStyle): + self.document.add_style(style) + logger.info(f"已添加样式: {style.name}") + + def get_statistics(self) -> Dict[str, Any]: + return { + 'segment_count': self.document.segment_count, + 'total_duration': self.document.total_duration, + 'total_characters': self.document.total_characters, + 'is_bilingual': self.document.is_bilingual, + 'original_language': self.document.original_language, + 'translated_language': self.document.translated_language, + 'style_count': len(self.document.styles), + 'history_items': len(self._history), + 'can_undo': self.can_undo, + 'can_redo': self.can_redo, + } + + def export_to_document(self) -> SubtitleDocument: + return SubtitleDocument.from_dict(self.document.to_dict()) + + +def create_editor_from_segments( + segments: List[Any], + is_bilingual: bool = False, +) -> SubtitleEditor: + """ + 从字幕段列表创建编辑器 + + Args: + segments: 字幕段列表（可以是字典、SubtitleSegment对象等） + is_bilingual: 是否为双语字幕 + + Returns: + SubtitleEditor实例 + """ + from subtitle_models import create_document_from_segments + + document = create_document_from_segments(segments, is_bilingual=is_bilingual) + return SubtitleEditor(document=document) diff --git a/subtitle_embedder.py b/subtitle_embedder.py new file mode 100644 index 0000000..b10fdd5 --- /dev/null +++ b/subtitle_embedder.py @@ -0,0 +1,491 @@ +""" +字幕硬嵌入模块 +使用FFmpeg将字幕烧录到视频文件中 +""" + +import os +import sys +import subprocess +import tempfile +import shutil +from pathlib import Path +from typing import List, Dict, Any, Optional, Tuple +from dataclasses import dataclass +from enum import Enum +import logging + +from subtitle_models import SubtitleStyle, SubtitleDocument + +logger = logging.getLogger(__name__) + + +class SubtitlePosition(Enum): + """字幕位置枚举""" + BOTTOM = "bottom" + TOP = "top" + MIDDLE = "middle" + + +@dataclass +class EmbeddingConfig: + """字幕硬嵌入配置""" + font_name: str = "Microsoft YaHei" + font_size: int = 24 + primary_color: str = "&H00FFFFFF" + outline_color: str = "&H00000000" + back_color: str = "&H00000000" + outline_width: int = 2 + shadow: int = 1 + position: SubtitlePosition = SubtitlePosition.BOTTOM + margin_v: int = 40 + margin_h: int = 20 + bold: bool = False + italic: bool = False + video_quality: str = "high" + video_codec: str = "libx264" + audio_codec: str = "aac" + audio_bitrate: str = "192k" + crf: int = 23 + preset: str = "medium" + use_gpu: bool = False + + +class SubtitleEmbedder: + """ + 字幕硬嵌入器 + 使用FFmpeg将字幕烧录到视频中 + """ + + PRESET_QUALITY = { + 'fast': {'crf': 28, 'preset': 'veryfast'}, + 'normal': {'crf': 23, 'preset': 'medium'}, + 'high': {'crf': 20, 'preset': 'slow'}, + 'lossless': {'crf': 0, 'preset': 'ultrafast'}, + } + + def __init__(self, ffmpeg_path: Optional[str] = None): + self.ffmpeg_path = ffmpeg_path or self._find_ffmpeg() + self.ffprobe_path = self._find_ffprobe() + + if not self.ffmpeg_path: + raise RuntimeError("未找到FFmpeg，请确保FFmpeg已安装并添加到系统PATH") + + def _find_ffmpeg(self) -> Optional[str]: + ffmpeg_exe = 'ffmpeg.exe' if sys.platform == 'win32' else 'ffmpeg' + + which_result = shutil.which('ffmpeg') + if which_result: + return which_result + + if sys.platform == 'win32': + common_paths = [ + r'C:\ffmpeg\bin\ffmpeg.exe', + r'C:\Program Files\ffmpeg\bin\ffmpeg.exe', + r'C:\Program Files (x86)\ffmpeg\bin\ffmpeg.exe', + ] + for path in common_paths: + if os.path.exists(path): + return path + + return None + + def _find_ffprobe(self) -> Optional[str]: + ffprobe_exe = 'ffprobe.exe' if sys.platform == 'win32' else 'ffprobe' + + which_result = shutil.which('ffprobe') + if which_result: + return which_result + + if self.ffmpeg_path: + ffmpeg_dir = Path(self.ffmpeg_path).parent + ffprobe_candidate = ffmpeg_dir / ffprobe_exe + if ffprobe_candidate.exists(): + return str(ffprobe_candidate) + + return None + + def _run_ffmpeg_command( + self, + cmd: List[str], + description: str = "执行FFmpeg命令", + timeout: Optional[int] = None, + ) -> Tuple[int, str, str]: + full_cmd = [self.ffmpeg_path] + cmd + + logger.debug(f"执行FFmpeg命令: {' '.join(full_cmd)}") + + try: + startupinfo = None + if sys.platform == 'win32': + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + startupinfo.wShowWindow = subprocess.SW_HIDE + + result = subprocess.run( + full_cmd, + capture_output=True, + text=True, + timeout=timeout, + startupinfo=startupinfo, + ) + + if result.returncode != 0: + logger.error(f"FFmpeg命令失败: {description}") + logger.error(f"命令: {' '.join(full_cmd)}") + logger.error(f"错误输出: {result.stderr}") + raise RuntimeError(f"FFmpeg执行失败: {result.stderr}") + + return result.returncode, result.stdout, result.stderr + + except FileNotFoundError as e: + raise RuntimeError(f"找不到FFmpeg: {e}") + except subprocess.TimeoutExpired: + raise RuntimeError(f"FFmpeg命令超时: {description}") + + def get_video_info(self, video_path: str) -> Dict[str, Any]: + video_path = Path(video_path).resolve() + + if not video_path.exists(): + raise FileNotFoundError(f"视频文件不存在: {video_path}") + + info = { + 'path': str(video_path), + 'name': video_path.name, + 'size_bytes': video_path.stat().st_size, + } + + if self.ffprobe_path: + try: + cmd = [ + self.ffprobe_path, + '-v', 'quiet', + '-print_format', 'json', + '-show_format', + '-show_streams', + str(video_path), + ] + + startupinfo = None + if sys.platform == 'win32': + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + startupinfo.wShowWindow = subprocess.SW_HIDE + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=30, + startupinfo=startupinfo, + ) + + if result.returncode == 0: + import json + probe_data = json.loads(result.stdout) + + format_info = probe_data.get('format', {}) + streams = probe_data.get('streams', []) + + video_stream = None + audio_stream = None + + for stream in streams: + if stream.get('codec_type') == 'video': + video_stream = stream + elif stream.get('codec_type') == 'audio': + audio_stream = stream + + info['duration_seconds'] = float(format_info.get('duration', 0)) + info['bitrate'] = int(format_info.get('bit_rate', 0)) + + if video_stream: + info['video_codec'] = video_stream.get('codec_name') + info['width'] = int(video_stream.get('width', 0)) + info['height'] = int(video_stream.get('height', 0)) + info['fps'] = self._parse_fps(video_stream.get('r_frame_rate', '0/0')) + + if audio_stream: + info['audio_codec'] = audio_stream.get('codec_name') + info['audio_channels'] = int(audio_stream.get('channels', 0)) + info['audio_sample_rate'] = int(audio_stream.get('sample_rate', 0)) + + except Exception as e: + logger.warning(f"使用ffprobe获取视频信息失败: {e}") + + return info + + def _parse_fps(self, fps_str: str) -> float: + try: + if '/' in fps_str: + num, den = fps_str.split('/') + if int(den) == 0: + return 0.0 + return float(num) / float(den) + return float(fps_str) + except: + return 0.0 + + def _build_style_options(self, config: EmbeddingConfig) -> Dict[str, str]: + options = {} + + options['FontName'] = config.font_name + options['FontSize'] = str(config.font_size) + options['PrimaryColour'] = config.primary_color + options['OutlineColour'] = config.outline_color + options['BackColour'] = config.back_color + options['Outline'] = str(config.outline_width) + options['Shadow'] = str(config.shadow) + options['Bold'] = '1' if config.bold else '0' + options['Italic'] = '1' if config.italic else '0' + + if config.position == SubtitlePosition.TOP: + options['Alignment'] = '8' + elif config.position == SubtitlePosition.MIDDLE: + options['Alignment'] = '5' + else: + options['Alignment'] = '2' + + options['MarginV'] = str(config.margin_v) + + return options + + def _build_filter_complex( + self, + subtitle_path: str, + config: EmbeddingConfig, + video_info: Dict[str, Any], + ) -> str: + subtitle_ext = Path(subtitle_path).suffix.lower() + + style_opts = self._build_style_options(config) + style_str = ','.join([f'{k}={v}' for k, v in style_opts.items()]) + + if subtitle_ext == '.ass': + filter_str = f"ass='{subtitle_path}'" + else: + filter_str = f"subtitles='{subtitle_path}':force_style='{style_str}'" + + return filter_str + + def embed_subtitles( + self, + video_path: str, + subtitle_path: str, + output_path: str, + config: Optional[EmbeddingConfig] = None, + progress_callback: Optional[callable] = None, + ) -> Dict[str, Any]: + video_path = Path(video_path).resolve() + subtitle_path = Path(subtitle_path).resolve() + output_path = Path(output_path).resolve() + + if not video_path.exists(): + raise FileNotFoundError(f"视频文件不存在: {video_path}") + + if not subtitle_path.exists(): + raise FileNotFoundError(f"字幕文件不存在: {subtitle_path}") + + output_path.parent.mkdir(parents=True, exist_ok=True) + + if config is None: + config = EmbeddingConfig() + + video_info = self.get_video_info(str(video_path)) + + subtitle_ext = subtitle_path.suffix.lower() + + temp_ass_path = None + use_temp_ass = False + + if subtitle_ext not in ['.ass', '.ssa'] and subtitle_ext != '.srt': + from subtitle_exporter import SubtitleImporter, SubtitleExporter + + importer = SubtitleImporter() + doc = importer.import_from_file(str(subtitle_path)) + + temp_dir = tempfile.gettempdir() + temp_ass_path = Path(temp_dir) / f"temp_subtitle_{os.getpid()}.ass" + + exporter = SubtitleExporter(doc) + exporter.export_to_file(str(temp_ass_path), format_type='ass') + + subtitle_path = temp_ass_path + use_temp_ass = True + subtitle_ext = '.ass' + + quality_settings = self.PRESET_QUALITY.get( + config.video_quality.lower(), + self.PRESET_QUALITY['normal'] + ) + + cmd = [] + + cmd.extend(['-i', str(video_path)]) + + subtitle_path_str = str(subtitle_path).replace(':', '\\\\:').replace('\\', '/') + + if subtitle_ext in ['.ass', '.ssa']: + filter_complex = f"ass='{subtitle_path_str}'" + else: + style_opts = self._build_style_options(config) + style_str = ','.join([f'{k}={v}' for k, v in style_opts.items()]) + filter_complex = f"subtitles='{subtitle_path_str}':force_style='{style_str}'" + + cmd.extend(['-vf', filter_complex]) + + if config.use_gpu: + cmd.extend(['-c:v', 'h264_nvenc']) + cmd.extend(['-preset', 'p6']) + cmd.extend(['-cq', '22']) + else: + cmd.extend(['-c:v', config.video_codec]) + cmd.extend(['-crf', str(config.crf if config.crf else quality_settings['crf'])]) + cmd.extend(['-preset', quality_settings['preset']]) + + cmd.extend(['-c:a', config.audio_codec]) + cmd.extend(['-b:a', config.audio_bitrate]) + + cmd.extend(['-y', str(output_path)]) + + logger.info(f"开始烧录字幕: {video_path.name} -> {output_path.name}") + logger.info(f"字幕文件: {subtitle_path.name}") + + try: + self._run_ffmpeg_command(cmd, "烧录字幕到视频") + + if not output_path.exists(): + raise RuntimeError(f"输出文件未生成: {output_path}") + + result = { + 'success': True, + 'input_video': str(video_path), + 'input_subtitle': str(subtitle_path), + 'output_video': str(output_path), + 'output_size_bytes': output_path.stat().st_size, + 'video_info': video_info, + 'config': { + 'font_name': config.font_name, + 'font_size': config.font_size, + 'position': config.position.value, + 'video_quality': config.video_quality, + }, + } + + logger.info(f"字幕烧录完成: {output_path}") + return result + + finally: + if use_temp_ass and temp_ass_path and temp_ass_path.exists(): + try: + os.remove(temp_ass_path) + except Exception: + pass + + def embed_subtitles_with_style( + self, + video_path: str, + subtitle_path: str, + output_path: str, + style: SubtitleStyle, + video_quality: str = "high", + use_gpu: bool = False, + ) -> Dict[str, Any]: + position_map = { + 1: SubtitlePosition.LEFT, + 2: SubtitlePosition.BOTTOM, + 3: SubtitlePosition.RIGHT, + 5: SubtitlePosition.MIDDLE, + 8: SubtitlePosition.TOP, + } + + position = position_map.get(style.alignment, SubtitlePosition.BOTTOM) + + config = EmbeddingConfig( + font_name=style.font_name, + font_size=style.font_size, + primary_color=style.primary_color, + outline_color=style.outline_color, + back_color=style.back_color, + outline_width=style.outline, + shadow=style.shadow, + position=position, + margin_v=style.margin_v, + bold=style.bold != 0, + italic=style.italic != 0, + video_quality=video_quality, + use_gpu=use_gpu, + ) + + return self.embed_subtitles( + video_path=video_path, + subtitle_path=subtitle_path, + output_path=output_path, + config=config, + ) + + def create_preset_config( + self, + preset_name: str = "default", + ) -> EmbeddingConfig: + presets = { + 'default': EmbeddingConfig(), + 'chinese': EmbeddingConfig( + font_name="Microsoft YaHei", + font_size=24, + outline_width=2, + position=SubtitlePosition.BOTTOM, + ), + 'english': EmbeddingConfig( + font_name="Arial", + font_size=20, + outline_width=1, + position=SubtitlePosition.BOTTOM, + ), + 'top': EmbeddingConfig( + font_name="Microsoft YaHei", + font_size=24, + position=SubtitlePosition.TOP, + ), + 'large': EmbeddingConfig( + font_name="Microsoft YaHei", + font_size=36, + outline_width=3, + ), + 'small': EmbeddingConfig( + font_name="Microsoft YaHei", + font_size=16, + outline_width=1, + ), + } + + return presets.get(preset_name.lower(), presets['default']) + + +def create_embedding_config_from_style( + style: SubtitleStyle, + video_quality: str = "high", + use_gpu: bool = False, +) -> EmbeddingConfig: + position_map = { + 1: SubtitlePosition.BOTTOM, + 2: SubtitlePosition.BOTTOM, + 3: SubtitlePosition.BOTTOM, + 5: SubtitlePosition.MIDDLE, + 8: SubtitlePosition.TOP, + } + + return EmbeddingConfig( + font_name=style.font_name, + font_size=style.font_size, + primary_color=style.primary_color, + outline_color=style.outline_color, + back_color=style.back_color, + outline_width=style.outline, + shadow=style.shadow, + position=position_map.get(style.alignment, SubtitlePosition.BOTTOM), + margin_v=style.margin_v, + bold=style.bold != 0, + italic=style.italic != 0, + video_quality=video_quality, + use_gpu=use_gpu, + ) diff --git a/subtitle_exporter.py b/subtitle_exporter.py new file mode 100644 index 0000000..186914b --- /dev/null +++ b/subtitle_exporter.py @@ -0,0 +1,778 @@ +""" +字幕导出模块 +支持SRT、ASS、VTT多格式字幕导出 +支持样式美化和双语字幕 +""" + +import logging +from typing import List, Dict, Any, Optional, Tuple +from pathlib import Path +from datetime import timedelta +import re + +from subtitle_models import ( + SubtitleSegment, + SubtitleDocument, + SubtitleStyle, + SubtitleFormat, +) + +logger = logging.getLogger(__name__) + + +SUPPORTED_EXPORT_FORMATS = { + 'srt': 'SubRip字幕格式', + 'ass': 'Advanced SubStation Alpha字幕格式', + 'ssa': 'SubStation Alpha字幕格式', + 'vtt': 'WebVTT字幕格式', +} + + +class SubtitleExporter: + """ + 字幕导出器类 + 支持多种字幕格式导出，包括样式美化 + """ + + def __init__(self, document: Optional[SubtitleDocument] = None): + self.document = document or SubtitleDocument() + + def load_document(self, document: SubtitleDocument): + self.document = document + + def load_from_segments(self, segments: List[Any]): + from subtitle_models import SubtitleSegment + + doc_segments = [] + for seg in segments: + if isinstance(seg, dict): + doc_segments.append(SubtitleSegment.from_dict(seg)) + else: + doc_segments.append(seg) + + self.document = SubtitleDocument(segments=doc_segments) + + def format_time_srt(self, seconds: float) -> str: + td = timedelta(seconds=seconds) + total_seconds = int(td.total_seconds()) + hours = total_seconds // 3600 + minutes = (total_seconds % 3600) // 60 + secs = total_seconds % 60 + milliseconds = int((td.total_seconds() - total_seconds) * 1000) + return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}" + + def format_time_vtt(self, seconds: float) -> str: + td = timedelta(seconds=seconds) + total_seconds = int(td.total_seconds()) + hours = total_seconds // 3600 + minutes = (total_seconds % 3600) // 60 + secs = total_seconds % 60 + milliseconds = int((td.total_seconds() - total_seconds) * 1000) + return f"{hours:02d}:{minutes:02d}:{secs:02d}.{milliseconds:03d}" + + def format_time_ass(self, seconds: float) -> str: + hours = int(seconds // 3600) + minutes = int((seconds % 3600) // 60) + secs = int(seconds % 60) + centiseconds = int((seconds - int(seconds)) * 100) + return f"{hours}:{minutes:02d}:{secs:02d}.{centiseconds:02d}" + + def _get_segment_display_text( + self, + segment: SubtitleSegment, + format_type: str, + bilingual_order: Optional[str] = None, + ) -> str: + if segment.is_bilingual: + order = bilingual_order or segment.bilingual_order or "zh_en" + + text_zh = segment.text_zh or segment.text + text_en = segment.text_en or segment.text + + if format_type == "ass": + if order == "zh_en": + return f"{text_zh}\\N{text_en}" + else: + return f"{text_en}\\N{text_zh}" + else: + if order == "zh_en": + return f"{text_zh}\n{text_en}" + else: + return f"{text_en}\n{text_zh}" + + return segment.text + + def generate_srt_content( + self, + segments: Optional[List[SubtitleSegment]] = None, + bilingual_order: Optional[str] = None, + ) -> str: + segs = segments or self.document.segments + + if not segs: + logger.warning("没有字幕段可生成") + return "" + + lines = [] + + for idx, seg in enumerate(segs, 1): + start_str = self.format_time_srt(seg.start) + end_str = self.format_time_srt(seg.end) + text = self._get_segment_display_text(seg, "srt", bilingual_order) + + lines.append(str(idx)) + lines.append(f"{start_str} --> {end_str}") + lines.append(text) + lines.append("") + + return '\n'.join(lines) + + def generate_vtt_content( + self, + segments: Optional[List[SubtitleSegment]] = None, + bilingual_order: Optional[str] = None, + add_style: bool = True, + ) -> str: + segs = segments or self.document.segments + + if not segs: + logger.warning("没有字幕段可生成") + return "" + + lines = [] + + lines.append("WEBVTT") + lines.append("") + + if add_style and self.document.styles: + lines.append("STYLE") + lines.append("::cue {") + lines.append(" background-color: transparent;") + lines.append(" color: white;") + lines.append(" font-family: 'Microsoft YaHei', sans-serif;") + lines.append(" font-size: 24px;") + lines.append(" text-shadow: 2px 2px 4px rgba(0,0,0,0.8);") + lines.append("}") + lines.append("") + + for idx, seg in enumerate(segs, 1): + start_str = self.format_time_vtt(seg.start) + end_str = self.format_time_vtt(seg.end) + text = self._get_segment_display_text(seg, "vtt", bilingual_order) + + lines.append(f"{idx:05d}") + lines.append(f"{start_str} --> {end_str}") + lines.append(text) + lines.append("") + + return '\n'.join(lines) + + def generate_ass_content( + self, + segments: Optional[List[SubtitleSegment]] = None, + title: str = "", + video_resolution: Tuple[int, int] = (1920, 1080), + bilingual_order: Optional[str] = None, + default_style: Optional[SubtitleStyle] = None, + ) -> str: + segs = segments or self.document.segments + + if not segs: + logger.warning("没有字幕段可生成") + return "" + + width, height = video_resolution + + lines = [] + + lines.append("[Script Info]") + lines.append(f"; Script generated by VideoSubtitleAI") + lines.append(f"Title: {title or self.document.title or 'Untitled'}") + lines.append("Original Script: VideoSubtitleAI") + lines.append("ScriptType: v4.00+") + lines.append("Collisions: Normal") + lines.append(f"PlayResX: {width}") + lines.append(f"PlayResY: {height}") + lines.append("Timer: 100.0000") + lines.append("") + + lines.append("[V4+ Styles]") + lines.append("Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding") + + styles = self.document.styles.copy() + + if default_style: + styles['Default'] = default_style + elif 'Default' not in styles: + styles['Default'] = SubtitleStyle.create_default_style() + + for style_name, style in styles.items(): + lines.append(style.to_ass_style_line()) + + lines.append("") + + lines.append("[Events]") + lines.append("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text") + + for seg in segs: + start_str = self.format_time_ass(seg.start) + end_str = self.format_time_ass(seg.end) + text = self._get_segment_display_text(seg, "ass", bilingual_order) + + style_name = seg.style_name or 'Default' + + margin_l = 0 + margin_r = 0 + margin_v = 0 + + lines.append(f"Dialogue: 0,{start_str},{end_str},{style_name},,{margin_l},{margin_r},{margin_v},,{text}") + + return '\n'.join(lines) + + def export_to_file( + self, + output_path: str, + format_type: Optional[str] = None, + encoding: str = 'utf-8', + **kwargs, + ) -> str: + output_path = Path(output_path).resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + + if format_type is None: + ext = output_path.suffix.lower() + if ext == '.srt': + format_type = 'srt' + elif ext == '.ass' or ext == '.ssa': + format_type = 'ass' + elif ext == '.vtt': + format_type = 'vtt' + else: + format_type = 'srt' + + format_type = format_type.lower() + + if format_type == 'srt': + content = self.generate_srt_content( + bilingual_order=kwargs.get('bilingual_order'), + ) + elif format_type == 'vtt': + content = self.generate_vtt_content( + bilingual_order=kwargs.get('bilingual_order'), + add_style=kwargs.get('add_style', True), + ) + elif format_type == 'ass': + content = self.generate_ass_content( + title=kwargs.get('title', ''), + video_resolution=kwargs.get('video_resolution', (1920, 1080)), + bilingual_order=kwargs.get('bilingual_order'), + default_style=kwargs.get('default_style'), + ) + else: + raise ValueError(f"不支持的字幕格式: {format_type}") + + try: + with open(output_path, 'w', encoding=encoding) as f: + f.write(content) + + logger.info(f"字幕文件已生成: {output_path}") + return str(output_path) + + except Exception as e: + logger.exception(f"写入字幕文件失败: {e}") + raise + + def export_multiple_formats( + self, + base_path: str, + formats: List[str] = ['srt', 'ass', 'vtt'], + encoding: str = 'utf-8', + **kwargs, + ) -> Dict[str, str]: + base_path = Path(base_path).resolve() + base_dir = base_path.parent + base_name = base_path.stem + + results = {} + + for fmt in formats: + fmt = fmt.lower() + ext = f".{fmt}" + output_path = base_dir / f"{base_name}{ext}" + + try: + result = self.export_to_file( + str(output_path), + format_type=fmt, + encoding=encoding, + **kwargs, + ) + results[fmt] = result + except Exception as e: + logger.error(f"导出 {fmt} 格式失败: {e}") + results[fmt] = None + + return results + + +class SubtitleImporter: + """ + 字幕导入器类 + 支持从SRT、ASS、VTT文件导入字幕 + """ + + def __init__(self): + pass + + def parse_srt(self, content: str) -> List[SubtitleSegment]: + segments = [] + + blocks = re.split(r'\n\n+', content.strip()) + + for block in blocks: + lines = block.strip().split('\n') + if len(lines) < 2: + continue + + seg_id = len(segments) + time_line_idx = 0 + + if lines[0].strip().isdigit(): + time_line_idx = 1 + + if time_line_idx >= len(lines): + continue + + time_line = lines[time_line_idx] + time_match = re.match( + r'(\d{2}:\d{2}:\d{2}[,.]\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}[,.]\d{3})', + time_line + ) + + if time_match: + start_time = self._parse_time_str(time_match.group(1)) + end_time = self._parse_time_str(time_match.group(2)) + + text_lines = lines[time_line_idx + 1:] + text = '\n'.join(text_lines).strip() + + is_bilingual = False + text_zh = "" + text_en = "" + + if '\n' in text: + parts = text.split('\n') + if len(parts) == 2: + is_bilingual = True + if self._is_chinese(parts[0]): + text_zh = parts[0] + text_en = parts[1] + else: + text_en = parts[0] + text_zh = parts[1] + + segment = SubtitleSegment( + id=seg_id, + start=start_time, + end=end_time, + text=text, + is_bilingual=is_bilingual, + text_zh=text_zh, + text_en=text_en, + ) + segments.append(segment) + + return segments + + def parse_vtt(self, content: str) -> List[SubtitleSegment]: + segments = [] + + lines = content.strip().split('\n') + + in_cue = False + current_id = 0 + start_time = 0.0 + end_time = 0.0 + text_lines = [] + + for line in lines: + line = line.strip() + + if not line: + if in_cue and text_lines: + text = '\n'.join(text_lines).strip() + + is_bilingual = False + text_zh = "" + text_en = "" + + if '\n' in text: + parts = text.split('\n') + if len(parts) == 2: + is_bilingual = True + if self._is_chinese(parts[0]): + text_zh = parts[0] + text_en = parts[1] + else: + text_en = parts[0] + text_zh = parts[1] + + segment = SubtitleSegment( + id=current_id, + start=start_time, + end=end_time, + text=text, + is_bilingual=is_bilingual, + text_zh=text_zh, + text_en=text_en, + ) + segments.append(segment) + current_id += 1 + + in_cue = False + text_lines = [] + continue + + if line.startswith("WEBVTT") or line.startswith("STYLE") or line.startswith("::cue"): + continue + + time_match = re.match( + r'(\d{2}:\d{2}:\d{2}[,.]\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}[,.]\d{3})', + line + ) + + if time_match: + start_time = self._parse_time_str(time_match.group(1)) + end_time = self._parse_time_str(time_match.group(2)) + in_cue = True + text_lines = [] + continue + + if in_cue: + if not line.isdigit(): + text_lines.append(line) + + if in_cue and text_lines: + text = '\n'.join(text_lines).strip() + segment = SubtitleSegment( + id=current_id, + start=start_time, + end=end_time, + text=text, + ) + segments.append(segment) + + return segments + + def parse_ass(self, content: str) -> List[SubtitleSegment]: + segments = [] + + styles = {} + in_events = False + + lines = content.split('\n') + + for line in lines: + line = line.strip() + + if line.startswith("[Events]"): + in_events = True + continue + + if line.startswith("[V4+ Styles]") or line.startswith("[V4 Styles]"): + in_events = False + continue + + if line.startswith("Style:"): + style_data = self._parse_ass_style(line) + if style_data: + styles[style_data['name']] = style_data + continue + + if not in_events: + continue + + if line.startswith("Dialogue:"): + seg_data = self._parse_ass_dialogue(line) + if seg_data: + segment = SubtitleSegment( + id=len(segments), + start=seg_data['start'], + end=seg_data['end'], + text=seg_data['text'], + style_name=seg_data.get('style'), + ) + segments.append(segment) + + return segments + + def _parse_ass_style(self, line: str) -> Optional[Dict[str, Any]]: + try: + if not line.startswith("Style:"): + return None + + style_part = line[6:].strip() + parts = style_part.split(',') + + if len(parts) < 22: + return None + + return { + 'name': parts[0], + 'font_name': parts[1], + 'font_size': int(parts[2]), + 'primary_color': parts[3], + 'secondary_color': parts[4], + 'outline_color': parts[5], + 'back_color': parts[6], + 'bold': int(parts[7]), + 'italic': int(parts[8]), + 'underline': int(parts[9]), + 'strike_out': int(parts[10]), + 'scale_x': int(parts[11]), + 'scale_y': int(parts[12]), + 'spacing': int(parts[13]), + 'angle': int(parts[14]), + 'border_style': int(parts[15]), + 'outline': int(parts[16]), + 'shadow': int(parts[17]), + 'alignment': int(parts[18]), + 'margin_l': int(parts[19]), + 'margin_r': int(parts[20]), + 'margin_v': int(parts[21]), + 'encoding': int(parts[22]) if len(parts) > 22 else 1, + } + except Exception: + return None + + def _parse_ass_dialogue(self, line: str) -> Optional[Dict[str, Any]]: + try: + if not line.startswith("Dialogue:"): + return None + + dialogue_part = line[9:].strip() + parts = dialogue_part.split(',', 9) + + if len(parts) < 10: + return None + + start_str = parts[1] + end_str = parts[2] + style = parts[3] + text = parts[9] + + text = text.replace('\\N', '\n').replace('\\n', ' ') + + import re + text = re.sub(r'\{[^}]+\}', '', text) + + return { + 'start': self._parse_ass_time(start_str), + 'end': self._parse_ass_time(end_str), + 'style': style, + 'text': text.strip(), + } + except Exception: + return None + + def _parse_ass_time(self, time_str: str) -> float: + time_str = time_str.strip() + + match = re.match(r'(\d+):(\d{2}):(\d{2})\.(\d{2})', time_str) + if match: + hours = int(match.group(1)) + minutes = int(match.group(2)) + seconds = int(match.group(3)) + centiseconds = int(match.group(4)) + return hours * 3600 + minutes * 60 + seconds + centiseconds / 100.0 + + return 0.0 + + def _parse_time_str(self, time_str: str) -> float: + time_str = time_str.strip() + time_str = time_str.replace(',', '.') + + parts = time_str.split(':') + if len(parts) == 3: + hours = int(parts[0]) + minutes = int(parts[1]) + seconds = float(parts[2]) + return hours * 3600 + minutes * 60 + seconds + elif len(parts) == 2: + minutes = int(parts[0]) + seconds = float(parts[1]) + return minutes * 60 + seconds + else: + return float(time_str) + + def _is_chinese(self, text: str) -> bool: + if not text: + return False + + chinese_count = 0 + for char in text: + if '\u4e00' <= char <= '\u9fff': + chinese_count += 1 + + if len(text) == 0: + return False + + return chinese_count / len(text) > 0.3 + + def import_from_file( + self, + file_path: str, + encoding: str = 'utf-8', + ) -> SubtitleDocument: + file_path = Path(file_path).resolve() + + if not file_path.exists(): + raise FileNotFoundError(f"字幕文件不存在: {file_path}") + + ext = file_path.suffix.lower() + + try: + with open(file_path, 'r', encoding=encoding) as f: + content = f.read() + except UnicodeDecodeError: + try: + with open(file_path, 'r', encoding='gbk') as f: + content = f.read() + except Exception: + raise ValueError(f"无法解码文件: {file_path}，尝试了 utf-8 和 gbk 编码") + + if ext == '.srt': + segments = self.parse_srt(content) + elif ext == '.vtt': + segments = self.parse_vtt(content) + elif ext in ['.ass', '.ssa']: + segments = self.parse_ass(content) + else: + segments = self.parse_srt(content) + + document = SubtitleDocument(segments=segments) + document.title = file_path.stem + + logger.info(f"成功导入字幕文件: {file_path}，共 {len(segments)} 段") + return document + + +def create_styled_ass_style( + font_name: str = "Microsoft YaHei", + font_size: int = 48, + primary_color: str = "#FFFFFF", + outline_color: str = "#000000", + outline_width: int = 2, + alignment: int = 2, + bold: bool = False, +) -> SubtitleStyle: + style = SubtitleStyle( + name="Custom", + font_name=font_name, + font_size=font_size, + border_style=1, + outline=outline_width, + alignment=alignment, + bold=1 if bold else 0, + ) + + style.update_color_from_hex(primary_color, "primary") + style.update_color_from_hex(outline_color, "outline") + + return style + + +def create_bilingual_style( + chinese_font: str = "Microsoft YaHei", + chinese_size: int = 48, + english_font: str = "Arial", + english_size: int = 36, +) -> Dict[str, SubtitleStyle]: + chinese_style = SubtitleStyle( + name="Chinese", + font_name=chinese_font, + font_size=chinese_size, + primary_color="&H00FFFFFF", + outline_color="&H00000000", + outline=2, + shadow=1, + alignment=2, + margin_v=60, + ) + + english_style = SubtitleStyle( + name="English", + font_name=english_font, + font_size=english_size, + primary_color="&H00FFFF00", + outline_color="&H00000000", + outline=2, + shadow=1, + alignment=2, + margin_v=20, + ) + + return { + 'Chinese': chinese_style, + 'English': english_style, + } + + +def hex_to_ass_color(hex_color: str, alpha: str = "00") -> str: + """ + 将十六进制RGB颜色转换为ASS格式颜色 + + ASS颜色格式: &HAABBGGRR + - AA: Alpha透明度 (00=完全透明, FF=完全不透明) + - BB: 蓝色分量 + - GG: 绿色分量 + - RR: 红色分量 + + Args: + hex_color: 十六进制颜色，如 "#FFFFFF" 或 "FFFFFF" + alpha: 透明度，默认为 "00" (完全不透明) + + Returns: + ASS格式的颜色字符串，如 "&H00FFFFFF" + """ + hex_color = hex_color.strip().lstrip('#') + + if len(hex_color) == 3: + hex_color = ''.join([c * 2 for c in hex_color]) + + if len(hex_color) != 6: + raise ValueError(f"无效的颜色格式: {hex_color}，应为6位十六进制") + + r = hex_color[0:2] + g = hex_color[2:4] + b = hex_color[4:6] + + return f"&H{alpha}{b}{g}{r}" + + +def ass_color_to_hex(ass_color: str) -> Tuple[str, str]: + """ + 将ASS格式颜色转换为十六进制RGB颜色 + + Args: + ass_color: ASS格式颜色，如 "&H00FFFFFF" 或 "00FFFFFF" + + Returns: + (hex_color, alpha) 元组，如 ("#FFFFFF", "00") + """ + ass_color = ass_color.strip().lstrip('&H').lstrip('&h') + + if len(ass_color) == 6: + alpha = "00" + bgr = ass_color + elif len(ass_color) == 8: + alpha = ass_color[0:2] + bgr = ass_color[2:8] + else: + return "#FFFFFF", "00" + + b = bgr[0:2] + g = bgr[2:4] + r = bgr[4:6] + + return f"#{r}{g}{b}", alpha diff --git a/subtitle_models.py b/subtitle_models.py new file mode 100644 index 0000000..b5ea68d --- /dev/null +++ b/subtitle_models.py @@ -0,0 +1,539 @@ +""" +字幕数据模型模块 +定义字幕段、样式、双语字幕等核心数据结构 +""" + +from dataclasses import dataclass, field +from typing import List, Dict, Any, Optional, Tuple +from datetime import timedelta +from enum import Enum +import re + + +class SubtitleFormat(Enum): + """字幕格式枚举""" + SRT = "srt" + ASS = "ass" + VTT = "vtt" + + +class Language(Enum): + """常用语言枚举""" + AUTO = "auto" + CHINESE = "zh" + ENGLISH = "en" + JAPANESE = "ja" + KOREAN = "ko" + FRENCH = "fr" + GERMAN = "de" + SPANISH = "es" + RUSSIAN = "ru" + PORTUGUESE = "pt" + ITALIAN = "it" + + +class TextStyle(Enum): + """文字样式枚举""" + NORMAL = "normal" + BOLD = "bold" + ITALIC = "italic" + UNDERLINE = "underline" + + +class HorizontalAlignment(Enum): + """水平对齐方式""" + LEFT = 1 + CENTER = 2 + RIGHT = 3 + + +class VerticalAlignment(Enum): + """垂直对齐方式""" + BOTTOM = 2 + MIDDLE = 5 + TOP = 8 + + +@dataclass +class SubtitleSegment: + """ + 字幕段数据模型 + 代表单个字幕条目，包含时间、文本、样式等信息 + """ + id: int = 0 + start: float = 0.0 + end: float = 0.0 + text: str = "" + language: str = "unknown" + + text_zh: str = "" + text_en: str = "" + + style_name: Optional[str] = None + is_bilingual: bool = False + bilingual_order: str = "zh_en" + + tokens: List[int] = field(default_factory=list) + avg_logprob: float = 0.0 + no_speech_prob: float = 0.0 + word_timestamps: List[Dict[str, Any]] = field(default_factory=list) + + metadata: Dict[str, Any] = field(default_factory=dict) + + @property + def duration(self) -> float: + return self.end - self.start + + @property + def display_text(self) -> str: + if not self.is_bilingual: + return self.text + + if self.bilingual_order == "zh_en": + zh = self.text_zh or self.text + en = self.text_en or "" + if en: + return f"{zh}\n{en}" + return zh + else: + en = self.text_en or self.text + zh = self.text_zh or "" + if zh: + return f"{en}\n{zh}" + return en + + def format_time_srt(self, seconds: float) -> str: + td = timedelta(seconds=seconds) + total_seconds = int(td.total_seconds()) + hours = total_seconds // 3600 + minutes = (total_seconds % 3600) // 60 + seconds = total_seconds % 60 + milliseconds = int((td.total_seconds() - total_seconds) * 1000) + return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}" + + def format_time_vtt(self, seconds: float) -> str: + td = timedelta(seconds=seconds) + total_seconds = int(td.total_seconds()) + hours = total_seconds // 3600 + minutes = (total_seconds % 3600) // 60 + seconds = total_seconds % 60 + milliseconds = int((td.total_seconds() - total_seconds) * 1000) + return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}" + + def to_dict(self) -> Dict[str, Any]: + return { + 'id': self.id, + 'start': self.start, + 'end': self.end, + 'text': self.text, + 'language': self.language, + 'text_zh': self.text_zh, + 'text_en': self.text_en, + 'style_name': self.style_name, + 'is_bilingual': self.is_bilingual, + 'bilingual_order': self.bilingual_order, + 'tokens': self.tokens, + 'avg_logprob': self.avg_logprob, + 'no_speech_prob': self.no_speech_prob, + 'word_timestamps': self.word_timestamps, + 'metadata': self.metadata, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> 'SubtitleSegment': + return cls( + id=data.get('id', 0), + start=data.get('start', 0.0), + end=data.get('end', 0.0), + text=data.get('text', ''), + language=data.get('language', 'unknown'), + text_zh=data.get('text_zh', ''), + text_en=data.get('text_en', ''), + style_name=data.get('style_name'), + is_bilingual=data.get('is_bilingual', False), + bilingual_order=data.get('bilingual_order', 'zh_en'), + tokens=data.get('tokens', []), + avg_logprob=data.get('avg_logprob', 0.0), + no_speech_prob=data.get('no_speech_prob', 0.0), + word_timestamps=data.get('word_timestamps', []), + metadata=data.get('metadata', {}), + ) + + @classmethod + def from_whisper_segment(cls, whisper_segment: Dict[str, Any]) -> 'SubtitleSegment': + return cls( + id=whisper_segment.get('id', 0), + start=round(float(whisper_segment.get('start', 0)), 3), + end=round(float(whisper_segment.get('end', 0)), 3), + text=whisper_segment.get('text', '').strip(), + language=whisper_segment.get('language', 'unknown'), + tokens=whisper_segment.get('tokens', []), + avg_logprob=whisper_segment.get('avg_logprob', 0.0), + no_speech_prob=whisper_segment.get('no_speech_prob', 0.0), + word_timestamps=whisper_segment.get('words', []), + ) + + +@dataclass +class SubtitleStyle: + """ + 字幕样式数据模型 + 用于ASS格式字幕的样式定义 + """ + name: str = "Default" + + font_name: str = "Microsoft YaHei" + font_size: int = 48 + primary_color: str = "&H00FFFFFF" + secondary_color: str = "&H000000FF" + outline_color: str = "&H00000000" + back_color: str = "&H00000000" + + bold: int = 0 + italic: int = 0 + underline: int = 0 + strike_out: int = 0 + scale_x: int = 100 + scale_y: int = 100 + spacing: int = 0 + angle: int = 0 + + border_style: int = 1 + outline: int = 2 + shadow: int = 1 + + alignment: int = 2 + margin_l: int = 10 + margin_r: int = 10 + margin_v: int = 40 + + encoding: int = 1 + + def to_ass_style_line(self) -> str: + return ( + f"Style: {self.name},{self.font_name},{self.font_size},{self.primary_color}," + f"{self.secondary_color},{self.outline_color},{self.back_color},{self.bold}," + f"{self.italic},{self.underline},{self.strike_out},{self.scale_x},{self.scale_y}," + f"{self.spacing},{self.angle},{self.border_style},{self.outline},{self.shadow}," + f"{self.alignment},{self.margin_l},{self.margin_r},{self.margin_v},{self.encoding}" + ) + + @classmethod + def create_default_style(cls) -> 'SubtitleStyle': + return cls() + + @classmethod + def create_chinese_style(cls) -> 'SubtitleStyle': + return cls( + name="Chinese", + font_name="Microsoft YaHei", + font_size=48, + primary_color="&H00FFFFFF", + outline_color="&H00000000", + outline=2, + shadow=1, + alignment=2, + margin_v=40, + ) + + @classmethod + def create_english_style(cls) -> 'SubtitleStyle': + return cls( + name="English", + font_name="Arial", + font_size=36, + primary_color="&H00FFFF00", + outline_color="&H00000000", + outline=2, + shadow=1, + alignment=2, + margin_v=80, + ) + + @classmethod + def create_top_style(cls) -> 'SubtitleStyle': + return cls( + name="Top", + font_name="Microsoft YaHei", + font_size=48, + alignment=8, + margin_v=40, + ) + + @classmethod + def create_bottom_style(cls) -> 'SubtitleStyle': + return cls( + name="Bottom", + font_name="Microsoft YaHei", + font_size=48, + alignment=2, + margin_v=40, + ) + + def update_color_from_hex(self, hex_color: str, color_type: str = "primary"): + hex_color = hex_color.lstrip('#') + if len(hex_color) == 6: + rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4)) + bgr_color = f"&H00{rgb[2]:02X}{rgb[1]:02X}{rgb[0]:02X}" + elif len(hex_color) == 8: + rgb = tuple(int(hex_color[i:i+2], 16) for i in (2, 4, 6)) + alpha = int(hex_color[0:2], 16) + bgr_color = f"&H{alpha:02X}{rgb[2]:02X}{rgb[1]:02X}{rgb[0]:02X}" + else: + raise ValueError(f"无效的颜色格式: {hex_color}") + + if color_type == "primary": + self.primary_color = bgr_color + elif color_type == "secondary": + self.secondary_color = bgr_color + elif color_type == "outline": + self.outline_color = bgr_color + elif color_type == "back": + self.back_color = bgr_color + + +@dataclass +class SubtitleDocument: + """ + 字幕文档数据模型 + 包含多个字幕段、样式集合和元数据 + """ + segments: List[SubtitleSegment] = field(default_factory=list) + styles: Dict[str, SubtitleStyle] = field(default_factory=dict) + + title: str = "" + original_language: str = "unknown" + translated_language: Optional[str] = None + is_bilingual: bool = False + + metadata: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if not self.styles: + self.styles['Default'] = SubtitleStyle.create_default_style() + + @property + def segment_count(self) -> int: + return len(self.segments) + + @property + def total_duration(self) -> float: + if not self.segments: + return 0.0 + return self.segments[-1].end - self.segments[0].start + + @property + def total_characters(self) -> int: + return sum(len(seg.text) for seg in self.segments) + + def add_segment(self, segment: SubtitleSegment) -> int: + segment.id = len(self.segments) + self.segments.append(segment) + return segment.id + + def add_style(self, style: SubtitleStyle): + self.styles[style.name] = style + + def get_segment(self, index: int) -> Optional[SubtitleSegment]: + if 0 <= index < len(self.segments): + return self.segments[index] + return None + + def remove_segment(self, index: int) -> bool: + if 0 <= index < len(self.segments): + self.segments.pop(index) + self._reindex_segments() + return True + return False + + def insert_segment(self, index: int, segment: SubtitleSegment) -> bool: + if 0 <= index <= len(self.segments): + self.segments.insert(index, segment) + self._reindex_segments() + return True + return False + + def _reindex_segments(self): + for idx, seg in enumerate(self.segments): + seg.id = idx + + def sort_segments_by_time(self): + self.segments.sort(key=lambda x: x.start) + self._reindex_segments() + + def merge_adjacent_segments(self, index: int, gap_threshold: float = 0.5) -> bool: + if index < 0 or index >= len(self.segments) - 1: + return False + + current = self.segments[index] + next_seg = self.segments[index + 1] + + if next_seg.start - current.end > gap_threshold: + return False + + merged = SubtitleSegment( + id=index, + start=current.start, + end=next_seg.end, + text=f"{current.text} {next_seg.text}", + language=current.language, + text_zh=f"{current.text_zh} {next_seg.text_zh}" if current.is_bilingual else "", + text_en=f"{current.text_en} {next_seg.text_en}" if current.is_bilingual else "", + is_bilingual=current.is_bilingual, + bilingual_order=current.bilingual_order, + style_name=current.style_name, + ) + + self.segments[index] = merged + self.segments.pop(index + 1) + self._reindex_segments() + return True + + def split_segment(self, index: int, split_seconds: float) -> Tuple[bool, int]: + if index < 0 or index >= len(self.segments): + return False, -1 + + segment = self.segments[index] + + if split_seconds <= segment.start or split_seconds >= segment.end: + return False, -1 + + duration = segment.end - segment.start + ratio = (split_seconds - segment.start) / duration + + text_1, text_2 = self._split_text_by_ratio(segment.text, ratio) + text_zh_1, text_zh_2 = self._split_text_by_ratio(segment.text_zh, ratio) + text_en_1, text_en_2 = self._split_text_by_ratio(segment.text_en, ratio) + + seg1 = SubtitleSegment( + id=index, + start=segment.start, + end=split_seconds, + text=text_1, + language=segment.language, + text_zh=text_zh_1, + text_en=text_en_1, + is_bilingual=segment.is_bilingual, + bilingual_order=segment.bilingual_order, + style_name=segment.style_name, + ) + + seg2 = SubtitleSegment( + id=index + 1, + start=split_seconds, + end=segment.end, + text=text_2, + language=segment.language, + text_zh=text_zh_2, + text_en=text_en_2, + is_bilingual=segment.is_bilingual, + bilingual_order=segment.bilingual_order, + style_name=segment.style_name, + ) + + self.segments[index] = seg1 + self.segments.insert(index + 1, seg2) + self._reindex_segments() + + return True, index + 1 + + def _split_text_by_ratio(self, text: str, ratio: float) -> Tuple[str, str]: + if not text: + return "", "" + + total_chars = len(text) + split_pos = int(total_chars * ratio) + + punctuations = ['。', '！', '？', '，', '；', '：', '.', '!', '?', ',', ';', ':', ' '] + + search_range = min(20, total_chars - split_pos) + for i in range(split_pos, min(split_pos + search_range, total_chars)): + if text[i] in punctuations: + split_pos = i + 1 + break + + search_range = min(20, split_pos) + for i in range(split_pos - 1, max(0, split_pos - search_range), -1): + if text[i] in punctuations: + split_pos = i + 1 + break + + return text[:split_pos].strip(), text[split_pos:].strip() + + def shift_all_timings(self, seconds: float): + for seg in self.segments: + seg.start = max(0.0, seg.start + seconds) + seg.end = max(0.0, seg.end + seconds) + + def to_dict(self) -> Dict[str, Any]: + return { + 'segments': [seg.to_dict() for seg in self.segments], + 'styles': {name: vars(style) for name, style in self.styles.items()}, + 'title': self.title, + 'original_language': self.original_language, + 'translated_language': self.translated_language, + 'is_bilingual': self.is_bilingual, + 'metadata': self.metadata, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> 'SubtitleDocument': + segments = [SubtitleSegment.from_dict(s) for s in data.get('segments', [])] + + styles = {} + for name, style_data in data.get('styles', {}).items(): + styles[name] = SubtitleStyle(**style_data) + + return cls( + segments=segments, + styles=styles, + title=data.get('title', ''), + original_language=data.get('original_language', 'unknown'), + translated_language=data.get('translated_language'), + is_bilingual=data.get('is_bilingual', False), + metadata=data.get('metadata', {}), + ) + + +def create_document_from_segments( + segments: List[Any], + is_bilingual: bool = False, +) -> SubtitleDocument: + """ + 从识别结果的段列表创建字幕文档 + + Args: + segments: 段列表（可以是字典或SubtitleSegment对象） + is_bilingual: 是否为双语字幕 + + Returns: + SubtitleDocument实例 + """ + doc_segments = [] + + for idx, seg in enumerate(segments): + if isinstance(seg, SubtitleSegment): + doc_segments.append(seg) + elif isinstance(seg, dict): + doc_seg = SubtitleSegment( + id=idx, + start=seg.get('start', 0.0), + end=seg.get('end', 0.0), + text=seg.get('text', '').strip(), + language=seg.get('language', 'unknown'), + tokens=seg.get('tokens', []), + avg_logprob=seg.get('avg_logprob', 0.0), + no_speech_prob=seg.get('no_speech_prob', 0.0), + is_bilingual=is_bilingual, + ) + doc_segments.append(doc_seg) + + document = SubtitleDocument(segments=doc_segments) + document.is_bilingual = is_bilingual + + return document + + +def create_default_style() -> SubtitleStyle: + """创建默认样式""" + return SubtitleStyle.create_default_style() diff --git a/subtitle_translator.py b/subtitle_translator.py new file mode 100644 index 0000000..20407e6 --- /dev/null +++ b/subtitle_translator.py @@ -0,0 +1,832 @@ +""" +字幕翻译模块 +支持双语字幕生成，提供多种翻译后端 +""" + +import logging +from typing import List, Dict, Any, Optional, Tuple, Callable +from dataclasses import dataclass +from enum import Enum +from abc import ABC, abstractmethod + +logger = logging.getLogger(__name__) + + +class TranslationBackend(Enum): + """翻译后端枚举""" + AUTO = "auto" + WHISPER_TRANSLATE = "whisper_translate" + ONLINE_API = "online_api" + LOCAL_MODEL = "local_model" + DUMMY = "dummy" + + +@dataclass +class TranslationResult: + """翻译结果数据类""" + original_text: str + translated_text: str + source_language: str + target_language: str + confidence: float = 1.0 + backend: str = "" + success: bool = True + error_message: Optional[str] = None + + +class BaseTranslator(ABC): + """翻译器基类""" + + backend: TranslationBackend = TranslationBackend.DUMMY + + def __init__(self, source_language: str = "auto", target_language: str = "en"): + self.source_language = source_language + self.target_language = target_language + self._initialized = False + + @abstractmethod + def translate(self, text: str) -> TranslationResult: + """翻译单个文本""" + pass + + def translate_batch(self, texts: List[str]) -> List[TranslationResult]: + """批量翻译""" + return [self.translate(text) for text in texts] + + @property + def is_initialized(self) -> bool: + return self._initialized + + +class WhisperTranslator(BaseTranslator): + """ + 使用Whisper的翻译功能 + Whisper本身支持将语音转录并翻译为英文 + """ + + def __init__( + self, + source_language: str = "auto", + target_language: str = "en", + whisper_model=None, + ): + super().__init__(source_language, target_language) + self.whisper_model = whisper_model + self._initialized = True + + def translate(self, text: str) -> TranslationResult: + """ + 注意：WhisperTranslator主要用于在语音识别阶段直接翻译 + 这个方法作为备用，仅返回原文本（因为Whisper不能翻译纯文本） + """ + logger.warning("WhisperTranslator不能翻译纯文本，建议在语音识别阶段使用task='translate'") + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.5, + backend="whisper_translate", + success=True, + ) + + +class OnlineAPITranslator(BaseTranslator): + """ + 在线API翻译器 + 支持多种翻译API（百度、谷歌、DeepL等） + """ + + SUPPORTED_APIS = ["baidu", "google", "deepl", "youdao"] + + def __init__( + self, + source_language: str = "auto", + target_language: str = "en", + api_type: str = "baidu", + api_key: Optional[str] = None, + app_id: Optional[str] = None, + ): + super().__init__(source_language, target_language) + self.api_type = api_type.lower() + self.api_key = api_key + self.app_id = app_id + self._initialized = self._check_credentials() + + def _check_credentials(self) -> bool: + if self.api_type == "baidu": + return self.api_key is not None and self.app_id is not None + elif self.api_type in ["google", "deepl", "youdao"]: + return self.api_key is not None + return False + + def translate(self, text: str) -> TranslationResult: + if not self._initialized: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend=f"online_{self.api_type}", + success=False, + error_message=f"API凭证未配置: {self.api_type}", + ) + + try: + if self.api_type == "baidu": + return self._translate_baidu(text) + elif self.api_type == "google": + return self._translate_google(text) + elif self.api_type == "deepl": + return self._translate_deepl(text) + elif self.api_type == "youdao": + return self._translate_youdao(text) + except Exception as e: + logger.error(f"在线翻译失败: {e}") + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend=f"online_{self.api_type}", + success=False, + error_message=str(e), + ) + + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend=f"online_{self.api_type}", + success=False, + error_message="不支持的API类型", + ) + + def _translate_baidu(self, text: str) -> TranslationResult: + try: + import requests + import hashlib + import uuid + import time + + url = "https://fanyi-api.baidu.com/api/trans/vip/translate" + + salt = str(uuid.uuid4()) + sign_str = f"{self.app_id}{text}{salt}{self.api_key}" + sign = hashlib.md5(sign_str.encode()).hexdigest() + + from_lang = self._map_language_baidu(self.source_language) + to_lang = self._map_language_baidu(self.target_language) + + params = { + 'q': text, + 'from': from_lang, + 'to': to_lang, + 'appid': self.app_id, + 'salt': salt, + 'sign': sign, + } + + response = requests.get(url, params=params, timeout=10) + result = response.json() + + if 'trans_result' in result: + translated_text = result['trans_result'][0]['dst'] + return TranslationResult( + original_text=text, + translated_text=translated_text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.9, + backend="online_baidu", + success=True, + ) + else: + error_msg = result.get('error_msg', '未知错误') + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_baidu", + success=False, + error_message=error_msg, + ) + + except ImportError: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_baidu", + success=False, + error_message="requests库未安装，请运行: pip install requests", + ) + + def _translate_google(self, text: str) -> TranslationResult: + try: + from googletrans import Translator + + translator = Translator() + src_lang = self._map_language_google(self.source_language) + dest_lang = self._map_language_google(self.target_language) + + result = translator.translate(text, src=src_lang, dest=dest_lang) + + return TranslationResult( + original_text=text, + translated_text=result.text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.85, + backend="online_google", + success=True, + ) + except ImportError: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_google", + success=False, + error_message="googletrans库未安装，请运行: pip install googletrans==4.0.0-rc1", + ) + except Exception as e: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_google", + success=False, + error_message=str(e), + ) + + def _translate_deepl(self, text: str) -> TranslationResult: + try: + import requests + + url = "https://api-free.deepl.com/v2/translate" + + source_lang = self._map_language_deepl(self.source_language) + target_lang = self._map_language_deepl(self.target_language) + + headers = { + 'Authorization': f'DeepL-Auth-Key {self.api_key}', + } + + data = { + 'text': text, + 'target_lang': target_lang.upper(), + } + + if source_lang and source_lang != 'auto': + data['source_lang'] = source_lang.upper() + + response = requests.post(url, headers=headers, data=data, timeout=10) + result = response.json() + + if 'translations' in result: + translated_text = result['translations'][0]['text'] + return TranslationResult( + original_text=text, + translated_text=translated_text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.95, + backend="online_deepl", + success=True, + ) + else: + error_msg = result.get('message', '未知错误') + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_deepl", + success=False, + error_message=error_msg, + ) + + except ImportError: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_deepl", + success=False, + error_message="requests库未安装", + ) + + def _translate_youdao(self, text: str) -> TranslationResult: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="online_youdao", + success=False, + error_message="有道翻译API暂未实现", + ) + + def _map_language_baidu(self, lang: str) -> str: + lang_map = { + 'auto': 'auto', + 'zh': 'zh', + 'en': 'en', + 'ja': 'jp', + 'ko': 'kor', + 'fr': 'fra', + 'de': 'de', + 'es': 'spa', + 'ru': 'ru', + 'pt': 'pt', + 'it': 'it', + } + return lang_map.get(lang, 'auto') + + def _map_language_google(self, lang: str) -> str: + lang_map = { + 'auto': 'auto', + 'zh': 'zh-cn', + 'en': 'en', + 'ja': 'ja', + 'ko': 'ko', + 'fr': 'fr', + 'de': 'de', + 'es': 'es', + 'ru': 'ru', + 'pt': 'pt', + 'it': 'it', + } + return lang_map.get(lang, 'auto') + + def _map_language_deepl(self, lang: str) -> str: + lang_map = { + 'auto': '', + 'zh': 'zh', + 'en': 'en', + 'ja': 'ja', + 'ko': 'ko', + 'fr': 'fr', + 'de': 'de', + 'es': 'es', + 'ru': 'ru', + 'pt': 'pt', + 'it': 'it', + } + return lang_map.get(lang, '') + + +class LocalModelTranslator(BaseTranslator): + """ + 本地模型翻译器 + 使用Helsinki-NLP的OPUS-MT模型或其他本地翻译模型 + """ + + def __init__( + self, + source_language: str = "zh", + target_language: str = "en", + model_name: Optional[str] = None, + device: str = "cpu", + ): + super().__init__(source_language, target_language) + self.model_name = model_name or self._get_default_model_name(source_language, target_language) + self.device = device + self.model = None + self.tokenizer = None + self._initialized = False + self._load_model() + + def _get_default_model_name(self, source: str, target: str) -> str: + model_map = { + ('zh', 'en'): "Helsinki-NLP/opus-mt-zh-en", + ('en', 'zh'): "Helsinki-NLP/opus-mt-en-zh", + ('ja', 'en'): "Helsinki-NLP/opus-mt-ja-en", + ('ko', 'en'): "Helsinki-NLP/opus-mt-ko-en", + } + return model_map.get((source, target), "Helsinki-NLP/opus-mt-zh-en") + + def _load_model(self): + try: + from transformers import MarianMTModel, MarianTokenizer + + logger.info(f"正在加载本地翻译模型: {self.model_name}") + + self.tokenizer = MarianTokenizer.from_pretrained(self.model_name) + self.model = MarianMTModel.from_pretrained(self.model_name) + self.model = self.model.to(self.device) + + self._initialized = True + logger.info(f"本地翻译模型加载成功: {self.model_name}") + + except ImportError as e: + logger.warning(f"无法加载本地翻译模型: {e}") + logger.warning("请安装: pip install transformers sentencepiece torch") + self._initialized = False + except Exception as e: + logger.error(f"加载本地翻译模型失败: {e}") + self._initialized = False + + def translate(self, text: str) -> TranslationResult: + if not self._initialized: + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="local_model", + success=False, + error_message="本地模型未初始化", + ) + + try: + import torch + + with torch.no_grad(): + inputs = self.tokenizer(text, return_tensors="pt", padding=True) + inputs = {k: v.to(self.device) for k, v in inputs.items()} + + outputs = self.model.generate(**inputs) + translated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True) + + return TranslationResult( + original_text=text, + translated_text=translated_text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.8, + backend="local_model", + success=True, + ) + + except Exception as e: + logger.error(f"本地翻译失败: {e}") + return TranslationResult( + original_text=text, + translated_text=text, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.0, + backend="local_model", + success=False, + error_message=str(e), + ) + + +class DummyTranslator(BaseTranslator): + """ + 模拟翻译器 + 用于测试或当没有可用翻译后端时，简单返回原文本 + """ + + def __init__( + self, + source_language: str = "auto", + target_language: str = "en", + mode: str = "same_text", + ): + super().__init__(source_language, target_language) + self.mode = mode + self._initialized = True + + def translate(self, text: str) -> TranslationResult: + if self.mode == "same_text": + translated = text + elif self.mode == "prefix": + translated = f"[{self.target_language}] {text}" + elif self.mode == "empty": + translated = "" + else: + translated = text + + return TranslationResult( + original_text=text, + translated_text=translated, + source_language=self.source_language, + target_language=self.target_language, + confidence=0.1, + backend="dummy", + success=True, + ) + + +class SubtitleTranslator: + """ + 字幕翻译管理器 + 统一管理多种翻译后端，提供字幕级别的翻译服务 + """ + + def __init__( + self, + source_language: str = "auto", + target_language: str = "en", + backend: TranslationBackend = TranslationBackend.AUTO, + whisper_model=None, + api_type: Optional[str] = None, + api_key: Optional[str] = None, + app_id: Optional[str] = None, + local_model_name: Optional[str] = None, + device: str = "cpu", + ): + self.source_language = source_language + self.target_language = target_language + self.backend = backend + + self.translators: Dict[str, BaseTranslator] = {} + self._default_translator: Optional[BaseTranslator] = None + + self._initialize_translators( + whisper_model=whisper_model, + api_type=api_type, + api_key=api_key, + app_id=app_id, + local_model_name=local_model_name, + device=device, + ) + + def _initialize_translators( + self, + whisper_model=None, + api_type: Optional[str] = None, + api_key: Optional[str] = None, + app_id: Optional[str] = None, + local_model_name: Optional[str] = None, + device: str = "cpu", + ): + if whisper_model is not None: + self.translators['whisper'] = WhisperTranslator( + source_language=self.source_language, + target_language=self.target_language, + whisper_model=whisper_model, + ) + + if api_type and api_key: + try: + self.translators['online'] = OnlineAPITranslator( + source_language=self.source_language, + target_language=self.target_language, + api_type=api_type, + api_key=api_key, + app_id=app_id, + ) + except Exception as e: + logger.warning(f"在线翻译API初始化失败: {e}") + + if local_model_name or self.backend == TranslationBackend.LOCAL_MODEL: + try: + self.translators['local'] = LocalModelTranslator( + source_language=self.source_language, + target_language=self.target_language, + model_name=local_model_name, + device=device, + ) + except Exception as e: + logger.warning(f"本地翻译模型初始化失败: {e}") + + self.translators['dummy'] = DummyTranslator( + source_language=self.source_language, + target_language=self.target_language, + ) + + self._select_default_translator() + + def _select_default_translator(self): + priority = ['local', 'online', 'whisper', 'dummy'] + + for key in priority: + if key in self.translators and self.translators[key].is_initialized: + self._default_translator = self.translators[key] + logger.info(f"选择翻译后端: {key}") + return + + self._default_translator = self.translators.get('dummy') + + def translate_text( + self, + text: str, + backend: Optional[str] = None, + ) -> TranslationResult: + translator = self._get_translator(backend) + return translator.translate(text) + + def translate_batch( + self, + texts: List[str], + backend: Optional[str] = None, + ) -> List[TranslationResult]: + translator = self._get_translator(backend) + return translator.translate_batch(texts) + + def _get_translator(self, backend: Optional[str] = None) -> BaseTranslator: + if backend and backend in self.translators: + return self.translators[backend] + return self._default_translator or self.translators['dummy'] + + def translate_segments( + self, + segments: List[Any], + source_language: str = "zh", + target_language: str = "en", + bilingual_order: str = "zh_en", + backend: Optional[str] = None, + ) -> List[Any]: + from subtitle_models import SubtitleSegment + + results = [] + + for seg in segments: + if isinstance(seg, dict): + segment = SubtitleSegment.from_dict(seg) + else: + segment = seg + + original_text = segment.text + + translation = self.translate_text(original_text, backend) + + if source_language == "zh" and target_language == "en": + segment.text_zh = original_text + segment.text_en = translation.translated_text + elif source_language == "en" and target_language == "zh": + segment.text_en = original_text + segment.text_zh = translation.translated_text + else: + if self.target_language == "en": + segment.text_en = translation.translated_text + segment.text_zh = original_text + else: + segment.text_zh = translation.translated_text + segment.text_en = original_text + + segment.is_bilingual = True + segment.bilingual_order = bilingual_order + + results.append(segment) + + return results + + def translate_document( + self, + document: Any, + source_language: Optional[str] = None, + target_language: Optional[str] = None, + bilingual_order: str = "zh_en", + backend: Optional[str] = None, + ) -> Any: + from subtitle_models import SubtitleDocument + + if source_language is None: + source_language = document.original_language or self.source_language + if target_language is None: + target_language = self.target_language + + translated_segments = self.translate_segments( + document.segments, + source_language=source_language, + target_language=target_language, + bilingual_order=bilingual_order, + backend=backend, + ) + + new_document = SubtitleDocument( + segments=translated_segments, + styles=document.styles, + title=document.title, + original_language=source_language, + translated_language=target_language, + is_bilingual=True, + metadata={ + **document.metadata, + 'translation_backend': self._default_translator.__class__.__name__ if self._default_translator else 'unknown', + }, + ) + + return new_document + + @property + def available_backends(self) -> List[str]: + return [key for key, translator in self.translators.items() if translator.is_initialized] + + def is_available(self) -> bool: + """检查是否有可用的翻译后端""" + return len(self.available_backends) > 0 + + @property + def preferred_backend(self) -> TranslationBackend: + """获取首选翻译后端""" + if self._default_translator: + return self._default_translator.backend + return TranslationBackend.DUMMY + + +_default_translator: Optional[SubtitleTranslator] = None + + +def get_translator( + source_language: str = "zh", + target_language: str = "en", +) -> SubtitleTranslator: + """ + 获取全局翻译器实例（单例模式） + + Args: + source_language: 源语言代码 + target_language: 目标语言代码 + + Returns: + SubtitleTranslator实例 + """ + global _default_translator + + if _default_translator is None: + _default_translator = SubtitleTranslator( + source_language=source_language, + target_language=target_language, + ) + + return _default_translator + + +def translate_segments_with_default( + segments: List[Any], + source_language: str = "zh", + target_language: str = "en", + bilingual_order: str = "zh_en", +) -> List[Any]: + """ + 使用默认翻译器翻译字幕段 + + Args: + segments: 字幕段列表 + source_language: 源语言 + target_language: 目标语言 + bilingual_order: 双语显示顺序 + + Returns: + 翻译后的字幕段列表 + """ + translator = get_translator(source_language, target_language) + return translator.translate_segments( + segments=segments, + source_language=source_language, + target_language=target_language, + bilingual_order=bilingual_order, + ) + + +def create_bilingual_document( + segments: List[Dict[str, Any]], + source_language: str = "zh", + target_language: str = "en", + bilingual_order: str = "zh_en", + translator: Optional[SubtitleTranslator] = None, +) -> Any: + """ + 创建双语文档 + + Args: + segments: 原始字幕段列表 + source_language: 源语言 + target_language: 目标语言 + bilingual_order: 双语显示顺序 + translator: 可选的翻译器实例 + + Returns: + 双语字幕文档 + """ + from subtitle_models import ( + SubtitleDocument, SubtitleSegment, + create_document_from_segments + ) + + if translator is None: + translator = get_translator(source_language, target_language) + + doc = create_document_from_segments(segments, is_bilingual=False) + + translated_doc = translator.translate_document( + document=doc, + source_language=source_language, + target_language=target_language, + bilingual_order=bilingual_order, + ) + + return translated_doc diff --git a/test_modules.py b/test_modules.py new file mode 100644 index 0000000..7460a51 --- /dev/null +++ b/test_modules.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- +""" +测试新模块导入和基本功能 +""" + +import sys +import os +sys.stdout.reconfigure(encoding='utf-8') +sys.stderr.reconfigure(encoding='utf-8') + +print("=" * 60) +print(" VideoSubtitleAI 新模块测试") +print("=" * 60) + +errors = [] + +# 测试 1: subtitle_models +print("\n[1/7] 测试 subtitle_models.py...") +try: + from subtitle_models import ( + SubtitleSegment, SubtitleStyle, SubtitleDocument, + create_document_from_segments, create_default_style + ) + print(" ✓ subtitle_models 导入成功") + + seg = SubtitleSegment(id=1, start=0.0, end=5.0, text="测试文本") + print(f" ✓ 创建 SubtitleSegment 成功: id={seg.id}, text={seg.text}") + + style = create_default_style() + print(f" ✓ 创建默认样式成功: font_name={style.font_name}, font_size={style.font_size}") + +except Exception as e: + print(f" ✗ subtitle_models 测试失败: {e}") + errors.append(f"subtitle_models: {e}") + +# 测试 2: subtitle_translator +print("\n[2/7] 测试 subtitle_translator.py...") +try: + from subtitle_translator import ( + SubtitleTranslator, TranslationBackend, get_translator, + translate_segments_with_default, create_bilingual_document + ) + print(" ✓ subtitle_translator 导入成功") + + translator = get_translator() + print(f" ✓ 获取翻译器成功: 后端={translator.preferred_backend.value}") + print(f" ✓ 翻译器可用: {'是' if translator.is_available() else '否 (仅测试模式)'}") + +except Exception as e: + print(f" ✗ subtitle_translator 测试失败: {e}") + errors.append(f"subtitle_translator: {e}") + +# 测试 3: subtitle_editor +print("\n[3/7] 测试 subtitle_editor.py...") +try: + from subtitle_editor import ( + SubtitleEditor, EditOperation, create_editor_from_segments + ) + print(" ✓ subtitle_editor 导入成功") + + test_segments = [ + SubtitleSegment(id=0, start=0.0, end=3.0, text="第一句"), + SubtitleSegment(id=1, start=3.0, end=6.0, text="第二句"), + ] + editor = create_editor_from_segments(test_segments) + print(f" ✓ 创建编辑器成功: segment_count={len(editor.document.segments)}") + + can_undo = editor.can_undo + can_redo = editor.can_redo + print(f" ✓ 撤销/重做功能可用: can_undo={can_undo}, can_redo={can_redo}") + +except Exception as e: + print(f" ✗ subtitle_editor 测试失败: {e}") + errors.append(f"subtitle_editor: {e}") + +# 测试 4: subtitle_exporter +print("\n[4/7] 测试 subtitle_exporter.py...") +try: + from subtitle_exporter import ( + SubtitleExporter, SubtitleImporter, + create_styled_ass_style, create_bilingual_style, + hex_to_ass_color, ass_color_to_hex, + SUPPORTED_EXPORT_FORMATS + ) + print(" ✓ subtitle_exporter 导入成功") + print(f" ✓ 支持的格式: {list(SUPPORTED_EXPORT_FORMATS.keys())}") + + ass_color = hex_to_ass_color("#FFFFFF") + print(f" ✓ 颜色转换成功: #FFFFFF -> {ass_color}") + + hex_color, alpha = ass_color_to_hex("&H00FFFFFF") + print(f" ✓ 逆颜色转换成功: &H00FFFFFF -> {hex_color} (alpha={alpha})") + + style = create_styled_ass_style( + font_name="Microsoft YaHei", + font_size=48, + primary_color="#FFD700", + outline_color="#000000", + outline_width=3, + ) + print(f" ✓ 创建自定义样式成功: font_size={style.font_size}") + +except Exception as e: + print(f" ✗ subtitle_exporter 测试失败: {e}") + errors.append(f"subtitle_exporter: {e}") + +# 测试 5: subtitle_embedder +print("\n[5/7] 测试 subtitle_embedder.py...") +try: + from subtitle_embedder import ( + SubtitleEmbedder, EmbeddingConfig, SubtitlePosition, + create_embedding_config_from_style + ) + print(" ✓ subtitle_embedder 导入成功") + + config = EmbeddingConfig( + font_name="Microsoft YaHei", + font_size=24, + position=SubtitlePosition.BOTTOM, + video_quality="high", + ) + print(f" ✓ 创建嵌入配置成功: position={config.position.value}") + + print(f" ✓ 字幕位置选项: {[p.value for p in SubtitlePosition]}") + print(f" ✓ 质量预设: {list(SubtitleEmbedder.PRESET_QUALITY.keys())}") + +except Exception as e: + print(f" ✗ subtitle_embedder 测试失败: {e}") + errors.append(f"subtitle_embedder: {e}") + +# 测试 6: language_detector +print("\n[6/7] 测试 language_detector.py...") +try: + from language_detector import ( + LanguageDetector, MultilingualSpeechRecognizer, + LANGUAGE_NAMES, is_chinese_text, is_english_text, classify_text_language + ) + print(" ✓ language_detector 导入成功") + + chinese_text = "这是一段中文文本" + english_text = "This is an English text" + + is_cn = is_chinese_text(chinese_text) + is_en = is_english_text(english_text) + print(f" ✓ 中文检测: '{chinese_text}' -> {'是中文' if is_cn else '不是中文'}") + print(f" ✓ 英文检测: '{english_text}' -> {'是英文' if is_en else '不是英文'}") + + lang, conf = classify_text_language(chinese_text) + print(f" ✓ 语言分类: '{chinese_text}' -> lang={lang}, confidence={conf:.2%}") + + detector = LanguageDetector() + result = detector.detect_language_from_text(chinese_text) + print(f" ✓ 增强检测: language={result.language}, confidence={result.confidence:.2%}") + + print(f" ✓ 支持语言: {list(LANGUAGE_NAMES.keys())}") + +except Exception as e: + print(f" ✗ language_detector 测试失败: {e}") + errors.append(f"language_detector: {e}") + +# 测试 7: video_subtitle_enhanced +print("\n[7/7] 测试 video_subtitle_enhanced.py...") +try: + from video_subtitle_enhanced import ( + VideoSubtitleAIEnhanced, EnhancedProcessingResult + ) + print(" ✓ video_subtitle_enhanced 导入成功") + + enhancer = VideoSubtitleAIEnhanced() + print(f" ✓ 创建增强版实例成功") + + preset_styles = enhancer.get_preset_styles() + print(f" ✓ 预设样式: {list(preset_styles.keys())}") + + supported_formats = enhancer.get_supported_formats() + print(f" ✓ 支持格式: {list(supported_formats.keys())}") + +except Exception as e: + print(f" ✗ video_subtitle_enhanced 测试失败: {e}") + errors.append(f"video_subtitle_enhanced: {e}") + +# 总结 +print("\n" + "=" * 60) +print(" 测试完成") +print("=" * 60) + +if errors: + print(f"\n✗ 检测到 {len(errors)} 个错误:") + for err in errors: + print(f" - {err}") + sys.exit(1) +else: + print("\n✓ 所有模块测试通过！") + print("\n新增功能模块已就绪:") + print(" 1. subtitle_models.py - 统一数据模型") + print(" 2. subtitle_translator.py - 双语字幕翻译") + print(" 3. subtitle_editor.py - 字幕编辑（撤销/重做）") + print(" 4. subtitle_exporter.py - 多格式导出（SRT/ASS/VTT）") + print(" 5. subtitle_embedder.py - 字幕硬嵌入（FFmpeg）") + print(" 6. language_detector.py - 语言检测优化") + print(" 7. video_subtitle_enhanced.py - 增强版主模块") + print(" 8. main.py (已更新) - 支持新命令行参数") + + print("\n使用示例:") + print(" python main.py -i video.mp4 --bilingual --formats srt,ass") + print(" python main.py -i video.mp4 --embed-video --style yellow") + print(" python main.py --list-styles") + print(" python main.py --list-formats") + print(" python main.py --export-only -s existing.srt --formats ass,vtt") + sys.exit(0) diff --git a/video_subtitle_enhanced.py b/video_subtitle_enhanced.py new file mode 100644 index 0000000..dc31c45 --- /dev/null +++ b/video_subtitle_enhanced.py @@ -0,0 +1,492 @@ +""" +VideoSubtitleAI 增强版主模块 +整合双语字幕、字幕编辑、多格式导出、硬嵌入等所有新功能 +""" + +import os +import sys +import logging +from pathlib import Path +from typing import List, Dict, Any, Optional, Tuple, Callable +from dataclasses import dataclass, field + +from subtitle_models import ( + SubtitleSegment, SubtitleStyle, SubtitleDocument, + create_document_from_segments, create_default_style +) +from subtitle_translator import ( + SubtitleTranslator, TranslationBackend, get_translator, + translate_segments_with_default, create_bilingual_document +) +from subtitle_editor import ( + SubtitleEditor, EditOperation, create_editor_from_segments +) +from subtitle_exporter import ( + SubtitleExporter, SubtitleImporter, + create_styled_ass_style, create_bilingual_style, + SUPPORTED_EXPORT_FORMATS +) +from subtitle_embedder import ( + SubtitleEmbedder, EmbeddingConfig, SubtitlePosition, + create_embedding_config_from_style +) +from language_detector import ( + LanguageDetector, MultilingualSpeechRecognizer, + LANGUAGE_NAMES, is_chinese_text, is_english_text, classify_text_language +) + + +logger = logging.getLogger(__name__) + + +@dataclass +class EnhancedProcessingResult: + """增强版处理结果""" + success: bool = False + input_path: str = '' + output_files: Dict[str, str] = field(default_factory=dict) + document: Optional[SubtitleDocument] = None + language_detection: Dict[str, Any] = field(default_factory=dict) + segment_count: int = 0 + statistics: Dict[str, Any] = field(default_factory=dict) + errors: List[str] = field(default_factory=list) + warnings: List[str] = field(default_factory=list) + + +class VideoSubtitleAIEnhanced: + """ + VideoSubtitleAI 增强版主类 + 整合所有新功能： + - 双语字幕生成 + - 字幕编辑（撤销/重做） + - 多格式导出（SRT/ASS/VTT） + - 字幕硬嵌入 + - 样式美化 + - 语言检测优化 + """ + + def __init__( + self, + base_recognizer=None, + whisper_model=None, + ffmpeg_path: Optional[str] = None, + verbose: bool = False, + ): + self.base_recognizer = base_recognizer + self.whisper_model = whisper_model + self.ffmpeg_path = ffmpeg_path + self.verbose = verbose + + self.language_detector = LanguageDetector(whisper_model=whisper_model) + self.multilingual_recognizer = MultilingualSpeechRecognizer( + base_recognizer=base_recognizer, + whisper_model=whisper_model + ) + + self._embedder: Optional[SubtitleEmbedder] = None + self._translator: Optional[SubtitleTranslator] = None + self._importer: Optional[SubtitleImporter] = None + + def get_embedder(self) -> SubtitleEmbedder: + if self._embedder is None: + try: + self._embedder = SubtitleEmbedder(ffmpeg_path=self.ffmpeg_path) + except Exception as e: + logger.warning(f"字幕嵌入器初始化失败: {e}") + raise + return self._embedder + + def get_translator(self) -> SubtitleTranslator: + if self._translator is None: + self._translator = get_translator() + return self._translator + + def get_importer(self) -> SubtitleImporter: + if self._importer is None: + self._importer = SubtitleImporter() + return self._importer + + def process_with_enhancement( + self, + recognition_result: Dict[str, Any], + source_language: str = 'auto', + enable_bilingual: bool = False, + target_language: str = 'en', + bilingual_order: str = 'zh_en', + translation_backend: str = 'auto', + ) -> Tuple[SubtitleDocument, Dict[str, Any]]: + """ + 处理识别结果，生成增强版字幕文档 + """ + segments = recognition_result.get('segments', []) + detected_lang = recognition_result.get('language', 'auto') + language_analysis = recognition_result.get('language_analysis', {}) + + if language_analysis.get('is_mixed', False): + logger.info(f"检测到多语言混合场景: {language_analysis.get('language_names', [])}") + + doc = create_document_from_segments(segments) + + for seg in doc.segments: + lang, conf = classify_text_language(seg.text) + seg.language = lang + + enhancement_info = { + 'detected_language': detected_lang, + 'language_analysis': language_analysis, + 'is_mixed': language_analysis.get('is_mixed', False), + 'languages': language_analysis.get('languages', []), + } + + if enable_bilingual: + doc = self._generate_bilingual_subtitles( + doc, + source_language=source_language if source_language != 'auto' else detected_lang, + target_language=target_language, + bilingual_order=bilingual_order, + translation_backend=translation_backend, + ) + enhancement_info['bilingual_enabled'] = True + enhancement_info['bilingual_order'] = bilingual_order + + return doc, enhancement_info + + def _generate_bilingual_subtitles( + self, + document: SubtitleDocument, + source_language: str = 'zh', + target_language: str = 'en', + bilingual_order: str = 'zh_en', + translation_backend: str = 'auto', + ) -> SubtitleDocument: + """ + 生成双语字幕 + """ + translator = self.get_translator() + + if not translator.is_available(): + logger.warning("翻译后端不可用，将仅保留原语言字幕") + return document + + segments_data = [] + for seg in document.segments: + segments_data.append({ + 'id': seg.id, + 'start': seg.start, + 'end': seg.end, + 'text': seg.text, + 'language': seg.language, + }) + + bilingual_doc = create_bilingual_document( + segments=segments_data, + source_language=source_language, + target_language=target_language, + bilingual_order=bilingual_order, + translator=translator, + ) + + return bilingual_doc + + def export_to_multiple_formats( + self, + document: SubtitleDocument, + base_output_path: str, + formats: List[str] = None, + default_style: Optional[SubtitleStyle] = None, + video_resolution: Tuple[int, int] = (1920, 1080), + ) -> Dict[str, str]: + """ + 导出为多种字幕格式 + """ + if formats is None: + formats = ['srt', 'ass'] + + base_path = Path(base_output_path) + output_files = {} + + exporter = SubtitleExporter(document) + + if default_style: + exporter.document.styles[default_style.name] = default_style + + for fmt in formats: + fmt = fmt.lower().strip() + if fmt not in SUPPORTED_EXPORT_FORMATS: + logger.warning(f"不支持的格式: {fmt}，跳过") + continue + + output_file = str(base_path.parent / f"{base_path.stem}.{fmt}") + + try: + if fmt == 'ass': + actual_path = exporter.export_to_file( + output_file, + format_type='ass', + video_resolution=video_resolution, + default_style=default_style, + ) + elif fmt == 'vtt': + actual_path = exporter.export_to_file( + output_file, + format_type='vtt', + add_style=True, + ) + else: + actual_path = exporter.export_to_file( + output_file, + format_type=fmt, + ) + + output_files[fmt] = actual_path + logger.info(f"已导出 {fmt.upper()} 格式: {actual_path}") + + except Exception as e: + logger.error(f"导出 {fmt} 格式失败: {e}") + + return output_files + + def embed_subtitles_to_video( + self, + video_path: str, + subtitle_path: str, + output_path: str, + style: Optional[SubtitleStyle] = None, + config: Optional[EmbeddingConfig] = None, + video_quality: str = 'high', + use_gpu: bool = False, + ) -> Dict[str, Any]: + """ + 将字幕硬嵌入到视频中 + """ + embedder = self.get_embedder() + + if style and not config: + config = create_embedding_config_from_style( + style=style, + video_quality=video_quality, + use_gpu=use_gpu, + ) + elif not config: + config = EmbeddingConfig( + video_quality=video_quality, + use_gpu=use_gpu, + ) + + result = embedder.embed_subtitles( + video_path=video_path, + subtitle_path=subtitle_path, + output_path=output_path, + config=config, + ) + + return result + + def analyze_language_from_segments( + self, + segments: List[Dict[str, Any]], + threshold: float = 0.3, + ) -> Dict[str, Any]: + """ + 分析字幕段的语言分布 + """ + return self.language_detector.detect_mixed_language_scenes( + segments=segments, + threshold=threshold, + ) + + def create_editor_from_document( + self, + document: SubtitleDocument, + ) -> SubtitleEditor: + """ + 从字幕文档创建编辑器 + """ + return SubtitleEditor(document=document) + + def create_editor_from_file( + self, + subtitle_path: str, + ) -> SubtitleEditor: + """ + 从字幕文件创建编辑器 + """ + importer = self.get_importer() + doc = importer.import_from_file(subtitle_path) + return SubtitleEditor(document=doc) + + def process_single_file_enhanced( + self, + recognition_result: Dict[str, Any], + input_path: str, + output_path: Optional[str] = None, + enable_bilingual: bool = False, + source_language: str = 'auto', + target_language: str = 'en', + bilingual_order: str = 'zh_en', + export_formats: List[str] = None, + enable_embedding: bool = False, + embedding_output_path: Optional[str] = None, + style: Optional[SubtitleStyle] = None, + video_quality: str = 'high', + use_gpu_for_embedding: bool = False, + ) -> EnhancedProcessingResult: + """ + 增强版单文件处理流程 + """ + result = EnhancedProcessingResult( + input_path=input_path, + ) + + try: + input_path_obj = Path(input_path) + + if output_path is None: + output_base = input_path_obj.parent / input_path_obj.stem + else: + output_path_obj = Path(output_path) + output_base = output_path_obj.parent / output_path_obj.stem + + doc, enhancement_info = self.process_with_enhancement( + recognition_result=recognition_result, + source_language=source_language, + enable_bilingual=enable_bilingual, + target_language=target_language, + bilingual_order=bilingual_order, + ) + + result.document = doc + result.language_detection = enhancement_info + result.segment_count = len(doc.segments) + + if export_formats is None: + export_formats = ['srt'] + + default_style = style or create_default_style() + + output_files = self.export_to_multiple_formats( + document=doc, + base_output_path=str(output_base), + formats=export_formats, + default_style=default_style, + ) + + result.output_files = output_files + + if enable_embedding: + if not self._is_video_file(input_path): + result.warnings.append("输入不是视频文件，跳过字幕嵌入") + else: + if 'ass' in output_files: + subtitle_for_embedding = output_files['ass'] + elif 'srt' in output_files: + subtitle_for_embedding = output_files['srt'] + else: + first_format = next(iter(output_files.keys()), None) + subtitle_for_embedding = output_files.get(first_format) if first_format else None + + if subtitle_for_embedding: + if embedding_output_path is None: + embedding_output = str(input_path_obj.parent / f"{input_path_obj.stem}_subtitled.mp4") + else: + embedding_output = embedding_output_path + + try: + embed_result = self.embed_subtitles_to_video( + video_path=input_path, + subtitle_path=subtitle_for_embedding, + output_path=embedding_output, + style=default_style, + video_quality=video_quality, + use_gpu=use_gpu_for_embedding, + ) + + result.output_files['embedded_video'] = embedding_output + logger.info(f"字幕硬嵌入完成: {embedding_output}") + + except Exception as e: + result.errors.append(f"字幕嵌入失败: {e}") + logger.error(f"字幕嵌入失败: {e}") + else: + result.errors.append("没有可用的字幕文件用于嵌入") + + result.statistics = { + 'segment_count': len(doc.segments), + 'is_bilingual': doc.is_bilingual, + 'languages': enhancement_info.get('languages', []), + 'exported_formats': list(output_files.keys()), + } + + result.success = True + + except Exception as e: + result.success = False + result.errors.append(str(e)) + logger.error(f"增强处理失败: {e}") + + return result + + def _is_video_file(self, path: str) -> bool: + video_extensions = {'.mp4', '.mkv', '.avi', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.ts', '.m2ts'} + return Path(path).suffix.lower() in video_extensions + + def get_supported_formats(self) -> Dict[str, str]: + return SUPPORTED_EXPORT_FORMATS.copy() + + def get_preset_styles(self) -> Dict[str, SubtitleStyle]: + return { + 'default': create_default_style(), + 'chinese_large': SubtitleStyle( + name='ChineseLarge', + font_name='Microsoft YaHei', + font_size=48, + primary_color='&H00FFFFFF', + outline_color='&H00000000', + outline=3, + shadow=2, + alignment=2, + margin_v=60, + ), + 'english': SubtitleStyle( + name='English', + font_name='Arial', + font_size=28, + primary_color='&H00FFFFFF', + outline_color='&H00000000', + outline=2, + shadow=1, + alignment=2, + margin_v=40, + ), + 'yellow': SubtitleStyle( + name='Yellow', + font_name='Microsoft YaHei', + font_size=48, + primary_color='&H0000FFFF', + outline_color='&H00000000', + outline=3, + shadow=2, + alignment=2, + ), + 'cyan': SubtitleStyle( + name='Cyan', + font_name='Microsoft YaHei', + font_size=48, + primary_color='&H00FFFF00', + outline_color='&H00000000', + outline=3, + shadow=2, + alignment=2, + ), + 'top': SubtitleStyle( + name='Top', + font_name='Microsoft YaHei', + font_size=40, + primary_color='&H00FFFFFF', + outline_color='&H00000000', + outline=2, + shadow=1, + alignment=8, + margin_v=40, + ), + }