OCR improvements

.gitignore

@@ -3,6 +3,9 @@
 .ruff_cache/
 *.pyc
 
+# Other
+.config
+
 # Build
 __pycache__/
 dist/

subtle_gui/config.py

@@ -51,6 +51,8 @@
     },
     "Settings": {
         "tessData": "",
+        "dialogLine": "-",
+        "dialogSpace": True,
     },
     "Fonts": {
         "guiFont": "",
@@ -158,6 +160,10 @@ def getSetting(self, key: str) -> str:
         """Get a generic string setting."""
         return str(self._data["Settings"].get(key, ""))
 
+    def getFlag(self, key: str) -> bool:
+        """Get a boolean setting."""
+        return bool(self._data["Settings"].get(key, False))
+
     def assetPath(self, resource: str, kind: str | None = None) -> Path:
         """Return the path to an asset."""
         path = self._appPath / "assets"
@@ -182,6 +188,14 @@ def setSizes(self, key: str, value: list[int]) -> None:
             except Exception as e:
                 logger.error("Problem when saving sizes list", exc_info=e)
 
+    def setSetting(self, key: str, value: str) -> None:
+        """Set a generic string setting."""
+        self._data["Settings"][key] = str(value)
+
+    def setFlag(self, key: str, value: bool) -> None:
+        """Set a boolean setting."""
+        self._data["Settings"][key] = bool(value)
+
     def setFontSpec(self, target: T_Fonts, font: QFont | str) -> None:
         """Set a font in config."""
         if isinstance(font, str):

subtle_gui/constants.py

@@ -50,3 +50,9 @@ class GuiLabels:
         MediaType.SUBS: QT_TRANSLATE_NOOP("Constant", "Subtitles"),
         MediaType.OTHER: QT_TRANSLATE_NOOP("Constant", "Other"),
     }
+
+
+class Constants:
+    """Various constants."""
+
+    DIALOG_LINES: Final[tuple[str, str, str]] = ("-", "\u2013", "\u2014")

subtle_gui/gui/imageviewer.py

@@ -55,6 +55,12 @@ def __init__(self, parent: QWidget) -> None:
     #  Public Slots
     ##
 
+    @pyqtSlot()
+    def processNewMediaLoaded(self) -> None:
+        """Clear previous content."""
+        self._imageSize = QRect(0, 0, 0, 0)
+        self.imageView.setScene(None)
+
     @pyqtSlot(FrameBase)
     def processFrameUpdate(self, frame: FrameBase) -> None:
         """Process frame update."""

subtle_gui/gui/texteditor.py

@@ -116,6 +116,12 @@ def __init__(self, parent: QWidget) -> None:
     #  Public Slots
     ##
 
+    @pyqtSlot()
+    def processNewMediaLoaded(self) -> None:
+        """Clear previous content."""
+        self._frame = None
+        self.textEdit.clear()
+
     @pyqtSlot(FrameBase)
     def setEditorText(self, frame: FrameBase) -> None:
         """Set the editor text."""

subtle_gui/gui/toolspanel.py

@@ -28,6 +28,7 @@
 from PyQt6.QtCore import Qt, pyqtSignal, pyqtSlot
 from PyQt6.QtWidgets import (
     QCheckBox,
+    QComboBox,
     QDoubleSpinBox,
     QFormLayout,
     QGroupBox,
@@ -40,8 +41,8 @@
     QWidget,
 )
 
-from subtle_gui import SHARED
-from subtle_gui.constants import MediaType
+from subtle_gui import CONFIG, SHARED
+from subtle_gui.constants import Constants, MediaType
 
 logger = logging.getLogger(__name__)
 
@@ -140,6 +141,28 @@ def __init__(self, parent: QWidget) -> None:
         self.srtFrame = QGroupBox(self.tr("SubRip / SRT"), self)
         self.srtFrame.setLayout(self.srtForm)
 
