mkv-episode-matcher 0.1.5__py3-none-any.whl → 0.1.10__py3-none-any.whl

mkv_episode_matcher/mkv_to_srt.py

@@ -1,174 +1,174 @@
-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
-
-import pytesseract
-from imagemaker import make_image
-from loguru import logger
-from pgsreader import PGSReader
-from PIL import Image, ImageOps
-
-from mkv_episode_matcher.__main__ import CONFIG_FILE
-from mkv_episode_matcher.config import get_config
-
-
-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):
-    """
-    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.
-    """
-
-    # 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")
-
-    # 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)
+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
+
+import pytesseract
+from imagemaker import make_image
+from loguru import logger
+from pgsreader import PGSReader
+from PIL import Image, ImageOps
+
+from mkv_episode_matcher.__main__ import CONFIG_FILE
+from mkv_episode_matcher.config import get_config
+
+
+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):
+    """
+    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.
+    """
+
+    # 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")
+
+    # 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)

mkv_episode_matcher/notebooks/get_subtitles_test.ipynb

@@ -0,0 +1,252 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load imports and create config directory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# __main__.py\n",
+    "import argparse\n",
+    "import os\n",
+    "\n",
+    "from loguru import logger\n",
+    "\n",
+    "from mkv_episode_matcher.config import get_config, set_config\n",
+    "from mkv_episode_matcher.utils import check_filename, cleanup_ocr_files, get_subtitles\n",
+    "from mkv_episode_matcher.tmdb_client import fetch_show_id,fetch_season_details\n",
+    "import os\n",
+    "import re\n",
+    "import shutil\n",
+    "from typing import Set\n",
+    "\n",
+    "import requests\n",
+    "from loguru import logger\n",
+    "from opensubtitlescom import OpenSubtitles\n",
+    "# Log the start of the application\n",
+    "logger.info(\"Starting the application\")\n",
+    "\n",
+    "\n",
+    "\n",
+    "# Check if the configuration directory exists, if not create it\n",
+    "if not os.path.exists(os.path.join(os.path.expanduser(\"~\"), \".mkv-episode-matcher\")):\n",
+    "    os.makedirs(os.path.join(os.path.expanduser(\"~\"), \".mkv-episode-matcher\"))\n",
+    "\n",
+    "# Define the paths for the configuration file and cache directory\n",
+    "CONFIG_FILE = os.path.join(\n",
+    "    os.path.expanduser(\"~\"), \".mkv-episode-matcher\", \"config.ini\"\n",
+    ")\n",
+    "CACHE_DIR = os.path.join(os.path.expanduser(\"~\"), \".mkv-episode-matcher\", \"cache\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load configuration settings from config.ini"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = get_config(CONFIG_FILE)\n",
+    "show_dir = config.get(\"show_dir\")\n",
+    "show_name = os.path.basename(show_dir)\n",
+    "series_name = os.path.basename(show_dir)\n",
+    "tmdb_api_key = config.get(\"tmdb_api_key\")\n",
+    "open_subtitles_api_key = config.get(\"open_subtitles_api_key\")\n",
+    "open_subtitles_user_agent = config.get(\"open_subtitles_user_agent\")\n",
+    "open_subtitles_username = config.get(\"open_subtitles_username\")\n",
+    "open_subtitles_password = config.get(\"open_subtitles_password\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Make sure all required info exists in config.ini"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for x in     [\n",
+    "        show_dir,\n",
+    "        tmdb_api_key,\n",
+    "        open_subtitles_api_key,\n",
+    "        open_subtitles_user_agent,\n",
+    "        open_subtitles_username,\n",
+    "        open_subtitles_password,\n",
+    "    ]:\n",
+    "    try:\n",
+    "        print(x)\n",
+    "    except:\n",
+    "        print('failed')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not all(\n",
+    "    [\n",
+    "        show_dir,\n",
+    "        tmdb_api_key,\n",
+    "        open_subtitles_api_key,\n",
+    "        open_subtitles_user_agent,\n",
+    "        open_subtitles_username,\n",
+    "        open_subtitles_password,\n",
+    "    ]\n",
+    "):\n",
+    "    logger.error(\"Missing configuration settings. Please run the setup script.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Make sure show can be found on TMDb\n",
+    "The show id is used to search on opensubtitles"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_id = fetch_show_id(show_name)\n",
+    "if show_id is None:\n",
+    "    logger.error(f\"Could not find show '{os.path.basename(show_name)}' on TMDb.\")\n",
+    "else:\n",
+    "    print(show_id)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Try getting the first season automatically"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "get_subtitles(show_id, seasons=set([1]))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Check if there's an issue in the get_subtitles function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "try:\n",
+    "    # Initialize the OpenSubtitles client\n",
+    "    subtitles = OpenSubtitles(open_subtitles_user_agent, open_subtitles_api_key)\n",
+    "\n",
+    "    # Log in (retrieve auth token)\n",
+    "    subtitles.login(open_subtitles_username, open_subtitles_password)\n",
+    "except Exception as e:\n",
+    "    logger.error(f\"Failed to log in to OpenSubtitles: {e}\")\n",
+    "\n",
+    "for season in [1]:\n",
+    "    episodes = fetch_season_details(show_id, season)\n",
+    "    logger.info(f\"Found {episodes} episodes in Season {season}\")\n",
+    "\n",
+    "    for episode in range(1, episodes + 1):\n",
+    "        logger.info(f\"Processing Season {season}, Episode {episode}...\")\n",
+    "        series_cache_dir =os.path.join(\n",
+    "            CACHE_DIR,\n",
+    "            \"data\",\n",
+    "            series_name)\n",
+    "        os.makedirs(series_cache_dir,exist_ok=True)\n",
+    "        srt_filepath = os.path.join(\n",
+    "            series_cache_dir,\n",
+    "            f\"{series_name} - S{season:02d}E{episode:02d}.srt\",\n",
+    "        )\n",
+    "        if not os.path.exists(srt_filepath):\n",
+    "            # get the episode info from TMDB\n",
+    "            url = f\"https://api.themoviedb.org/3/tv/{show_id}/season/{season}/episode/{episode}?api_key={tmdb_api_key}\"\n",
+    "            response = requests.get(url)\n",
+    "            response.raise_for_status()\n",
+    "            episode_data = response.json()\n",
+    "            episode_name = episode_data[\"name\"]\n",
+    "            episode_id = episode_data[\"id\"]\n",
+    "            # search for the subtitle\n",
+    "            response = subtitles.search(tmdb_id=episode_id, languages=\"en\")\n",
+    "            if len(response.data) == 0:\n",
+    "                logger.warning(\n",
+    "                    f\"No subtitles found for {series_name} - S{season:02d}E{episode:02d}\"\n",
+    "                )\n",
+    "\n",
+    "            for subtitle in response.data:\n",
+    "                subtitle_dict = subtitle.to_dict()\n",
+    "                # Remove special characters and convert to uppercase\n",
+    "                filename_clean = re.sub(\n",
+    "                    r\"\\W+\", \" \", subtitle_dict[\"file_name\"]\n",
+    "                ).upper()\n",
+    "                if f\"E{episode:02d}\" in filename_clean:\n",
+    "                    logger.info(f\"Original filename: {subtitle_dict['file_name']}\")\n",
+    "                    srt_file = subtitles.download_and_save(subtitle)\n",
+    "                    series_name = series_name.replace(\":\", \" -\")\n",
+    "                    shutil.move(os.path.join(os.getcwd(),srt_file), srt_filepath)\n",
+    "                    logger.info(f\"Subtitle saved to {srt_filepath}\")\n",
+    "                    break\n",
+    "                else:\n",
+    "                    continue\n",
+    "        else:\n",
+    "            logger.info(\n",
+    "                f\"Subtitle already exists for {series_name} - S{season:02d}E{episode:02d}\"\n",
+    "            )\n",
+    "            continue"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mkv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

mkv_episode_matcher/requirements.txt

@@ -1,8 +1,7 @@
-requests
-loguru
-pillow
-imagehash
-configparser
-tmdb_client
-pytesseract
+requests
+loguru
+pillow
+configparser
+tmdb_client
+pytesseract
 opensubtitlescom

mkv_episode_matcher/utils.py

@@ -164,10 +164,13 @@ def get_subtitles(show_id, seasons: Set[int]):
 
         for episode in range(1, episodes + 1):
             logger.info(f"Processing Season {season}, Episode {episode}...")
-            srt_filepath = os.path.join(
+            series_cache_dir =os.path.join(
                 CACHE_DIR,
                 "data",
-                series_name,
+                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",
             )
             if not os.path.exists(srt_filepath):