Merge pull request #22 from Jsakkos/19-add-support-non-pgs-subtitles-ie-from-dvds

19 add support non pgs subtitles ie from dvds

Author: Jsakkos

mkv_episode_matcher/episode_matcher.py

@@ -74,28 +74,36 @@ def process_show(season=None, dry_run=False, get_subs=False):
             logger.warning(f"Season {season} has no .mkv files to process")
             return
 
+        season_paths_to_process = [season_path]
+    else:
+        # Process all valid seasons
+        season_paths_to_process = valid_season_paths
+
+    # Process each season
+    for season_path in season_paths_to_process:
+        logger.info(f"Processing season path: {season_path}")
         mkv_files = [
             os.path.join(season_path, f)
             for f in os.listdir(season_path)
             if f.endswith(".mkv")
         ]
-    else:
-        # Process all valid seasons
-        for season_path in valid_season_paths:
-            mkv_files = [
-                os.path.join(season_path, f)
-                for f in os.listdir(season_path)
-                if f.endswith(".mkv")
-            ]
-    # Filter out files that have already been processed
-    for f in mkv_files:
-        if check_filename(f):
-            logger.info(f"Skipping {f}, already processed")
-            mkv_files.remove(f)
-    if len(mkv_files) == 0:
-        logger.info("No new files to process")
-        return
-    convert_mkv_to_srt(season_path, mkv_files)
+        
+        # Filter out files that have already been processed
+        unprocessed_files = []
+        for f in mkv_files:
+            if check_filename(f):
+                logger.info(f"Skipping {f}, already processed")
+            else:
+                unprocessed_files.append(f)
+
+        if not unprocessed_files:
+            logger.info(f"No new files to process in {season_path}")
+            continue
+
+        logger.info(f"Processing {len(unprocessed_files)} files in {season_path}")
+        convert_mkv_to_srt(season_path, unprocessed_files)
+        
+    # Process reference and SRT files after all seasons are converted
     reference_text_dict = process_reference_srt_files(show_name)
     srt_text_dict = process_srt_files(show_dir)
     compare_and_rename_files(srt_text_dict, reference_text_dict, dry_run=dry_run)

mkv_episode_matcher/mkv_to_srt.py

@@ -9,15 +9,28 @@
 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

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = mkv_episode_matcher
-version = 0.1.13
+version = 0.2.0
 author = Jonathan Sakkos
 author_email = [email protected]
 description = The MKV Episode Matcher is a tool for identifying TV series episodes from MKV files and renaming the files accordingly.