diff --git a/.gitignore b/.gitignore index 7c5a3c0..ed94f8d 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,9 @@ .ruff_cache/ *.pyc +# Other +.config + # Build __pycache__/ dist/ diff --git a/subtle_gui/config.py b/subtle_gui/config.py index 186ef3e..8a9479c 100644 --- a/subtle_gui/config.py +++ b/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): diff --git a/subtle_gui/constants.py b/subtle_gui/constants.py index 271c00a..83cf2a7 100644 --- a/subtle_gui/constants.py +++ b/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") diff --git a/subtle_gui/gui/imageviewer.py b/subtle_gui/gui/imageviewer.py index 7486d45..4446e04 100644 --- a/subtle_gui/gui/imageviewer.py +++ b/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.""" diff --git a/subtle_gui/gui/texteditor.py b/subtle_gui/gui/texteditor.py index 6a9733a..ea24ee3 100644 --- a/subtle_gui/gui/texteditor.py +++ b/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.""" diff --git a/subtle_gui/gui/toolspanel.py b/subtle_gui/gui/toolspanel.py index 15f1fd9..66a3c09 100644 --- a/subtle_gui/gui/toolspanel.py +++ b/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 ## diff --git a/subtle_gui/guimain.py b/subtle_gui/guimain.py index 2acb0c7..4f0f99c 100644 --- a/subtle_gui/guimain.py +++ b/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) diff --git a/subtle_gui/ocr/base.py b/subtle_gui/ocr/base.py index f9e6562..1536765 100644 --- a/subtle_gui/ocr/base.py +++ b/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 diff --git a/subtle_gui/ocr/tesseract.py b/subtle_gui/ocr/tesseract.py index df48082..a272f00 100644 --- a/subtle_gui/ocr/tesseract.py +++ b/subtle_gui/ocr/tesseract.py @@ -28,6 +28,8 @@ 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 @@ -35,7 +37,7 @@ 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"(? 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 @@ -97,16 +103,38 @@ def processImage(self, index: int, image: QImage, lang: list[str]) -> list[str]: # 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: @@ -114,7 +142,7 @@ def _callTesseract(self, file: Path, lang: list[str]) -> str: 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)