engines/trocr_engine.py

"""
TrOCR Engine Plugin

Wraps the TrOCR HTR inference system as a plugin for the unified GUI.
"""

from pathlib import Path
from typing import Dict, Any, Optional
import numpy as np

from htr_engine_base import HTREngine, TranscriptionResult

try:
    from PyQt6.QtWidgets import (
        QWidget, QVBoxLayout, QHBoxLayout, QLabel, QComboBox,
        QPushButton, QSpinBox, QCheckBox, QLineEdit, QFileDialog,
        QGroupBox, QRadioButton, QButtonGroup
    )
    from PyQt6.QtCore import Qt
    PYQT_AVAILABLE = True
except ImportError:
    PYQT_AVAILABLE = False
    QWidget = object

try:
    from inference_page import TrOCRInference
    TROCR_AVAILABLE = True
except ImportError:
    TROCR_AVAILABLE = False


class TrOCREngine(HTREngine):
    """TrOCR HTR engine plugin."""

    def __init__(self):
        self.model: Optional[TrOCRInference] = None
        self._config_widget: Optional[QWidget] = None
        self._current_model_path: Optional[str] = None  # Store loaded model path

        # Widget references (set when config widget is created)
        self._model_source_combo: Optional[QComboBox] = None
        self._local_model_combo: Optional[QComboBox] = None
        self._hf_preset_combo: Optional[QComboBox] = None
        self._hf_model_edit: Optional[QLineEdit] = None
        self._beam_spin: Optional[QSpinBox] = None
        self._normalize_check: Optional[QCheckBox] = None
        self._flip_rtl_check: Optional[QCheckBox] = None

    def get_name(self) -> str:
        return "TrOCR"

    def get_description(self) -> str:
        return "Transformer-based OCR optimized for handwritten manuscripts"

    def is_available(self) -> bool:
        return TROCR_AVAILABLE

    def get_unavailable_reason(self) -> str:
        if not TROCR_AVAILABLE:
            return "TrOCR inference module not available. Check that inference_page.py exists."
        return ""

    def get_config_widget(self) -> QWidget:
        """Create TrOCR configuration panel."""
        if self._config_widget is not None:
            return self._config_widget

        widget = QWidget()
        layout = QVBoxLayout()

        # Model source selection
        source_group = QGroupBox("Model Source")
        source_layout = QVBoxLayout()

        self._model_source_combo = QComboBox()
        self._model_source_combo.addItems(["Local Models", "HuggingFace Hub"])
        self._model_source_combo.currentTextChanged.connect(self._on_model_source_changed)
        source_layout.addWidget(self._model_source_combo)

        source_group.setLayout(source_layout)
        layout.addWidget(source_group)

        # Local models group
        self._local_group = QGroupBox("Local Model")
        local_layout = QVBoxLayout()

        self._local_model_combo = QComboBox()
        self._populate_local_models()
        local_layout.addWidget(QLabel("Model:"))
        local_layout.addWidget(self._local_model_combo)

        browse_btn = QPushButton("Browse...")
        browse_btn.clicked.connect(self._browse_local_model)
        local_layout.addWidget(browse_btn)

        self._local_group.setLayout(local_layout)
        layout.addWidget(self._local_group)

        # HuggingFace models group
        self._hf_group = QGroupBox("HuggingFace Model")
        hf_layout = QVBoxLayout()

        # Predefined models dropdown
        hf_layout.addWidget(QLabel("Preset Models:"))
        self._hf_preset_combo = QComboBox()
        self._hf_preset_combo.addItems([
            "Custom (enter below)",
            "cyrillic-trocr/trocr-handwritten-cyrillic",
            "kazars24/trocr-base-handwritten-ru",
            "dh-unibe/trocr-kurrent",
            "dh-unibe/trocr-kurrent-XVI-XVII",
            "microsoft/trocr-base-handwritten",
            "microsoft/trocr-large-handwritten",
            "microsoft/trocr-base-printed",
            "microsoft/trocr-large-printed"
        ])
        self._hf_preset_combo.currentTextChanged.connect(self._on_preset_changed)
        hf_layout.addWidget(self._hf_preset_combo)

        # Custom model ID entry
        hf_layout.addWidget(QLabel("Model ID:"))
        self._hf_model_edit = QLineEdit()
        self._hf_model_edit.setPlaceholderText("e.g., cyrillic-trocr/trocr-handwritten-cyrillic")
        self._hf_model_edit.setText("cyrillic-trocr/trocr-handwritten-cyrillic")
        hf_layout.addWidget(self._hf_model_edit)

        self._hf_group.setLayout(hf_layout)
        self._hf_group.setVisible(False)  # Hidden by default
        layout.addWidget(self._hf_group)

        # Inference settings
        settings_group = QGroupBox("Inference Settings")
        settings_layout = QVBoxLayout()

        # Beam search
        beam_layout = QHBoxLayout()
        beam_layout.addWidget(QLabel("Beam Search:"))
        self._beam_spin = QSpinBox()
        self._beam_spin.setRange(1, 10)
        self._beam_spin.setValue(4)
        self._beam_spin.setToolTip("Higher = better quality but slower (1 = greedy, 4 = recommended)")
        beam_layout.addWidget(self._beam_spin)
        beam_layout.addStretch()
        settings_layout.addLayout(beam_layout)

        # Normalize background
        self._normalize_check = QCheckBox("Normalize Background")
        self._normalize_check.setToolTip("Apply CLAHE normalization (use if model was trained with it)")
        settings_layout.addWidget(self._normalize_check)

        self._flip_rtl_check = QCheckBox("RTL manuscript (flip line images)")
        self._flip_rtl_check.setToolTip("Flip line images horizontally for RTL scripts (Ottoman, Arabic, Hebrew)")
        settings_layout.addWidget(self._flip_rtl_check)

        settings_group.setLayout(settings_layout)
        layout.addWidget(settings_group)

        layout.addStretch()
        widget.setLayout(layout)

        self._config_widget = widget
        return widget

    def _populate_local_models(self):
        """Scan models directory for TrOCR checkpoints."""
        if self._local_model_combo is None:
            return

        self._local_model_combo.clear()

        models_dir = Path("models")
        if not models_dir.exists():
            self._local_model_combo.addItem("No models found")
            return

        # Find checkpoint directories
        checkpoints = []
        for model_dir in models_dir.iterdir():
            if not model_dir.is_dir():
                continue

            # Check if it's a TrOCR model (has pytorch_model.bin or model.safetensors)
            if (model_dir / "pytorch_model.bin").exists() or (model_dir / "model.safetensors").exists():
                checkpoints.append(str(model_dir))

            # Also check for checkpoint subdirectories
            for subdir in model_dir.glob("checkpoint-*"):
                if subdir.is_dir():
                    if (subdir / "pytorch_model.bin").exists() or (subdir / "model.safetensors").exists():
                        checkpoints.append(str(subdir))

        if checkpoints:
            self._local_model_combo.addItems(sorted(checkpoints))
        else:
            self._local_model_combo.addItem("No local models found")

    def _browse_local_model(self):
        """Open file dialog to select model directory."""
        directory = QFileDialog.getExistingDirectory(
            self._config_widget,
            "Select Model Directory",
            "models"
        )

        if directory:
            # Add to combo box if not already present
            if self._local_model_combo.findText(directory) == -1:
                self._local_model_combo.addItem(directory)
            self._local_model_combo.setCurrentText(directory)

    def _on_model_source_changed(self, source: str):
        """Toggle between local and HuggingFace model selection."""
        is_local = (source == "Local Models")
        self._local_group.setVisible(is_local)
        self._hf_group.setVisible(not is_local)

    def _on_preset_changed(self, preset: str):
        """Update model ID when preset is selected."""
        if preset != "Custom (enter below)":
            self._hf_model_edit.setText(preset)

    def get_config(self) -> Dict[str, Any]:
        """Extract configuration from widget controls."""
        if self._config_widget is None:
            return {}

        is_local = (self._model_source_combo.currentText() == "Local Models")

        return {
            "model_source": "local" if is_local else "huggingface",
            "model_path": self._local_model_combo.currentText() if is_local else self._hf_model_edit.text(),
            "beam_search": self._beam_spin.value(),
            "normalize_background": self._normalize_check.isChecked(),
            "flip_rtl": self._flip_rtl_check.isChecked(),
        }

    def set_config(self, config: Dict[str, Any]):
        """Restore configuration to widget controls."""
        if self._config_widget is None:
            return

        model_source = config.get("model_source", "local")
        self._model_source_combo.setCurrentText("Local Models" if model_source == "local" else "HuggingFace Hub")

        model_path = config.get("model_path", "")
        if model_source == "local":
            idx = self._local_model_combo.findText(model_path)
            if idx >= 0:
                self._local_model_combo.setCurrentIndex(idx)
        else:
            self._hf_model_edit.setText(model_path)

        self._beam_spin.setValue(config.get("beam_search", 4))
        self._normalize_check.setChecked(config.get("normalize_background", False))
        if self._flip_rtl_check:
            self._flip_rtl_check.setChecked(config.get("flip_rtl", False))

    def load_model(self, config: Dict[str, Any]) -> bool:
        """Load TrOCR model."""
        try:
            model_path = config.get("model_path", "")
            if not model_path or model_path == "No local models found" or model_path == "No models found":
                return False

            normalize = config.get("normalize_background", False)
            flip_rtl = config.get("flip_rtl", False)
            model_source = config.get("model_source", "local")
            is_hf = (model_source == "huggingface")

            self.model = TrOCRInference(
                model_path=model_path,
                normalize_bg=normalize,
                flip_rtl=flip_rtl,
                is_huggingface=is_hf
            )

            # Store the model path for metadata
            self._current_model_path = model_path

            return True

        except Exception as e:
            print(f"Error loading TrOCR model: {e}")
            self.model = None
            self._current_model_path = None
            return False

    def unload_model(self):
        """Unload model from memory."""
        if self.model is not None:
            del self.model
            self.model = None
            self._current_model_path = None

            # Free GPU memory
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def is_model_loaded(self) -> bool:
        """Check if model is loaded."""
        return self.model is not None

    def transcribe_line(self, image: np.ndarray, config: Optional[Dict[str, Any]] = None) -> TranscriptionResult:
        """Transcribe a line image with TrOCR."""
        if self.model is None:
            return TranscriptionResult(text="[Model not loaded]", confidence=0.0)

        # Use provided config or fall back to current config
        if config is None:
            config = self.get_config()

        beam_search = config.get("beam_search", 4)

        try:
            # TrOCRInference expects PIL Image, convert from numpy
            from PIL import Image
            pil_image = Image.fromarray(image)

            text = self.model.transcribe_line(pil_image, num_beams=beam_search)

            # Build metadata with model information
            metadata = {
                "beam_search": beam_search,
                "engine": "TrOCR"
            }

            # Add model path if available
            if self._current_model_path:
                metadata["model"] = self._current_model_path

            return TranscriptionResult(
                text=text,
                confidence=1.0,  # TrOCR doesn't provide confidence scores
                metadata=metadata
            )

        except Exception as e:
            print(f"Error in TrOCR transcription: {e}")
            return TranscriptionResult(text=f"[Error: {e}]", confidence=0.0)

    def get_capabilities(self) -> Dict[str, bool]:
        """TrOCR capabilities."""
        return {
            "batch_processing": False,  # Could be implemented in future
            "confidence_scores": False,  # TrOCR doesn't provide per-token confidence
            "beam_search": True,
            "language_model": False,  # Implicit in decoder, not explicit LM
            "preprocessing": True,  # Has built-in normalization
        }