mkv-episode-matcher 0.3.3__py3-none-any.whl → 1.0.0__py3-none-any.whl

mkv_episode_matcher/libraries/pgs2srt/pgsreader.py

@@ -1,225 +0,0 @@
-#!/usr/bin/env python3
-
-from collections import namedtuple
-from os.path import split as pathsplit
-
-# Constants for Segments
-PDS = int("0x14", 16)
-ODS = int("0x15", 16)
-PCS = int("0x16", 16)
-WDS = int("0x17", 16)
-END = int("0x80", 16)
-
-# Named tuple access for static PDS palettes
-Palette = namedtuple("Palette", "Y Cr Cb Alpha")
-
-
-class InvalidSegmentError(Exception):
-    """Raised when a segment does not match PGS specification"""
-
-
-class PGSReader:
-    def __init__(self, filepath):
-        self.filedir, self.file = pathsplit(filepath)
-        with open(filepath, "rb") as f:
-            self.bytes = f.read()
-
-    def make_segment(self, bytes_):
-        cls = SEGMENT_TYPE[bytes_[10]]
-        return cls(bytes_)
-
-    def iter_segments(self):
-        bytes_ = self.bytes[:]
-        while bytes_:
-            size = 13 + int(bytes_[11:13].hex(), 16)
-            yield self.make_segment(bytes_[:size])
-            bytes_ = bytes_[size:]
-
-    def iter_displaysets(self):
-        ds = []
-        for s in self.iter_segments():
-            ds.append(s)
-            if s.type == "END":
-                yield DisplaySet(ds)
-                ds = []
-
-    @property
-    def segments(self):
-        if not hasattr(self, "_segments"):
-            self._segments = list(self.iter_segments())
-        return self._segments
-
-    @property
-    def displaysets(self):
-        if not hasattr(self, "_displaysets"):
-            self._displaysets = list(self.iter_displaysets())
-        return self._displaysets
-
-
-class BaseSegment:
-    SEGMENT = {PDS: "PDS", ODS: "ODS", PCS: "PCS", WDS: "WDS", END: "END"}
-
-    def __init__(self, bytes_):
-        self.bytes = bytes_
-        if bytes_[:2] != b"PG":
-            raise InvalidSegmentError
-        self.pts = int(bytes_[2:6].hex(), base=16) / 90
-        self.dts = int(bytes_[6:10].hex(), base=16) / 90
-        self.type = self.SEGMENT[bytes_[10]]
-        self.size = int(bytes_[11:13].hex(), base=16)
-        self.data = bytes_[13:]
-
-    def __len__(self):
-        return self.size
-
-    @property
-    def presentation_timestamp(self):
-        return self.pts
-
-    @property
-    def decoding_timestamp(self):
-        return self.dts
-
-    @property
-    def segment_type(self):
-        return self.type
-
-
-class PresentationCompositionSegment(BaseSegment):
-    class CompositionObject:
-        def __init__(self, bytes_):
-            self.bytes = bytes_
-            self.object_id = int(bytes_[0:2].hex(), base=16)
-            self.window_id = bytes_[2]
-            self.cropped = bool(bytes_[3])
-            self.x_offset = int(bytes_[4:6].hex(), base=16)
-            self.y_offset = int(bytes_[6:8].hex(), base=16)
-            if self.cropped:
-                self.crop_x_offset = int(bytes_[8:10].hex(), base=16)
-                self.crop_y_offset = int(bytes_[10:12].hex(), base=16)
-                self.crop_width = int(bytes_[12:14].hex(), base=16)
-                self.crop_height = int(bytes_[14:16].hex(), base=16)
-
-    STATE = {
-        int("0x00", base=16): "Normal",
-        int("0x40", base=16): "Acquisition Point",
-        int("0x80", base=16): "Epoch Start",
-    }
-
-    def __init__(self, bytes_):
-        BaseSegment.__init__(self, bytes_)
-        self.width = int(self.data[0:2].hex(), base=16)
-        self.height = int(self.data[2:4].hex(), base=16)
-        self.frame_rate = self.data[4]
-        self._num = int(self.data[5:7].hex(), base=16)
-        self._state = self.STATE[self.data[7]]
-        self.palette_update = bool(self.data[8])
-        self.palette_id = self.data[9]
-        self._num_comps = self.data[10]
-
-    @property
-    def composition_number(self):
-        return self._num
-
-    @property
-    def composition_state(self):
-        return self._state
-
-    @property
-    def composition_objects(self):
-        if not hasattr(self, "_composition_objects"):
-            self._composition_objects = self.get_composition_objects()
-            if len(self._composition_objects) != self._num_comps:
-                print(
-                    "Warning: Number of composition objects asserted "
-                    "does not match the amount found."
-                )
-        return self._composition_objects
-
-    def get_composition_objects(self):
-        bytes_ = self.data[11:]
-        comps = []
-        while bytes_:
-            length = 8 * (1 + bool(bytes_[3]))
-            comps.append(self.CompositionObject(bytes_[:length]))
-            bytes_ = bytes_[length:]
-        return comps
-
-
-class WindowDefinitionSegment(BaseSegment):
-    def __init__(self, bytes_):
-        BaseSegment.__init__(self, bytes_)
-        self.num_windows = self.data[0]
-        self.window_id = self.data[1]
-        self.x_offset = int(self.data[2:4].hex(), base=16)
-        self.y_offset = int(self.data[4:6].hex(), base=16)
-        self.width = int(self.data[6:8].hex(), base=16)
-        self.height = int(self.data[8:10].hex(), base=16)
-
-
-class PaletteDefinitionSegment(BaseSegment):
-    def __init__(self, bytes_):
-        BaseSegment.__init__(self, bytes_)
-        self.palette_id = self.data[0]
-        self.version = self.data[1]
-        self.palette = [Palette(0, 0, 0, 0)] * 256
-        # Slice from byte 2 til end of segment. Divide by 5 to determine number of palette entries
-        # Iterate entries. Explode the 5 bytes into namedtuple Palette. Must be exploded
-        for entry in range(len(self.data[2:]) // 5):
-            i = 2 + entry * 5
-            self.palette[self.data[i]] = Palette(*self.data[i + 1 : i + 5])
-
-
-class ObjectDefinitionSegment(BaseSegment):
-    SEQUENCE = {
-        int("0x40", base=16): "Last",
-        int("0x80", base=16): "First",
-        int("0xc0", base=16): "First and last",
-    }
-
-    def __init__(self, bytes_):
-        BaseSegment.__init__(self, bytes_)
-        self.id = int(self.data[0:2].hex(), base=16)
-        self.version = self.data[2]
-        self.in_sequence = self.SEQUENCE[self.data[3]]
-        self.data_len = int(self.data[4:7].hex(), base=16)
-        self.width = int(self.data[7:9].hex(), base=16)
-        self.height = int(self.data[9:11].hex(), base=16)
-        self.img_data = self.data[11:]
-        if len(self.img_data) != self.data_len - 4:
-            print(
-                "Warning: Image data length asserted does not match the length found."
-            )
-
-
-class EndSegment(BaseSegment):
-    @property
-    def is_end(self):
-        return True
-
-
-SEGMENT_TYPE = {
-    PDS: PaletteDefinitionSegment,
-    ODS: ObjectDefinitionSegment,
-    PCS: PresentationCompositionSegment,
-    WDS: WindowDefinitionSegment,
-    END: EndSegment,
-}
-
-
-class DisplaySet:
-    def __init__(self, segments):
-        self.segments = segments
-        self.segment_types = [s.type for s in segments]
-        self.has_image = "ODS" in self.segment_types
-
-
-def segment_by_type_getter(type_):
-    def f(self):
-        return [s for s in self.segments if s.type == type_]
-
-    return f
-
-
-for type_ in BaseSegment.SEGMENT.values():
-    setattr(DisplaySet, type_.lower(), property(segment_by_type_getter(type_)))

mkv_episode_matcher/libraries/pgs2srt/requirements.txt

@@ -1,4 +0,0 @@
-pytesseract==0.3.7
-numpy==1.19.4
-Pillow==8.2.0
-tld~=0.12.3

mkv_episode_matcher/mkv_to_srt.py

@@ -1,302 +0,0 @@
-import os
-import subprocess
-import sys
-
-# Get the absolute path of the parent directory of the current script.
-parent_dir = os.path.dirname(os.path.abspath(__file__))
-# Add the 'pgs2srt' directory to the Python path.
-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):
-    """
-    Convert an .mkv file to a .sup file using FFmpeg and pgs2srt.
-
-    Args:
-        mkv_file (str): Path to the .mkv file.
-        output_dir (str): Path to the directory where the .sup file will be saved.
-
-    Returns:
-        str: Path to the converted .sup file.
-    """
-    # Get the base name of the .mkv file without the extension
-    base_name = os.path.splitext(os.path.basename(mkv_file))[0]
-
-    # Construct the output .sup file path
-    sup_file = os.path.join(output_dir, f"{base_name}.sup")
-    if not os.path.exists(sup_file):
-        logger.info(f"Processing {mkv_file} to {sup_file}")
-        # FFmpeg command to convert .mkv to .sup
-        ffmpeg_cmd = ["ffmpeg", "-i", mkv_file, "-map", "0:s:0", "-c", "copy", sup_file]
-        try:
-            subprocess.run(ffmpeg_cmd, check=True)
-            logger.info(f"Converted {mkv_file} to {sup_file}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"Error converting {mkv_file}: {e}")
-    else:
-        logger.info(f"File {sup_file} already exists, skipping")
-    return sup_file
-
-
-@logger.catch
-def perform_ocr(sup_file_path: str) -> Optional[str]:
-    """
-    Perform OCR on a .sup file and save the extracted text to a .srt 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)
-
-    # Set index
-    i = 0
-
-    # Complete subtitle track index
-    si = 0
-
-    tesseract_lang = "eng"
-    tesseract_config = f"-c tessedit_char_blacklist=[] --psm 6 --oem {1}"
-
-    config = get_config(CONFIG_FILE)
-    tesseract_path = config.get("tesseract_path")
-    logger.debug(f"Setting Teesseract Path to {tesseract_path}")
-    pytesseract.pytesseract.tesseract_cmd = str(tesseract_path)
-
-    # SubRip output
-    output = ""
-
-    if not os.path.exists(srt_file):
-        # Iterate the pgs generator
-        for ds in pgs.iter_displaysets():
-            # If set has image, parse the image
-            if ds.has_image:
-                # Get Palette Display Segment
-                pds = ds.pds[0]
-                # Get Object Display Segment
-                ods = ds.ods[0]
-
-                if pds and ods:
-                    # Create and show the bitmap image and convert it to RGBA
-                    src = make_image(ods, pds).convert("RGBA")
-
-                    # Create grayscale image with black background
-                    img = Image.new("L", src.size, "BLACK")
-                    # Paste the subtitle bitmap
-                    img.paste(src, (0, 0), src)
-                    # Invert images so the text is readable by Tesseract
-                    img = ImageOps.invert(img)
-
-                    # Parse the image with tesesract
-                    text = pytesseract.image_to_string(
-                        img, lang=tesseract_lang, config=tesseract_config
-                    ).strip()
-
-                    # Replace "|" with "I"
-                    # Works better than blacklisting "|" in Tesseract,
-                    # which results in I becoming "!" "i" and "1"
-                    text = re.sub(r"[|/\\]", "I", text)
-                    text = re.sub(r"[_]", "L", text)
-                    start = datetime.fromtimestamp(ods.presentation_timestamp / 1000)
-                    start = start + timedelta(hours=-1)
-
-            else:
-                # Get Presentation Composition Segment
-                pcs = ds.pcs[0]
-
-                if pcs:
-                    end = datetime.fromtimestamp(pcs.presentation_timestamp / 1000)
-                    end = end + timedelta(hours=-1)
-
-                    if (
-                        isinstance(start, datetime)
-                        and isinstance(end, datetime)
-                        and len(text)
-                    ):
-                        si = si + 1
-                        sub_output = str(si) + "\n"
-                        sub_output += (
-                            start.strftime("%H:%M:%S,%f")[0:12]
-                            + " --> "
-                            + end.strftime("%H:%M:%S,%f")[0:12]
-                            + "\n"
-                        )
-                        sub_output += text + "\n\n"
-
-                        output += sub_output
-                        start = end = text = None
-            i = i + 1
-        with open(srt_file, "w") as f:
-            f.write(output)
-        logger.info(f"Saved to: {srt_file}")
-
-
-# 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.
-
-#     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)
-    
-    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

@@ -1,90 +0,0 @@
-# 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