ass_srt_extractor.py

import os
import subprocess
import re
import time
from tqdm import tqdm

def list_subtitle_tracks(filename):
    cmd = ['mkvmerge', '-i', filename]
    result = subprocess.run(cmd, capture_output=True, text=True)
    tracks = []

    for line in result.stdout.splitlines():
        if "subtitles" in line:
            match = re.match(r'Track ID (\d+): subtitles \(([^)]+)\)(?: \[(.+?)\])?', line)
            if match:
                track_id = match.group(1)
                codec = match.group(2)
                lang = match.group(3) if match.group(3) else "eng"  # padrão eng se faltar
                tracks.append((track_id, codec, lang))
    return tracks

def extract_track(input_file, track_id, output_file):
    cmd = [
        'mkvextract', 'tracks',
        input_file, f'{track_id}:{output_file}'
    ]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def convert_ass_to_srt(input_file, output_file):
    cmd = [
        'ffmpeg', '-y', '-loglevel', 'error',
        '-i', input_file, output_file
    ]
    subprocess.run(cmd)

def process_file(filename):
    base = os.path.splitext(filename)[0]
    tracks = list_subtitle_tracks(filename)

    if not tracks:
        print(f"⚠️ Nenhuma legenda encontrada em {filename}")
        return

    print(f"🎬 Processando '{filename}' - Aguarde...")

    start_time = time.time()
    for track_id, codec, lang in tracks:
        lang = lang.lower() if lang else "eng"

        # Ignorar legendas que não sejam ass, ssa ou srt
        codec_lower = codec.lower()
        print("codec lower:", codec.lower())
        valid_codecs = ['ass', 'ssa', 'substationalpha', 'srt']

        if not any(v in codec_lower for v in valid_codecs):
            print(f"   ⚠️ Pulando trilha {track_id} com codec não suportado: {codec}")
            continue

        raw_ext = 'ass' if ('ass' in codec_lower or 'ssa' in codec_lower) else 'srt'
        raw_output = f"{base}_track{track_id}.{raw_ext}"
        final_output = f"{base}.track{track_id}.srt"

        print(f"   Extraindo trilha {track_id} ({codec}, lang={lang})")

        extract_track(filename, track_id, raw_output)

        if raw_ext == 'ass':
            convert_ass_to_srt(raw_output, final_output)
            os.remove(raw_output)
        else:
            os.rename(raw_output, final_output)

    elapsed = time.time() - start_time
    print(f"✔️ Concluído em {elapsed:.1f} segundos.\n")

def main():
    mkv_files = []
    for root, dirs, files in os.walk('.'):
        for f in files:
            if f.lower().endswith('.mkv'):
                mkv_files.append(os.path.join(root, f))

    if not mkv_files:
        print("Nenhum arquivo .mkv encontrado.")
        return

    pbar = tqdm(mkv_files, desc="Arquivos MKV", unit="arquivo")

    for file in pbar:
        process_file(file)

    print("✅ Legendas extraídas com sucesso.")

if __name__ == "__main__":
    main()