capcut.py - master - MuchachMan/cupcat

capcut.py
313 строк15 KB
MuchachMan
upload files
23 фев 2026, 20:10
23 фев 2026, 20:10fd41972
import logging
import json
import subprocess
import sys
import os
from pathlib import Path
from typing import Optional, List, Dict
from faster_whisper import WhisperModel
from tqdm import tqdm  # <!-- ADD: Добавь эту строку -->
from vidcut import VideoCutter
from subtitle_generator import SubtitleGenerator

# --- Настройка логирования ---
def setup_logging(log_file="transcribe.log"):
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler(sys.stdout)
        ]
    )
    return logging.getLogger(__name__)

logger = setup_logging()

class VideoTranscriber:
    def __init__(self, model_size: str = "large-v3-turbo", device: str = "auto"):
        """
        Инициализация модели.
        Рекомендации для RU: 'large-v3-turbo' (баланс) или 'medium' (скорость).
        """
        logger.info(f"Инициализация модели Whisper ({model_size})...")
        try:
            # device="auto" сам выберет CUDA, если есть
            # compute_type="float16" для GPU, "int8" для CPU
            self.model = WhisperModel(model_size, device=device, compute_type="auto")
            logger.info("Модель успешно загружена.")
        except Exception as e:
            logger.error(f"Ошибка загрузки модели: {e}")
            raise e

    def extract_audio(self, video_path: Path, output_audio_path: Path) -> bool:
        """Извлечение аудио через FFmpeg (Mono, 16kHz)"""
        logger.info(f"Извлечение аудио из {video_path.name}...")
        
        command = [
            'ffmpeg', '-y',
            '-i', str(video_path),
            '-vn',
            '-acodec', 'pcm_s16le',
            '-ar', '16000',
            '-ac', '1',
            str(output_audio_path)
        ]

        try:
            subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            logger.info("Аудио успешно извлечено.")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"FFmpeg ошибка: {e.stderr.decode('utf-8')}")
            return False
        except FileNotFoundError:
            logger.error("FFmpeg не найден в системе.")
            return False

    def transcribe(self, audio_path: Path, language: str = "ru") -> List[Dict]:
        logger.info(f"Начало транскрибации (язык: {language})...")
        
        # Запускаем транскрибацию (возвращает генератор)
        segments_gen, info = self.model.transcribe(
            str(audio_path),
            language=language,
            beam_size=5,
            word_timestamps=True,
            vad_filter=True,
            vad_parameters=dict(min_silence_duration_ms=500)
        )

        logger.info(f"Детекция подтверждена: {info.language} ({info.language_probability:.2f})")
        
        words_data = []
        
        # --- ЛОГИКА ПРОГРЕСС-БАРА В РЕАЛЬНОМ ВРЕМЕНИ ---
        
        # Оценка: в среднем один сегмент речи длится около 3-4 секунд.
        # Делим общую длительность на 3.5, чтобы получить примерное кол-во сегментов.
        estimated_segments = max(1, int(info.duration / 3.5))
        
        # Оборачиваем генератор в tqdm. 
        # total=estimated_segments дает нам шкалу процентов.
        # disable=False включаем бар.
        # unit="сек" показывает, сколько секунд аудио обработано (примерно)
        progress_bar = tqdm(
            segments_gen, 
            total=estimated_segments, 
            desc="Обработка", 
            unit="сегм",
            ncols=80
        )

        for segment in progress_bar:
            # Обновляем описание прогресс-бара текущим временем
            progress_bar.set_postfix({"time": f"{segment.end:.1f}s"})
            
            if not hasattr(segment, 'words') or segment.words is None:
                continue
                
            for word in segment.words:
                word_text = word.word.strip()
                if not word_text:
                    continue
                    
                words_data.append({
                    "word": word_text,
                    "start": round(word.start, 3),
                    "end": round(word.end, 3),
                    "confidence": round(word.probability, 3)
                })
        
        progress_bar.close()
        # -----------------------------------------------

        logger.info(f"Транскрибация завершена. Слов: {len(words_data)}")
        return words_data

    def process(self, input_file: str, output_json: str, model: str, device: str):
        """Полный пайплайн"""
        input_path = Path(input_file)
        if not input_path.exists():
            logger.error(f"Файл не найден: {input_file}")
            return

        # Инициализация модели внутри процесса (чтобы можно было менять в runtime)
        self.model = WhisperModel(model, device=device, compute_type="auto")

        temp_audio = input_path.parent / "temp_audio.wav"
        output_path = Path(output_json)

        try:
            if not self.extract_audio(input_path, temp_audio):
                return

            words = self.transcribe(temp_audio, language="ru")

           # --- ЗАЩИТА ОТ ОШИБОК ТИПА ---
            # Если вдруг пришел set или кортеж, превращаем в список
            if isinstance(words, set):
                logger.warning("Получен набор (set), конвертируем в список...")
                words = list(words)
            elif not isinstance(words, list):
                logger.error(f"Неожиданный тип данных words: {type(words)}")
                return 
            # -----------------------------
            # --- COMPACT MODE START ---
            # Оставляем только нужное для LLM: слово, старт, конец
            compact_words = [
                {"w": w["word"], "s": w["start"], "e": w["end"]} 
                for w in words
            ]

            result_data = {
                "d": words[-1]['end'] if words else 0,
                "t": compact_words
            }

            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(result_data, f, ensure_ascii=False, indent=None, separators=(',', ':'))
            # --- COMPACT MODE END ---
            
            logger.info(f"Результат сохранен в {output_json} (Compact Mode)")

        except Exception as e:
            logger.exception(f"Критическая ошибка: {e}")
        
        finally:
            if temp_audio.exists():
                try:
                    temp_audio.unlink()
                except:
                    pass

