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

import logging

import cv2

# from arcaea_offline_ocr_device_creation_wizard.implements.wizard import Wizard
from arcaea_offline_ocr.device.v1.definition import DeviceV1
from arcaea_offline_ocr.device.v2.definition import DeviceV2
from arcaea_offline_ocr.sift_db import SIFTDatabase
from PySide6.QtCore import Qt, Slot
from PySide6.QtWidgets import QApplication, QFileDialog, QWidget

from ui.designer.tabs.tabOcr.tabOcr_Device_ui import Ui_TabOcr_Device
from ui.extends.components.ocrQueue import OcrQueueModel
from ui.extends.shared.language import LanguageChangeEventFilter
from ui.extends.shared.settings import (
    DEVICES_JSON_FILE,
    KNN_MODEL_FILE,
    SIFT_DATABASE_FILE,
    Settings,
)
from ui.extends.tabs.tabOcr.tabOcr_Device import (
    ScoreConverter,
    TabDeviceV2AutoRoisOcrRunnable,
    TabDeviceV2OcrRunnable,
)

logger = logging.getLogger(__name__)


class TabOcr_Device(Ui_TabOcr_Device, QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setupUi(self)
        self.openWizardButton.setEnabled(False)

        self.languageChangeEventFilter = LanguageChangeEventFilter(self)
        self.installEventFilter(self.languageChangeEventFilter)

        self.deviceFileSelector.filesSelected.connect(self.deviceFileSelected)
        self.knnModelSelector.filesSelected.connect(self.knnModelFileSelected)
        self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseFileSelected)

        self.deviceFileSelector.connectSettings(DEVICES_JSON_FILE)
        self.knnModelSelector.connectSettings(KNN_MODEL_FILE)
        self.siftDatabaseSelector.connectSettings(SIFT_DATABASE_FILE)

        settings = Settings()
        logger.info("Applying default settings...")
        self.deviceFileSelector.selectFile(settings.devicesJsonFile())
        self.tesseractFileSelector.selectFile(settings.tesseractPath())
        self.deviceComboBox.selectDevice(settings.deviceUuid())
        self.knnModelSelector.selectFile(settings.knnModelFile())
        self.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())

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

    @Slot()
    def on_openWizardButton_clicked(self):
        # wizard = Wizard(self)
        # wizard.open()
        pass

    @Slot()
    def on_deviceUseAutoFactorCheckBox_stateChanged(self):
        checkState = self.deviceUseAutoFactorCheckBox.checkState()
        if checkState == Qt.CheckState.Checked:
            self.deviceDependenciesStackedWidget.setCurrentIndex(1)
            self.deviceComboBox.setCurrentIndex(-1)
            self.deviceFileSelector.setEnabled(False)
            self.deviceComboBox.setEnabled(False)
        else:
            self.deviceFileSelector.setEnabled(True)
            self.deviceComboBox.setEnabled(True)

    @Slot()
    def on_deviceComboBox_currentIndexChanged(self):
        self.changeDeviceDepStackedWidget()

    def changeDeviceDepStackedWidget(self):
        device = self.deviceComboBox.currentData()
        if isinstance(device, (DeviceV1, DeviceV2)):
            self.deviceDependenciesStackedWidget.setCurrentIndex(device.version - 1)

    def deviceFileSelected(self):
        if selectedFiles := self.deviceFileSelector.selectedFiles():
            file = selectedFiles[0]
            self.deviceComboBox.loadDevicesJson(file)

    def knnModelFileSelected(self):
        if selectedFiles := self.knnModelSelector.selectedFiles():
            self.knnModel = cv2.ml.KNearest.load(selectedFiles[0])

    def siftDatabaseFileSelected(self):
        if selectedFiles := self.siftDatabaseSelector.selectedFiles():
            self.siftDatabase = SIFTDatabase(selectedFiles[0])

    @Slot()
    def on_ocr_addImageButton_clicked(self):
        files, _filter = QFileDialog.getOpenFileNames(
            self, None, "", "Image Files (*.png *.jpg *.jpeg *.bmp *.webp);;*"
        )
        for file in files:
            self.ocrQueueModel.addItem(file)
            QApplication.processEvents()
        self.ocrQueue.resizeTableView()

    @Slot()
    def on_ocr_startButton_clicked(self):
        for row in range(self.ocrQueueModel.rowCount()):
            index = self.ocrQueueModel.index(row, 0)
            imagePath = index.data(OcrQueueModel.ImagePathRole)
            if self.deviceUseAutoFactorCheckBox.checkState() == Qt.CheckState.Checked:
                runnable = TabDeviceV2AutoRoisOcrRunnable(
                    imagePath,
                    self.knnModel,
                    self.siftDatabase,
                )
            else:
                runnable = TabDeviceV2OcrRunnable(
                    imagePath,
                    self.deviceComboBox.currentData(),
                    self.knnModel,
                    self.siftDatabase,
                )
            self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
            self.ocrQueueModel.setData(
                index,
                ScoreConverter.deviceV2,
                OcrQueueModel.ProcessOcrResultFuncRole,
            )
        self.ocrQueueModel.startQueue()

    @Slot()
    def on_ocr_removeAllButton_clicked(self):
        self.ocrQueueModel.clear()