mkv-episode-matcher 0.1.13__py3-none-any.whl → 0.3.0__py3-none-any.whl

mkv_episode_matcher/mkv_to_srt.py

@@ -9,15 +9,28 @@ sys.path.append(os.path.join(parent_dir, "libraries", "pgs2srt"))
 import re
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta
-
+from pathlib import Path
 import pytesseract
 from imagemaker import make_image
 from loguru import logger
 from pgsreader import PGSReader
 from PIL import Image, ImageOps
-
+from typing import Optional
 from mkv_episode_matcher.__main__ import CONFIG_FILE
 from mkv_episode_matcher.config import get_config
+def check_if_processed(filename: str) -> bool:
+    """
+    Check if the file has already been processed (has SxxExx format)
+    
+    Args:
+        filename (str): Filename to check
+        
+    Returns:
+        bool: True if file is already processed
+    """
+    import re
+    match = re.search(r"S\d+E\d+", filename)
+    return bool(match)
 
 
 def convert_mkv_to_sup(mkv_file, output_dir):
@@ -51,21 +64,23 @@ def convert_mkv_to_sup(mkv_file, output_dir):
 
 
 @logger.catch
-def perform_ocr(sup_file_path):
+def perform_ocr(sup_file_path: str) -> Optional[str]:
     """
     Perform OCR on a .sup file and save the extracted text to a .srt file.
-
-    Args:
-        sup_file_path (str): Path to the .sup file.
+    Returns the path to the created SRT file.
     """
-
     # Get the base name of the .sup file without the extension
     base_name = os.path.splitext(os.path.basename(sup_file_path))[0]
     output_dir = os.path.dirname(sup_file_path)
     logger.info(f"Performing OCR on {sup_file_path}")
+    
     # Construct the output .srt file path
     srt_file = os.path.join(output_dir, f"{base_name}.srt")
 
+    if os.path.exists(srt_file):
+        logger.info(f"SRT file {srt_file} already exists, skipping OCR")
+        return srt_file
+
     # Load a PGS/SUP file.
     pgs = PGSReader(sup_file_path)
 
@@ -151,24 +166,137 @@ def perform_ocr(sup_file_path):
         logger.info(f"Saved to: {srt_file}")
 
 
-def convert_mkv_to_srt(season_path, mkv_files):
-    """
-    Converts MKV files to SRT format.
+# def convert_mkv_to_srt(season_path, mkv_files):
+#     """
+#     Converts MKV files to SRT format.
 
-    Args:
-        season_path (str): The path to the season directory.
-        mkv_files (list): List of MKV files to convert.
+#     Args:
+#         season_path (str): The path to the season directory.
+#         mkv_files (list): List of MKV files to convert.
 
-    Returns:
-        None
+#     Returns:
+#         None
+#     """
+#     logger.info(f"Converting {len(mkv_files)} files to SRT")
+#     output_dir = os.path.join(season_path, "ocr")
+#     os.makedirs(output_dir, exist_ok=True)
+#     sup_files = []
+#     for mkv_file in mkv_files:
+#         sup_file = convert_mkv_to_sup(mkv_file, output_dir)
+#         sup_files.append(sup_file)
+#     with ThreadPoolExecutor() as executor:
+#         for sup_file in sup_files:
+#             executor.submit(perform_ocr, sup_file)
+
+        
+
+def extract_subtitles(mkv_file: str, output_dir: str) -> Optional[str]:
+    """
+    Extract subtitles from MKV file based on detected subtitle type.
+    """
+    subtitle_type, stream_index = detect_subtitle_type(mkv_file)
+    if not subtitle_type:
+        logger.error(f"No supported subtitle streams found in {mkv_file}")
+        return None
+        
+    base_name = Path(mkv_file).stem
+    
+    if subtitle_type == 'subrip':
+        # For SRT subtitles, extract directly to .srt
+        output_file = os.path.join(output_dir, f"{base_name}.srt")
+        if not os.path.exists(output_file):
+            cmd = [
+                "ffmpeg", "-i", mkv_file,
+                "-map", f"0:{stream_index}",
+                output_file
+            ]
+    else:
+        # For DVD or PGS subtitles, extract to SUP format first
+        output_file = os.path.join(output_dir, f"{base_name}.sup")
+        if not os.path.exists(output_file):
+            cmd = [
+                "ffmpeg", "-i", mkv_file,
+                "-map", f"0:{stream_index}",
+                "-c", "copy",
+                output_file
+            ]
+    
+    if not os.path.exists(output_file):
+        try:
+            subprocess.run(cmd, check=True)
+            logger.info(f"Extracted subtitles from {mkv_file} to {output_file}")
+            return output_file
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Error extracting subtitles: {e}")
+            return None
+    else:
+        logger.info(f"Subtitle file {output_file} already exists, skipping extraction")
+        return output_file
+
+def convert_mkv_to_srt(season_path: str, mkv_files: list[str]) -> None:
+    """
+    Convert subtitles from MKV files to SRT format.
     """
     logger.info(f"Converting {len(mkv_files)} files to SRT")