+        # OCR Panel
+        # =========
+        self.ocrForm = QFormLayout()
+
+        self.ocrDialogLine = QComboBox(self)
+        self.ocrDialogLine.addItems(Constants.DIALOG_LINES)
+        self.ocrDialogLine.setCurrentText(CONFIG.getSetting("dialogLine"))
+        self.ocrDialogLine.currentTextChanged.connect(self._updateOcrDialogLine)
+
+        self.ocrSpace = QCheckBox(self.tr("Add space"), self)
+        self.ocrSpace.setChecked(CONFIG.getFlag("dialogSpace"))
+        self.ocrSpace.clicked.connect(self._updateOcrSpace)
+
+        self.ocrDialogOpt = QHBoxLayout()
+        self.ocrDialogOpt.addWidget(self.ocrDialogLine)
+        self.ocrDialogOpt.addWidget(self.ocrSpace)
+
+        self.ocrForm.addRow(self.tr("Dialogue Line"), self.ocrDialogOpt)
+
+        self.ocrFrame = QGroupBox(self.tr("OCR Options"), self)
+        self.ocrFrame.setLayout(self.ocrForm)
+
         # Layout
         # ======
 
@@ -150,6 +173,7 @@ def __init__(self, parent: QWidget) -> None:
 
         self.rightBox = QVBoxLayout()
         self.rightBox.addWidget(self.subsFrame)
+        self.rightBox.addWidget(self.ocrFrame)
         self.rightBox.addStretch(1)
 
         self.outerBox = QHBoxLayout()
@@ -223,6 +247,16 @@ def _updateTrackInfo(self) -> None:
             self.srtSaveDir.setText(str(folder))
             self.srtFileName.setText(".".join(bits))
 
+    @pyqtSlot()
+    def _updateOcrDialogLine(self) -> None:
+        """Update the dialog line setting."""
+        CONFIG.setSetting("dialogLine", self.ocrDialogLine.currentText())
+
+    @pyqtSlot()
+    def _updateOcrSpace(self) -> None:
+        """Update the dialog space setting."""
+        CONFIG.setFlag("dialogSpace", self.ocrSpace.isChecked())
+
     ##
     #  Internal Functions
     ##

subtle_gui/guimain.py

@@ -98,6 +98,8 @@ def __init__(self) -> None:
         SHARED.media.newMediaLoaded.connect(self.mediaView.processNewMediaLoaded)
         SHARED.media.newMediaLoaded.connect(self.subsView.processNewMediaLoaded)
         SHARED.media.newMediaLoaded.connect(self.toolsPanel.processNewMediaLoaded)
+        SHARED.media.newMediaLoaded.connect(self.imageViewer.processNewMediaLoaded)
+        SHARED.media.newMediaLoaded.connect(self.textEditor.processNewMediaLoaded)
         SHARED.media.newTrackSelected.connect(self.subsView.processNewTrackLoaded)
         SHARED.media.newTrackSelected.connect(self.toolsPanel.processNewTrackLoaded)
 

subtle_gui/ocr/base.py

@@ -26,6 +26,9 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
+from subtle_gui import CONFIG
+from subtle_gui.constants import Constants
+
 if TYPE_CHECKING:
     from PyQt6.QtGui import QImage
 
@@ -42,3 +45,22 @@ def __init__(self) -> None:
     def processImage(self, index: int, image: QImage, lang: list[str]) -> list[str]:
         """Process an image and return the recognized text."""
         raise NotImplementedError
+
+    def postProcessText(self, text: list[str]) -> list[str]:
+        """Run standard text post-processing tasks."""
+        dialogLine = CONFIG.getSetting("dialogLine")
+        dialogSpace = CONFIG.getFlag("dialogSpace")
+        result = []
+        for line in text:
+            if line.startswith(Constants.DIALOG_LINES):
+                post = dialogLine + (" " if dialogSpace else "") + line[1:].lstrip()
+            else:
+                post = line
+
+            if post != line:
+                logger.debug("Post Before: '%s'", line)
+                logger.debug("Post Result: '%s'", post)
+
+            result.append(post)
+
+        return result

subtle_gui/ocr/tesseract.py

@@ -28,14 +28,16 @@
 
 from typing import TYPE_CHECKING
 
+from PyQt6.QtGui import QImage
+
 from subtle_gui import CONFIG
 from subtle_gui.common import regexCleanup, simplified
 from subtle_gui.ocr.base import OCRBase
 
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from PyQt6.QtGui import QImage
+BINARY_THRESHOLD = 128
 
 logger = logging.getLogger(__name__)
 
