Add video input support with ffmpeg audio extraction

This commit is contained in:
Кирилл Блинов 2026-05-29 10:53:16 +03:00
parent 09d0a74520
commit ddee721bea
3 changed files with 169 additions and 14 deletions

64
run.py
View File

@ -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("При первом запуске будут скачаны модели ИИ (~45 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,

109
src/audio_utils.py Normal file
View File

@ -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

View File

@ -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)