+    
+    # Filter out already processed files
+    unprocessed_files = []
+    for mkv_file in mkv_files:
+        if check_if_processed(os.path.basename(mkv_file)):
+            logger.info(f"Skipping {mkv_file} - already processed")
+            continue
+        unprocessed_files.append(mkv_file)
+    
+    if not unprocessed_files:
+        logger.info("No new files to process")
+        return
+        
+    # Create OCR directory
     output_dir = os.path.join(season_path, "ocr")
     os.makedirs(output_dir, exist_ok=True)
-    sup_files = []
-    for mkv_file in mkv_files:
-        sup_file = convert_mkv_to_sup(mkv_file, output_dir)
-        sup_files.append(sup_file)
-    with ThreadPoolExecutor() as executor:
-        for sup_file in sup_files:
-            executor.submit(perform_ocr, sup_file)
+    
+    for mkv_file in unprocessed_files:
+        subtitle_file = extract_subtitles(mkv_file, output_dir)
+        if not subtitle_file:
+            continue
+            
+        if subtitle_file.endswith('.srt'):
+            # Already have SRT, keep it in OCR directory
+            logger.info(f"Extracted SRT subtitle to {subtitle_file}")
+        else:
+            # For SUP files (DVD or PGS), perform OCR
+            srt_file = perform_ocr(subtitle_file)
+            if srt_file:
+                logger.info(f"Created SRT from OCR: {srt_file}")
+            
+def detect_subtitle_type(mkv_file: str) -> tuple[Optional[str], Optional[int]]:
+    """
+    Detect the type and index of subtitle streams in an MKV file.
+    """
+    cmd = ["ffmpeg", "-i", mkv_file]
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        
+        subtitle_streams = []
+        for line in result.stderr.split('\n'):
+            if 'Subtitle' in line:
+                stream_index = int(line.split('#0:')[1].split('(')[0])
+                if 'subrip' in line:
+                    subtitle_streams.append(('subrip', stream_index))
+                elif 'dvd_subtitle' in line:
+                    subtitle_streams.append(('dvd_subtitle', stream_index))
+                elif 'hdmv_pgs_subtitle' in line:
+                    subtitle_streams.append(('hdmv_pgs_subtitle', stream_index))
+        
+        # Prioritize subtitle formats: SRT > DVD > PGS
+        for format_priority in ['subrip', 'dvd_subtitle', 'hdmv_pgs_subtitle']:
+            for format_type, index in subtitle_streams:
+                if format_type == format_priority:
+                    return format_type, index
+                    
+        return None, None
+        
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Error detecting subtitle type: {e}")
+        return None, None

mkv_episode_matcher/speech_to_text.py