@@ -66,6 +68,8 @@
         # Misinterpreted words
         (re.compile(r"\b(tt)\b", re.UNICODE), "it"),
         (re.compile(r"\b(fo)\b", re.UNICODE), "to"),
+        (re.compile(r"\b(lf)\b", re.UNICODE), "If"),
+        (re.compile(r"\b(l)\b", re.UNICODE), "I"),
         # Wrong capitalisation at the start of words
         (re.compile(r"(?<![.!?\)\]-])\s(K)now", re.UNICODE), "k"),
         (re.compile(r"(?<![.!?\)\]-])\s(I)t+", re.UNICODE), "i"),
@@ -74,6 +78,7 @@
         (re.compile(r"\b[D|d]id(nt)\b", re.UNICODE), "n't"),
         (re.compile(r"\b[T|t]hey(re)\b", re.UNICODE), "'re"),
         (re.compile(r"\b[Y|y]ou(ll)\b", re.UNICODE), "'ll"),
+        (re.compile(r"\b[W|w]ei(ll)\b", re.UNICODE), "'ll"),
         (re.compile(r"\b(l'll)\b", re.UNICODE), "I'll"),
     ],
 }
@@ -88,33 +93,56 @@ def __init__(self) -> None:
     def processImage(self, index: int, image: QImage, lang: list[str]) -> list[str]:
         """Perform OCR on a QImage."""
         tmpFile = CONFIG.dumpPath / f"{uuid.uuid4()!s}.png"
-        image.save(str(tmpFile), quality=100)
+        self._toMonochrome(image).save(str(tmpFile), quality=100)
         result = self._processText(self._callTesseract(tmpFile, lang), lang)
+        result = self.postProcessText(result)
         tmpFile.unlink(missing_ok=True)
         return result
 
     ##
     #  Internal Functions
     ##
 
+    def _toMonochrome(self, image: QImage, threshold: int = BINARY_THRESHOLD) -> QImage:
+        """Convert image to black text on white background for OCR."""
+        gray = image.convertToFormat(QImage.Format.Format_Grayscale8)
+        buf = gray.constBits()
+        buf.setsize(gray.sizeInBytes())
+        table = bytes(0 if lum > threshold else 255 for lum in range(256))
+        data = bytes(buf).translate(table)  # type: ignore
+        result = QImage(data, gray.width(), gray.height(), gray.bytesPerLine(), QImage.Format.Format_Grayscale8)
+        return result.copy()
+
     def _callTesseract(self, file: Path, lang: list[str]) -> str:
         """Call tesseract on an image file."""
+        out = self._runTesseract(file, lang, "6")
+        if not out.strip():
+            # PSM 6 assumes a block of text, which can cause tiny/sparse
+            # crops to be discarded as noise before they're even read.
+            # PSM 7 (single line, no layout analysis) can recover those,
+            # but mangles multi-line text, so it's only used as a fallback.
+            logger.debug("Tesseract returned no text with PSM 6, trying PSM 7")
+            out = self._runTesseract(file, lang, "7")
+        return out
+
+    def _runTesseract(self, file: Path, lang: list[str], psm: str) -> str:
+        """Run tesseract on an image file with a given page segmentation mode."""
         try:
-            cmd = ["tesseract", str(file), "-", "-l", "+".join(lang)]
+            cmd = [
+                "tesseract", str(file), "-", "-l", "+".join(lang),
+                "--oem", "1", "--psm", psm,
+            ]  # fmt: off
             if tessData := CONFIG.getSetting("tessData"):
                 cmd += ["--tessdata-dir", tessData]
-            p = subprocess.Popen(
-                ["tesseract", str(file), "-", "-l", "+".join(lang)],
-                stdout=subprocess.PIPE,
-            )
+            p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
             out, _ = p.communicate()
             return out.decode("utf-8")
         except Exception as e:
             logger.error("Failed to extract text with tesseract", exc_info=e)
         return ""
 
     def _processText(self, text: str, lang: list[str]) -> list[str]:
-        """Post-process text returned from tesseract."""
+        """Process text returned from tesseract."""
         temp = text.strip()
         for a, b in TXT_REPLACE.items():
             temp = temp.replace(a, b)