arcaea-offline-pyside-ui/ui/implements/tabs/tabOcr/tabOcr_B30.py

import logging
from pathlib import Path

import cv2
from arcaea_offline_ocr.b30.chieri.v4.ocr import ChieriBotV4Ocr
from arcaea_offline_ocr.sift_db import SIFTDatabase
from arcaea_offline_ocr.utils import imread_unicode

# from paddleocr import PaddleOCR
from PySide6.QtCore import Signal, Slot
from PySide6.QtWidgets import QWidget

from ui.designer.tabs.tabOcr.tabOcr_B30_ui import Ui_TabOcr_B30
from ui.extends.components.ocrQueue import OcrQueueModel
from ui.extends.shared.cv2_utils import cv2BgrMatToQImage, qImageToCvMatBgr
from ui.extends.shared.settings import Settings
from ui.extends.tabs.tabOcr.tabOcr_B30 import ChieriV4OcrRunnable, b30ResultToScore

logger = logging.getLogger(__name__)


class TabOcr_B30(Ui_TabOcr_B30, QWidget):
    tryPrepareOcr = Signal()

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setupUi(self)

        self.b30TypeComboBox.addItem("ChieriV4", "chieri_v4")
        self.b30TypeComboBox.setCurrentIndex(0)
        self.b30TypeComboBox.setEnabled(False)

        # self.paddleFolderSelector.setMode(
        #     self.paddleFolderSelector.getExistingDirectory
        # )

        self.imageSelector.filesSelected.connect(self.imageSelected)
        self.knnModelSelector.filesSelected.connect(self.knnModelSelected)
        self.b30KnnModelSelector.filesSelected.connect(self.b30KnnModelSelected)
        # self.paddleFolderSelector.filesSelected.connect(self.paddleFolderSelected)
        self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseSelected)

        self.imagePath = None  # for checking only
        self.img = None
        self.paddleFolder = None
        self.paddle = None
        self.knnModel = None
        self.b30KnnModel = None
        self.siftDatabase = None

        self.ocr = None

        self.tryPrepareOcr.connect(self.prepareOcr)

        settings = Settings()
        logger.info("Applying default settings...")
        self.knnModelSelector.selectFile(settings.knnModelFile())
        self.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())

        self.ocrQueueModel = OcrQueueModel(self)
        self.ocrQueue.setModel(self.ocrQueueModel)

    def imageSelected(self):
        selectedFiles = self.imageSelector.selectedFiles()
        if selectedFiles:
            imagePath = selectedFiles[0]
            self.imagePath = imagePath
            self.img = imread_unicode(imagePath)
            self.tryPrepareOcr.emit()

    def knnModelSelected(self):
        selectedFiles = self.knnModelSelector.selectedFiles()
        if selectedFiles:
            knnModelPath = selectedFiles[0]
            self.knnModel = cv2.ml.KNearest.load(knnModelPath)
            self.tryPrepareOcr.emit()

    def b30KnnModelSelected(self):
        selectedFiles = self.b30KnnModelSelector.selectedFiles()
        if selectedFiles:
            b30KnnModelPath = selectedFiles[0]
            self.b30KnnModel = cv2.ml.KNearest.load(b30KnnModelPath)
            self.tryPrepareOcr.emit()

    def siftDatabaseSelected(self):
        selectedFiles = self.siftDatabaseSelector.selectedFiles()
        if selectedFiles:
            siftDatabasePath = selectedFiles[0]
            self.siftDatabase = SIFTDatabase(siftDatabasePath)
            self.tryPrepareOcr.emit()

    def paddleFolderSelected(self):
        selectedFiles = self.paddleFolderSelector.selectedFiles()
        if selectedFiles:
            self.paddleFolder = selectedFiles[0]
            self.initPaddle()
            self.tryPrepareOcr.emit()

    def initPaddle(self):
        paddleFolder = Path(self.paddleFolder)
        paddleDetFolder = paddleFolder / "det"
        paddleClsFolder = paddleFolder / "cls"
        paddleRecFolder = paddleFolder / "rec"

        if not (paddleDetFolder.exists() and paddleRecFolder.exists()):
            logger.warning("paddleocr folder incomplete, aborting.")
            return

        self.paddle = PaddleOCR(
            show_log=False,
            use_angle_cls=False,
            det_model_dir=str(paddleDetFolder),
            cls_model_dir=str(paddleClsFolder),
            rec_model_dir=str(paddleRecFolder),
        )

    def prepareOcr(self):
        b30Type = self.b30TypeComboBox.currentData()
        if not b30Type:
            return

        if b30Type == "chieri_v4":
            if (
                not self.imagePath
                or not self.knnModel
                or not self.b30KnnModel
                or not self.siftDatabase
            ):
                return

            self.ocrQueueModel.clear()

            ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.siftDatabase)
            ocr.set_factor(self.img)
            self.ocr = ocr

            roi = ocr.rois
            for component in roi.components(self.img):
                qImage = cv2BgrMatToQImage(component.copy())
                self.ocrQueueModel.addItem(qImage)
        self.ocrQueue.resizeTableView()

    @Slot()
    def on_ocr_startButton_clicked(self):
        if (
            not self.imagePath
            or not self.knnModel
            or not self.b30KnnModel
            or not self.siftDatabase
        ):
            return

        for row in range(self.ocrQueueModel.rowCount()):
            index = self.ocrQueueModel.index(row, 0)
            qImage = index.data(OcrQueueModel.ImageQImageRole)
            cv2Mat = qImageToCvMatBgr(qImage)
            runnable = ChieriV4OcrRunnable(self.ocr, cv2Mat)
            self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
            self.ocrQueueModel.setData(
                index,
                b30ResultToScore,
                OcrQueueModel.ProcessOcrResultFuncRole,
            )
        self.ocrQueueModel.startQueue()