four-level-schedule/duomotai/utils/ffmpeg_wrapper.py

import subprocess
import os
from typing import Optional


class FFmpegWrapper:
    """FFmpeg命令行包装工具"""
    
    def __init__(self):
        self.ffmpeg_path = "ffmpeg"  # 假设ffmpeg已在系统PATH中
    
    def extract_audio(self, video_path: str, output_audio_path: str) -> bool:
        """
        从视频中提取音频轨道
        
        Args:
            video_path: 视频文件路径
            output_audio_path: 输出音频文件路径
            
        Returns:
            bool: 操作是否成功
        """
        try:
            cmd = [
                self.ffmpeg_path,
                "-i", video_path,
                "-vn",  # 禁用视频
                "-acodec", "pcm_s16le",  # 16位PCM
                "-ar", "16000",  # 16kHz采样率
                "-ac", "1",  # 单声道
                "-y",  # 覆盖输出文件
                output_audio_path
            ]
            
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            return True
        except subprocess.CalledProcessError as e:
            raise Exception(f"音频提取失败: {e.stderr}")
    
    def convert_audio(self, input_audio_path: str, output_audio_path: str) -> bool:
        """
        转换音频格式为16k/16bit/mono wav
        
        Args:
            input_audio_path: 输入音频文件路径
            output_audio_path: 输出音频文件路径
            
        Returns:
            bool: 操作是否成功
        """
        try:
            cmd = [
                self.ffmpeg_path,
                "-i", input_audio_path,
                "-acodec", "pcm_s16le",
                "-ar", "16000",
                "-ac", "1",
                "-y",
                output_audio_path
            ]
            
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            return True
        except subprocess.CalledProcessError as e:
            raise Exception(f"音频转换失败: {e.stderr}")
    
    def extract_keyframes(self, video_path: str, output_dir: str, interval: int = 60) -> list:
        """
        从视频中提取关键帧
        
        Args:
            video_path: 视频文件路径
            output_dir: 输出目录
            interval: 提取间隔（秒）
            
        Returns:
            list: 提取的关键帧文件路径列表
        """
        try:
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            # 提取关键帧
            output_pattern = os.path.join(output_dir, "frame_%06d.jpg")
            cmd = [
                self.ffmpeg_path,
                "-i", video_path,
                "-vf", f"fps=1/{interval}",  # 每分钟一张
                "-y",
                output_pattern
            ]
            
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            
            # 收集提取的关键帧
            keyframes = []
            for file in os.listdir(output_dir):
                if file.startswith("frame_") and file.endswith(".jpg"):
                    keyframes.append(os.path.join(output_dir, file))
            
            return sorted(keyframes)
        except subprocess.CalledProcessError as e:
            raise Exception(f"关键帧提取失败: {e.stderr}")