@@ -0,0 +1,90 @@
+# mkv_episode_matcher/speech_to_text.py
+
+import os
+import subprocess
+from pathlib import Path
+import whisper
+import torch
+from loguru import logger
+
+def process_speech_to_text(mkv_file, output_dir):
+    """
+    Convert MKV file to transcript using Whisper.
+    
+    Args:
+        mkv_file (str): Path to MKV file
+        output_dir (str): Directory to save transcript files
+    """
+    # Extract audio if not already done
+    wav_file = extract_audio(mkv_file, output_dir)
+    if not wav_file:
+        return None
+
+    # Load model
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if device == "cuda":
+        logger.info(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
+    else:
+        logger.info("CUDA not available. Using CPU.")
+    
+    model = whisper.load_model("base", device=device)
+    
+    # Generate transcript
+    segments_file = os.path.join(output_dir, f"{Path(mkv_file).stem}.segments.json")
+    if not os.path.exists(segments_file):
+        try:
+            result = model.transcribe(
+                wav_file,
+                task="transcribe",
+                language="en",
+            )
+            
+            # Save segments
+            import json
+            with open(segments_file, 'w', encoding='utf-8') as f:
+                json.dump(result["segments"], f, indent=2)
+                
+            logger.info(f"Transcript saved to {segments_file}")
+            
+        except Exception as e:
+            logger.error(f"Error during transcription: {e}")
+            return None
+    else:
+        logger.info(f"Using existing transcript: {segments_file}")
+    
+    return segments_file
+
+def extract_audio(mkv_file, output_dir):
+    """
+    Extract audio from MKV file using FFmpeg.
+    
+    Args:
+        mkv_file (str): Path to MKV file
+        output_dir (str): Directory to save WAV file
+        
+    Returns:
+        str: Path to extracted WAV file
+    """
+    wav_file = os.path.join(output_dir, f"{Path(mkv_file).stem}.wav")
+    
+    if not os.path.exists(wav_file):
+        logger.info(f"Extracting audio from {mkv_file}")
+        try:
+            cmd = [
+                'ffmpeg',
+                '-i', mkv_file,
+                '-vn',  # Disable video
+                '-acodec', 'pcm_s16le',  # Convert to PCM format
+                '-ar', '16000',  # Set sample rate to 16kHz
+                '-ac', '1',  # Convert to mono
+                wav_file
+            ]
+            subprocess.run(cmd, check=True, capture_output=True)
+            logger.info(f"Audio extracted to {wav_file}")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Error extracting audio: {e}")
+            return None
+    else:
+        logger.info(f"Audio file {wav_file} already exists, skipping extraction")
+    
+    return wav_file

mkv_episode_matcher/utils.py

@@ -2,7 +2,6 @@
 import os
 import re
 import shutil
-from typing import Set
 
 import requests
 from loguru import logger
@@ -12,32 +11,51 @@ from mkv_episode_matcher.__main__ import CACHE_DIR, CONFIG_FILE
 from mkv_episode_matcher.config import get_config
 from mkv_episode_matcher.tmdb_client import fetch_season_details
 
-
-def check_filename(filename, series_title, season_number, episode_number):
+def get_valid_seasons(show_dir):
     """
-    Check if a filename matches the expected naming convention for a series episode.
+    Get all season directories that contain MKV files.
 
     Args:
-        filename (str): The filename to be checked.
-        series_title (str): The title of the series.
-        season_number (int): The season number of the episode.
-        episode_number (int): The episode number of the episode.
+        show_dir (str): Base directory for the TV show
 
     Returns:
-        bool: True if the filename matches the expected naming convention, False otherwise.
+        list: List of paths to valid season directories
+    """
+    # Get all season directories
+    season_paths = [
+        os.path.join(show_dir, d)
+        for d in os.listdir(show_dir)
+        if os.path.isdir(os.path.join(show_dir, d))
+    ]
+
+    # Filter seasons to only include those with .mkv files
+    valid_season_paths = []
+    for season_path in season_paths:
+        mkv_files = [f for f in os.listdir(season_path) if f.endswith(".mkv")]
+        if mkv_files:
+            valid_season_paths.append(season_path)
+
+    if not valid_season_paths:
+        logger.warning(f"No seasons with .mkv files found in show '{os.path.basename(show_dir)}'")
+    else:
+        logger.info(
+            f"Found {len(valid_season_paths)} seasons with .mkv files in '{os.path.basename(show_dir)}'"
+        )
 
-    This function checks if the given filename matches the expected naming convention for a series episode.
-    The expected naming convention is '{series_title} - S{season_number:02d}E{episode_number:02d}.mkv'.
-    If the filename matches the expected pattern, it returns True; otherwise, it returns False.
+    return valid_season_paths
+def check_filename(filename):
+    """
+    Check if the filename is in the correct format (S01E02).
 
-    Example:
-        If filename = 'Example - S01E03.mkv', series_title = 'Example', season_number = 1, and episode_number = 3,
-        the function will return True because the filename matches the expected pattern.
+    Args:
+        filename (str): The filename to check.
+
+    Returns:
+        bool: True if the filename matches the expected pattern.
     """
-    pattern = re.compile(
-        f"{re.escape(series_title)} - S{season_number:02d}E{episode_number:02d}.mkv"
-    )
-    return bool(pattern.match(filename))
+    # Check if the filename matches the expected format
+    match = re.search(r'.*S\d+E\d+', filename)
+    return bool(match)
 
 
 def scramble_filename(original_file_path, file_number):
@@ -64,60 +82,44 @@ def scramble_filename(original_file_path, file_number):
         os.rename(original_file_path, new_file_path)
 
 
-def rename_episode_file(original_file_path, season_number, episode_number):
+def rename_episode_file(original_file_path, new_filename):
     """
     Rename an episode file with a standardized naming convention.
 
     Args:
         original_file_path (str): The original file path of the episode.
-        season_number (int): The season number of the episode.
-        episode_number (int): The episode number of the episode.
+        new_filename (str): The new filename including season/episode info.
 
     Returns:
-        None
-
-    This function renames an episode file with a standardized naming convention based on the series title, season number,
-    and episode number. If a file with the intended new name already exists, it appends a numerical suffix to the filename
-    until it finds a unique name.
-
-    Example:
-        If original_file_path = '/path/to/episode.mkv', season_number = 1, and episode_number = 3, and the series title is 'Example',
-        the function will rename the file to 'Example - S01E03.mkv' if no file with that name already exists. If a file with that
-        name already exists, it will be renamed to 'Example - S01E03_2.mkv', and so on.
+        str: Path to the renamed file, or None if rename failed.
     """
-    series_title = os.path.basename(
-        os.path.dirname(os.path.dirname(original_file_path))
-    )
-    original_file_name = os.path.basename(original_file_path)
-    extension = os.path.splitext(original_file_path)[-1]
-    new_file_name = (
-        f"{series_title} - S{season_number:02d}E{episode_number:02d}{extension}"
-    )
-    new_file_path = os.path.join(os.path.dirname(original_file_path), new_file_name)
-
-    # Check if the new file path already exists
+    original_dir = os.path.dirname(original_file_path)
+    new_file_path = os.path.join(original_dir, new_filename)
+    
+    # Check if new filepath already exists
     if os.path.exists(new_file_path):
-        logger.warning(f"Filename already exists: {new_file_name}.")
-
-        # If the file already exists, find a unique name by appending a numerical suffix
+        logger.warning(f"File already exists: {new_filename}")
+        
+        # Add numeric suffix if file exists
+        base, ext = os.path.splitext(new_filename)
         suffix = 2
         while True:
-            new_file_name = f"{series_title} - S{season_number:02d}E{episode_number:02d}_{suffix}{extension}"
-            new_file_path = os.path.join(
-                os.path.dirname(original_file_path), new_file_name
-            )
+            new_filename = f"{base}_{suffix}{ext}"
+            new_file_path = os.path.join(original_dir, new_filename)
             if not os.path.exists(new_file_path):
                 break
             suffix += 1
-
-        logger.info(f"Renaming {original_file_name} -> {new_file_name}")
-        os.rename(original_file_path, new_file_path)
-    else:
-        logger.info(f"Renaming {original_file_name} -> {new_file_name}")
+    
+    try:
         os.rename(original_file_path, new_file_path)
+        logger.info(f"Renamed {os.path.basename(original_file_path)} -> {new_filename}")
+        return new_file_path
+    except OSError as e:
+        logger.error(f"Failed to rename file: {e}")
+        return None
 
 
-def get_subtitles(show_id, seasons: Set[int]):
+def get_subtitles(show_id, seasons: set[int]):
     """
     Retrieves and saves subtitles for a given TV show and seasons.
 
@@ -138,16 +140,14 @@ def get_subtitles(show_id, seasons: Set[int]):
     open_subtitles_user_agent = config.get("open_subtitles_user_agent")
     open_subtitles_username = config.get("open_subtitles_username")
     open_subtitles_password = config.get("open_subtitles_password")
-    if not all(
-        [
-            show_dir,
-            tmdb_api_key,
-            open_subtitles_api_key,
-            open_subtitles_user_agent,
-            open_subtitles_username,
-            open_subtitles_password,
-        ]
-    ):
+    if not all([
+        show_dir,
+        tmdb_api_key,
+        open_subtitles_api_key,
+        open_subtitles_user_agent,
+        open_subtitles_username,
+        open_subtitles_password,
+    ]):
         logger.error("Missing configuration settings. Please run the setup script.")
     try:
         # Initialize the OpenSubtitles client
@@ -164,11 +164,8 @@ def get_subtitles(show_id, seasons: Set[int]):
 
         for episode in range(1, episodes + 1):
             logger.info(f"Processing Season {season}, Episode {episode}...")
-            series_cache_dir =os.path.join(
-                CACHE_DIR,
-                "data",
-                series_name)
-            os.makedirs(series_cache_dir,exist_ok=True)
+            series_cache_dir = os.path.join(CACHE_DIR, "data", series_name)
+            os.makedirs(series_cache_dir, exist_ok=True)
             srt_filepath = os.path.join(
                 series_cache_dir,
                 f"{series_name} - S{season:02d}E{episode:02d}.srt",
@@ -179,7 +176,7 @@ def get_subtitles(show_id, seasons: Set[int]):
                 response = requests.get(url)
                 response.raise_for_status()
                 episode_data = response.json()
-                episode_name = episode_data["name"]
+                episode_data["name"]
                 episode_id = episode_data["id"]
                 # search for the subtitle
                 response = subtitles.search(tmdb_id=episode_id, languages="en")
@@ -229,8 +226,159 @@ def cleanup_ocr_files(show_dir):
         if os.path.exists(ocr_dir_path):
             logger.info(f"Cleaning up OCR files in {ocr_dir_path}")
             shutil.rmtree(ocr_dir_path)
+
+
 def clean_text(text):
     # Remove brackets, parentheses, and their content
-    cleaned_text = re.sub(r'\[.*?\]|\(.*?\)|\{.*?\}', '', text)
+    cleaned_text = re.sub(r"\[.*?\]|\(.*?\)|\{.*?\}", "", text)
     # Strip leading/trailing whitespace
-    return cleaned_text.strip()
+    return cleaned_text.strip()
+# mkv_episode_matcher/utils.py
+
+# Add this to your existing utils.py, keeping all other functions
+
+def process_reference_srt_files(series_name):
+    """
+    Process reference SRT files for a given series.
+
+    Args:
+        series_name (str): The name of the series.
+
+    Returns:
+        dict: A dictionary containing the reference files where the keys are the MKV filenames
+              and the values are the corresponding SRT texts.
+    """
+    from mkv_episode_matcher.__main__ import CACHE_DIR
+    import os
+    
+    reference_files = {}
+    reference_dir = os.path.join(CACHE_DIR, "data", series_name)
+    
+    for dirpath, _, filenames in os.walk(reference_dir):
+        for filename in filenames:
+            if filename.lower().endswith(".srt"):
+                srt_file = os.path.join(dirpath, filename)
+                logger.info(f"Processing {srt_file}")
+                srt_text = extract_srt_text(srt_file)
+                season, episode = extract_season_episode(filename)
+                mkv_filename = f"{series_name} - S{season:02}E{episode:02}.mkv"
+                reference_files[mkv_filename] = srt_text
+                
+    return reference_files
+
+def extract_srt_text(filepath):
+    """
+    Extracts text content from an SRT file.
+
+    Args:
+        filepath (str): Path to the SRT file.
+
+    Returns:
+        list: List of text lines from the SRT file.
+    """
+    # Read the file content
+    with open(filepath) as f:
+        content = f.read()
+        
+    # Split into subtitle blocks
+    blocks = content.strip().split('\n\n')
+    
+    text_lines = []
+    for block in blocks:
+        lines = block.split('\n')
+        if len(lines) < 3:
+            continue
+            
+        # Skip index and timestamp, get all remaining lines as text
+        text = ' '.join(lines[2:])
+        # Remove stage directions and tags
+        text = re.sub(r'\[.*?\]|\<.*?\>', '', text)
+        if text:
+            text_lines.append(text)
+            
+    return text_lines
+
+def extract_season_episode(filename):
+    """
+    Extract season and episode numbers from filename.
+    
+    Args:
+        filename (str): Filename to parse
+        
+    Returns:
+        tuple: (season_number, episode_number)
+    """
+    match = re.search(r'S(\d+)E(\d+)', filename)
+    if match:
+        return int(match.group(1)), int(match.group(2))
+    return None, None
+def process_srt_files(show_dir):
+    """
+    Process all SRT files in the given directory and its subdirectories.
+
+    Args:
+        show_dir (str): The directory path where the SRT files are located.
+
+    Returns:
+        dict: A dictionary containing the SRT file paths as keys and their corresponding text content as values.
+    """
+    srt_files = {}
+    for dirpath, _, filenames in os.walk(show_dir):
+        for filename in filenames:
+            if filename.lower().endswith(".srt"):
+                srt_file = os.path.join(dirpath, filename)
+                logger.info(f"Processing {srt_file}")
+                srt_text = extract_srt_text(srt_file)
+                srt_files[srt_file] = srt_text
+    return srt_files
+def compare_and_rename_files(srt_files, reference_files, dry_run=False):
+    """
+    Compare the srt files with the reference files and rename the matching mkv files.
+
+    Args:
+        srt_files (dict): A dictionary containing the srt files as keys and their contents as values.
+        reference_files (dict): A dictionary containing the reference files as keys and their contents as values.
+        dry_run (bool, optional): If True, the function will only log the renaming actions without actually renaming the files. Defaults to False.
+    """
+    logger.info(
+        f"Comparing {len(srt_files)} srt files with {len(reference_files)} reference files"
+    )
+    for srt_text in srt_files.keys():
+        parent_dir = os.path.dirname(os.path.dirname(srt_text))
+        for reference in reference_files.keys():
+            _season, _episode = extract_season_episode(reference)
+            mkv_file = os.path.join(
+                parent_dir, os.path.basename(srt_text).replace(".srt", ".mkv")
+            )
+            matching_lines = compare_text(
+                reference_files[reference], srt_files[srt_text]
+            )
+            if matching_lines >= int(len(reference_files[reference]) * 0.1):
+                logger.info(f"Matching lines: {matching_lines}")
+                logger.info(f"Found matching file: {mkv_file} ->{reference}")
+                new_filename = os.path.join(parent_dir, reference)
+                if not os.path.exists(new_filename):
+                    if os.path.exists(mkv_file) and not dry_run:
+                        logger.info(f"Renaming {mkv_file} to {new_filename}")
+                        os.rename(mkv_file, new_filename)
+                else:
+                    logger.info(f"File {new_filename} already exists, skipping")
+
+def compare_text(text1, text2):
+    """
+    Compare two lists of text lines and return the number of matching lines.
+
+    Args:
+        text1 (list): List of text lines from the first source.
+        text2 (list): List of text lines from the second source.
+
+    Returns:
+        int: Number of matching lines between the two sources.
+    """
+    # Flatten the list of text lines
+    flat_text1 = [line for lines in text1 for line in lines]
+    flat_text2 = [line for lines in text2 for line in lines]
+
+    # Compare the two lists of text lines
+    matching_lines = set(flat_text1).intersection(flat_text2)
+    return len(matching_lines)