From ddee721beaad0056ef631e91db0523d9b73bb00a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=91=D0=BB=D0=B8?= =?UTF-8?q?=D0=BD=D0=BE=D0=B2?= Date: Fri, 29 May 2026 10:53:16 +0300 Subject: [PATCH] Add video input support with ffmpeg audio extraction --- run.py | 64 +++++++++++++++++++++----- src/audio_utils.py | 109 +++++++++++++++++++++++++++++++++++++++++++++ src/pipeline.py | 10 ++++- 3 files changed, 169 insertions(+), 14 deletions(-) create mode 100644 src/audio_utils.py diff --git a/run.py b/run.py index 375a261..cb981a2 100644 --- a/run.py +++ b/run.py @@ -5,6 +5,7 @@ import os import sys from pathlib import Path +from src.audio_utils import check_ffmpeg, is_audio_file, is_video_file from src.config import get_profile, load_config from src.document import build_document from src.pipeline import run_pipeline @@ -20,11 +21,28 @@ def resolve_device(preferred: str) -> str: return "cpu" +def print_first_run_info(): + """Показывает информацию о первом запуске и скачивании моделей.""" + print("=" * 60) + print("ПЕРВЫЙ ЗАПУСК: скачивание моделей") + print("=" * 60) + print("При первом запуске будут скачаны модели ИИ (~4–5 GB):") + print(" • Whisper large-v3 ~3.0 GB (распознавание речи)") + print(" • Pyannote диаризация ~0.4 GB (разделение спикеров)") + print(" • Wav2Vec2 alignment ~1.0 GB (точные таймкоды)") + print("") + print("Это займёт время в зависимости от скорости интернета.") + print("При последующих запусках модели загружаются с диска.") + print("=" * 60) + print("") + + def main(): parser = argparse.ArgumentParser( - description="Транскрибация совещаний с диаризацией и таймкодами." + description="Транскрибация совещаний с диаризацией и таймкодами. " + "Поддерживает аудио (wav, mp3, m4a, ogg, flac) и видео (mp4, avi, mkv, mov, etc.)." ) - parser.add_argument("--input", "-i", required=True, help="Путь к аудиофайлу") + parser.add_argument("--input", "-i", required=True, help="Путь к аудио или видео файлу") parser.add_argument("--output", "-o", default=None, help="Путь к выходному файлу (docx/md/txt)") parser.add_argument("--profile", "-p", default=None, help="Профиль конфигурации (mac_m4, gpu_8gb, cpu_best)") parser.add_argument("--config", "-c", default=None, help="Путь к config.yaml") @@ -35,10 +53,28 @@ def main(): args = parser.parse_args() - if not Path(args.input).exists(): + input_path = Path(args.input) + if not input_path.exists(): print(f"Ошибка: файл не найден: {args.input}", file=sys.stderr) sys.exit(1) + # Проверка ffmpeg для видео + if is_video_file(args.input): + if not check_ffmpeg(): + print( + "Ошибка: для обработки видео нужен ffmpeg.\n" + "Установите ffmpeg:\n" + " Mac: brew install ffmpeg\n" + " Linux: sudo apt-get install ffmpeg\n" + " Windows: https://ffmpeg.org/download.html", + file=sys.stderr, + ) + sys.exit(1) + print(f"[Info] Обнаружен видео файл: {input_path.suffix}") + print("[Info] Аудио будет извлечено автоматически.") + elif not is_audio_file(args.input): + print(f"[Warning] Неизвестный формат: {input_path.suffix}. Попытка обработать...") + # Загрузка конфига config = load_config(args.config) profile = get_profile(config, args.profile) @@ -59,7 +95,7 @@ def main(): if args.output: output_path = args.output else: - stem = Path(args.input).stem + stem = input_path.stem output_dir = Path(config.get("paths", {}).get("output_dir", "./output")) output_path = str(output_dir / f"{stem}.{fmt}") @@ -68,22 +104,26 @@ def main(): if profile.get("diarize", True) and not hf_token: print( "Ошибка: для диаризации нужен HuggingFace токен.\n" - "Установите env HF_TOKEN или укажите hf_token в config.yaml", + "Установите env HF_TOKEN или укажите hf_token в config.yaml\n" + "Инструкция: https://huggingface.co/docs/hub/security-tokens", file=sys.stderr, ) sys.exit(1) - print(f"Профиль: {args.profile or config.get('active_profile')}") - print(f"Устройство: {profile['device']}") - print(f"Модель: {profile['model']}") - print(f"Язык: {profile['language']}") - print(f"Вход: {args.input}") - print(f"Выход: {output_path}") + # Информация о первом запуске + print_first_run_info() + + print(f"Профиль: {args.profile or config.get('active_profile')}") + print(f"Устройство: {profile['device']}") + print(f"Модель: {profile['model']}") + print(f"Язык: {profile['language']}") + print(f"Вход: {args.input}") + print(f"Выход: {output_path}") print("-" * 40) # Запуск пайплайна result = run_pipeline( - audio_path=args.input, + input_path=args.input, profile_name=args.profile, config_path=args.config, output_path=output_path, diff --git a/src/audio_utils.py b/src/audio_utils.py new file mode 100644 index 0000000..770f978 --- /dev/null +++ b/src/audio_utils.py @@ -0,0 +1,109 @@ +"""Утилиты для обработки входных файлов (видео → аудио, ffmpeg).""" + +import shutil +import subprocess +import tempfile +from pathlib import Path +from typing import Optional + + +VIDEO_EXTENSIONS = {".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm", ".m4v", ".mpeg", ".mpg"} +AUDIO_EXTENSIONS = {".wav", ".mp3", ".m4a", ".ogg", ".flac", ".aac", ".wma"} + + +def check_ffmpeg() -> bool: + """Проверяет, установлен ли ffmpeg.""" + return shutil.which("ffmpeg") is not None + + +def is_video_file(file_path: str) -> bool: + """Проверяет, является ли файл видео.""" + return Path(file_path).suffix.lower() in VIDEO_EXTENSIONS + + +def is_audio_file(file_path: str) -> bool: + """Проверяет, является ли файл аудио.""" + return Path(file_path).suffix.lower() in AUDIO_EXTENSIONS + + +def extract_audio_from_video( + video_path: str, + output_audio_path: Optional[str] = None, + sample_rate: int = 16000, + channels: int = 1, +) -> str: + """ + Извлекает аудио из видео файла через ffmpeg. + + Args: + video_path: Путь к видео + output_audio_path: Куда сохранить аудио (если None — создаётся временный файл) + sample_rate: Частота дискретизации (16kHz оптимально для Whisper) + channels: Количество каналов (1 = mono) + + Returns: + Путь к извлечённому аудио файлу (wav) + """ + if not check_ffmpeg(): + raise RuntimeError( + "ffmpeg не найден. Установите ffmpeg:\n" + " Mac: brew install ffmpeg\n" + " Ubuntu/Debian: sudo apt-get install ffmpeg\n" + " Windows: скачайте с https://ffmpeg.org/download.html" + ) + + video = Path(video_path) + if not video.exists(): + raise FileNotFoundError(f"Видео файл не найден: {video_path}") + + if output_audio_path is None: + # Временный файл в папке tmp проекта + tmp_dir = Path("tmp") + tmp_dir.mkdir(exist_ok=True) + output_audio_path = str(tmp_dir / f"{video.stem}_audio.wav") + + cmd = [ + "ffmpeg", + "-y", # перезаписать если существует + "-i", str(video), + "-vn", # без видео + "-acodec", "pcm_s16le", # 16-bit PCM + "-ar", str(sample_rate), # частота дискретизации + "-ac", str(channels), # моно + output_audio_path, + ] + + print(f"[Extract] Извлечение аудио из {video.name}...") + result = subprocess.run( + cmd, + capture_output=True, + text=True, + ) + + if result.returncode != 0: + raise RuntimeError(f"Ошибка ffmpeg:\n{result.stderr}") + + print(f"[Extract] Аудио сохранено: {output_audio_path}") + return output_audio_path + + +def prepare_audio_input(input_path: str) -> str: + """ + Подготавливает входной файл: если видео — извлекает аудио, + если аудио — возвращает как есть. + + Returns: + Путь к аудио файлу готовому для обработки + """ + input_path = str(input_path) + + if is_audio_file(input_path): + return input_path + + if is_video_file(input_path): + return extract_audio_from_video(input_path) + + # Неизвестное расширение — попробуем передать как есть, + # ffmpeg/whisperx сам разберётся + print(f"[Warning] Неизвестный формат файла: {Path(input_path).suffix}. Попытка обработать как аудио...") + return input_path diff --git a/src/pipeline.py b/src/pipeline.py index 01e9a5d..fa52b96 100644 --- a/src/pipeline.py +++ b/src/pipeline.py @@ -7,17 +7,19 @@ from typing import Any, Dict, List, Optional import whisperx from whisperx.diarize import DiarizationPipeline +from src.audio_utils import prepare_audio_input from src.config import get_profile, load_config, resolve_hf_token def run_pipeline( - audio_path: str, + input_path: str, profile_name: Optional[str] = None, config_path: Optional[str] = None, output_path: Optional[str] = None, ) -> Dict[str, Any]: """ Запускает полный пайплайн транскрибации. + Принимает видео или аудио файлы. Returns: Словарь с сегментами (speaker, text, start, end). @@ -34,7 +36,11 @@ def run_pipeline( hf_token = resolve_hf_token(config) - # 1. Загрузка аудио + # 1. Подготовка входного файла (видео → аудио) + audio_path = prepare_audio_input(input_path) + + # 2. Загрузка аудио + print("[Pipeline] Загрузка аудио...") audio = whisperx.load_audio(audio_path) # 2. Транскрибация (ASR)