wip: DeviceV2 ocr support

This commit is contained in:
2023-08-15 15:33:11 +08:00
parent 0126ce6b9a
commit 044a8eb07a
5 changed files with 335 additions and 168 deletions

View File

@ -1,14 +1,17 @@
import cv2
import pytesseract
# from arcaea_offline_ocr_device_creation_wizard.implements.wizard import Wizard
from PySide6.QtCore import QModelIndex, Qt, Slot
from PySide6.QtGui import QColor
from PySide6.QtWidgets import QFileDialog, QHeaderView, QWidget
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 Slot
from PySide6.QtWidgets import QFileDialog, QWidget
from ui.designer.tabs.tabOcr_ui import Ui_TabOcr
from ui.extends.components.ocrQueue import OcrQueueModel
from ui.extends.settings import Settings
from ui.extends.tabs.tabOcr import TabDeviceV2OcrRunnable, ScoreInsertConverter
from ui.extends.tabs.tabOcr import ScoreInsertConverter, TabDeviceV2OcrRunnable
class TabOcr(Ui_TabOcr, QWidget):
@ -17,10 +20,16 @@ class TabOcr(Ui_TabOcr, QWidget):
self.setupUi(self)
self.openWizardButton.setEnabled(False)
self.deviceComboBox.currentIndexChanged.connect(
self.changeDeviceDepStackedWidget
)
self.deviceFileSelector.filesSelected.connect(self.deviceFileSelected)
self.knnModelSelector.filesSelected.connect(self.knnModelFileSelected)
self.tesseractFileSelector.filesSelected.connect(
self.tesseractFileSelectorFilesSelected
)
self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseFileSelected)
settings = Settings()
self.deviceFileSelector.selectFile(settings.devicesJsonFile())
@ -37,17 +46,32 @@ class TabOcr(Ui_TabOcr, QWidget):
# wizard.open()
pass
def changeDeviceDepStackedWidget(self):
device = self.deviceComboBox.currentData()
if isinstance(device, (DeviceV1, DeviceV2)):
self.deviceDependenciesStackedWidget.setCurrentIndex(device.version - 1)
def deviceFileSelected(self):
selectedFiles = self.deviceFileSelector.selectedFiles()
if selectedFiles:
file = selectedFiles[0]
self.deviceComboBox.loadDevicesJson(file)
def knnModelFileSelected(self):
selectedFiles = self.knnModelSelector.selectedFiles()
if selectedFiles:
self.knnModel = cv2.ml.KNearest.load(selectedFiles[0])
def tesseractFileSelectorFilesSelected(self):
selectedFiles = self.tesseractFileSelector.selectedFiles()
if selectedFiles:
pytesseract.pytesseract.tesseract_cmd = selectedFiles[0]
def siftDatabaseFileSelected(self):
selectedFiles = self.siftDatabaseSelector.selectedFiles()
if selectedFiles:
self.siftDatabase = SIFTDatabase(selectedFiles[0])
@Slot()
def on_ocr_addImageButton_clicked(self):
files, _filter = QFileDialog.getOpenFileNames(
@ -63,7 +87,10 @@ class TabOcr(Ui_TabOcr, QWidget):
index = self.ocrQueueModel.index(row, 0)
imagePath = index.data(OcrQueueModel.ImagePathRole)
runnable = TabDeviceV2OcrRunnable(
imagePath, self.deviceComboBox.currentData(), self.knn, self.siftDb
imagePath,
self.deviceComboBox.currentData(),
self.knnModel,
self.siftDatabase,
)
self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
self.ocrQueueModel.setData(