if __name__ == "__main__":
    import argparse
    # --- АРГУМЕНТЫ ДЛЯ ТРАНСКРИБАЦИИ ---
    parser = argparse.ArgumentParser(description="Транскрибация видео/аудио (RU)")
    parser.add_argument("input", nargs="?", help="Путь к видео/аудио файлу")
    #parser.add_argument("-o", "--output", default="transcription.json", help="Путь к выходному JSON")
    parser.add_argument("-m", "--model", default="large-v3-turbo", 
                        choices=["tiny", "base", "small", "medium", "large-v3", "large-v3-turbo"],
                        help="Модель Whisper (default: large-v3-turbo)")
    parser.add_argument("-d", "--device", default="auto", choices=["auto", "cpu", "cuda"],
                        help="Устройство для вычислений")
    # ----------------------------------                    
    # --- АРГУМЕНТЫ ДЛЯ НАРЕЗКИ ---
    parser.add_argument("--cut", action="store_true", help="Режим нарезки видео")
    # Принимаем строку вида "1.5-25.0; 85.0-120.5"
    parser.add_argument("--time", type=str, help="Таймкоды в формате 'start-end; start-end'")
    parser.add_argument("--concat", action="store_true", help="Склеить фрагменты в один файл")
    parser.add_argument("--output", "-o", default="output.mp4", help="Выходной файл видео")
    parser.add_argument("--transcript", "-t", type=str, default="transcription.json", 
                        help="Файл транскрибации для создания JSON сегментов")
    # ----------------------------------
    # --- АРГУМЕНТЫ ДЛЯ СУБТИТРОВ ---
    # parser.add_argument("--subs", action="store_true", help="Режим наложения субтитров")
    # parser.add_argument("--video", type=str, help="Видео файл для субтитров")
    # parser.add_argument("--json", type=str, help="JSON транскрибации для субтитров")
    # parser.add_argument("--font", type=str, default="Arial", help="Шрифт субтитров")
    # parser.add_argument("--color", type=str, default="&H00FFFF&", help="Цвет текста (ASS формат)")
    # ----------------------------------
    # --- АРГУМЕНТЫ ДЛЯ СУБТИТРОВ ---
    parser.add_argument("--subs", action="store_true", help="Режим наложения субтитров")
    parser.add_argument("--video", type=str, help="Видео файл для субтитров")
    parser.add_argument("--json", type=str, help="JSON транскрибации для субтитров")
    parser.add_argument("--font", type=str, default="Arial", help="Шрифт субтитров")
    parser.add_argument("--color", type=str, default="&H00FFFF&", help="Цвет текста (ASS формат)")
    parser.add_argument("--pos", type=str, default="bottom", help="Позиция: bottom, center, top")
    parser.add_argument("--x", type=int, default=None, help="Точная координата X (пиксели)")
    parser.add_argument("--y", type=int, default=None, help="Точная координата Y (пиксели)")
    parser.add_argument("--size", type=int, default=48, help="Размер шрифта")
    parser.add_argument("--margin", type=int, default=10, help="Вертикальный отступ")
    # ----------------------------------
    args = parser.parse_args()
    args = parser.parse_args()


    # --- РЕЖИМ СУБТИТРОВ ---
    if args.subs:
        if not args.video or not args.json:
            logger.error("Укажите --video и --json для режима субтитров")
            sys.exit(1)
        
        subgen = SubtitleGenerator(
            font=args.font, 
            primary_color=args.color,
            position=args.pos,
            pos_x=args.x,  # Новые параметры
            pos_y=args.y,
            font_size=args.size,
            margin_v=args.margin
        )
        success = subgen.process(
            video_path=args.video,
            transcription_json=args.json,
            output_path=args.output
        )
        sys.exit(0 if success else 1)
        
    # --- РЕЖИМ НАРЕЗКИ ---
    if args.cut:
        cutter = VideoCutter()
        if not args.input:
            logger.error("Укажите входной видеофайл (input)")
            sys.exit(1)        
        
        
        if not args.time:
            logger.error("Для нарезки укажите --time (например: --time '1.5-25.0; 85.0-120.5')")
            sys.exit(1)
        
        # Парсим строку таймкодов в список кортежей [(1.5, 25.0), (85.0, 120.5)]
        timestamps = []
        try:
            parts = args.time.split(";")
            for part in parts:
                if "-" in part:
                    s, e = part.strip().split("-")
                    timestamps.append((float(s), float(e)))
        except Exception as e:
            logger.error(f"Ошибка формата таймкодов. Используйте 'start-end; start-end'. {e}")
            sys.exit(1)

        if not timestamps:
            logger.error("Не удалось распознать ни одного таймкода")
            sys.exit(1)
         # -----------------Парсим----------------------------------#           
        success = cutter.cut_by_timestamps(
            video_path=args.input,
            timestamps=timestamps,
            output_path=args.output,
            concatenate=args.concat,  # Передаем флаг склейки
            transcription_json=args.transcript  # Передаём путь к JSON            
        )       
        sys.exit(0 if success else 1)
    # --- РЕЖИМ ТРАНСКРИБАЦИИ (по умолчанию) ---
    else:
        if not args.input:
            logger.error("Укажите входной видеофайл (input)")
            sys.exit(1)

        app = VideoTranscriber()        
        # Исправление: для транскрибации всегда используем .json
        json_output = "transcription.json"       
        # Если пользователь явно указал файл с расширением .json, используем его
        if args.output.endswith('.json'):
            json_output = args.output
            
        app.process(args.input, json_output, args.model, args.device)



    # else:
        # app = VideoTranscriber()
                
        # # Если вывод не указан, по умолчанию сохраняем в transcription.json
        # json_output = args.output if args.output.endswith('.json') else "transcription.json"
        
        # # Если пользователь указал .mp4 в режиме транскрибации, игнорируем это для JSON
        # if not args.output.endswith('.json'):
            # json_output = "transcription.json"
            
        # app.process(args.input, args.output, args.model, args.device)