diff --git a/ui/designer/tabs/tabOcr/tabOcr_B30.ui b/ui/designer/tabs/tabOcr/tabOcr_B30.ui index 5965a79..6a2df22 100644 --- a/ui/designer/tabs/tabOcr/tabOcr_B30.ui +++ b/ui/designer/tabs/tabOcr/tabOcr_B30.ui @@ -13,7 +13,7 @@ TabOcr_B30 - + @@ -27,60 +27,87 @@ - - - - - knnModelSelector.title - - - - - - - - - - - - b30KnnModelSelector.title - - - - - - - - - - - - - - - - phashDatabaseSelector.title - - - - - - - - - - - - imageSelector.title - - - - - - - - - + + + dependencies.title + + + + + + dependencies.knnModel + + + Qt::AlignRight|Qt::AlignTrailing|Qt::AlignVCenter + + + + + + + ... + + + + + + + Qt::Vertical + + + + + + + dependencies.b30KnnModel + + + Qt::AlignRight|Qt::AlignTrailing|Qt::AlignVCenter + + + + + + + ... + + + + + + + Qt::Vertical + + + + + + + dependencies.phashDatabase + + + Qt::AlignRight|Qt::AlignTrailing|Qt::AlignVCenter + + + + + + + ... + + + + + + + + + + + + + + diff --git a/ui/designer/tabs/tabOcr/tabOcr_B30_ui.py b/ui/designer/tabs/tabOcr/tabOcr_B30_ui.py index be0f290..ba0343c 100644 --- a/ui/designer/tabs/tabOcr/tabOcr_B30_ui.py +++ b/ui/designer/tabs/tabOcr/tabOcr_B30_ui.py @@ -15,8 +15,9 @@ from PySide6.QtGui import (QBrush, QColor, QConicalGradient, QCursor, QFont, QFontDatabase, QGradient, QIcon, QImage, QKeySequence, QLinearGradient, QPainter, QPalette, QPixmap, QRadialGradient, QTransform) -from PySide6.QtWidgets import (QApplication, QComboBox, QGroupBox, QHBoxLayout, - QSizePolicy, QVBoxLayout, QWidget) +from PySide6.QtWidgets import (QApplication, QComboBox, QFrame, QGridLayout, + QGroupBox, QLabel, QSizePolicy, QVBoxLayout, + QWidget) from ui.implements.components.fileSelector import FileSelector from ui.implements.components.ocrQueue import OcrQueue @@ -27,8 +28,8 @@ class Ui_TabOcr_B30(object): TabOcr_B30.setObjectName(u"TabOcr_B30") TabOcr_B30.resize(555, 461) TabOcr_B30.setWindowTitle(u"TabOcr_B30") - self.verticalLayout_3 = QVBoxLayout(TabOcr_B30) - self.verticalLayout_3.setObjectName(u"verticalLayout_3") + self.verticalLayout_2 = QVBoxLayout(TabOcr_B30) + self.verticalLayout_2.setObjectName(u"verticalLayout_2") self.groupBox = QGroupBox(TabOcr_B30) self.groupBox.setObjectName(u"groupBox") self.verticalLayout = QVBoxLayout(self.groupBox) @@ -39,65 +40,80 @@ class Ui_TabOcr_B30(object): self.verticalLayout.addWidget(self.b30TypeComboBox) - self.verticalLayout_3.addWidget(self.groupBox) + self.verticalLayout_2.addWidget(self.groupBox) - self.horizontalLayout = QHBoxLayout() - self.horizontalLayout.setObjectName(u"horizontalLayout") - self.groupBox_3 = QGroupBox(TabOcr_B30) - self.groupBox_3.setObjectName(u"groupBox_3") - self.verticalLayout_4 = QVBoxLayout(self.groupBox_3) - self.verticalLayout_4.setObjectName(u"verticalLayout_4") - self.knnModelSelector = FileSelector(self.groupBox_3) - self.knnModelSelector.setObjectName(u"knnModelSelector") + self.groupBox_6 = QGroupBox(TabOcr_B30) + self.groupBox_6.setObjectName(u"groupBox_6") + self.gridLayout = QGridLayout(self.groupBox_6) + self.gridLayout.setObjectName(u"gridLayout") + self.label = QLabel(self.groupBox_6) + self.label.setObjectName(u"label") + self.label.setAlignment(Qt.AlignRight|Qt.AlignTrailing|Qt.AlignVCenter) - self.verticalLayout_4.addWidget(self.knnModelSelector) + self.gridLayout.addWidget(self.label, 0, 0, 1, 1) + self.dependencies_knnModelStatusLabel = QLabel(self.groupBox_6) + self.dependencies_knnModelStatusLabel.setObjectName(u"dependencies_knnModelStatusLabel") + self.dependencies_knnModelStatusLabel.setText(u"...") - self.horizontalLayout.addWidget(self.groupBox_3) + self.gridLayout.addWidget(self.dependencies_knnModelStatusLabel, 0, 2, 1, 1) - self.groupBox_5 = QGroupBox(TabOcr_B30) - self.groupBox_5.setObjectName(u"groupBox_5") - self.verticalLayout_6 = QVBoxLayout(self.groupBox_5) - self.verticalLayout_6.setObjectName(u"verticalLayout_6") - self.b30KnnModelSelector = FileSelector(self.groupBox_5) - self.b30KnnModelSelector.setObjectName(u"b30KnnModelSelector") + self.line_2 = QFrame(self.groupBox_6) + self.line_2.setObjectName(u"line_2") + self.line_2.setFrameShape(QFrame.VLine) + self.line_2.setFrameShadow(QFrame.Sunken) - self.verticalLayout_6.addWidget(self.b30KnnModelSelector) + self.gridLayout.addWidget(self.line_2, 0, 3, 3, 1) + self.label_2 = QLabel(self.groupBox_6) + self.label_2.setObjectName(u"label_2") + self.label_2.setAlignment(Qt.AlignRight|Qt.AlignTrailing|Qt.AlignVCenter) - self.horizontalLayout.addWidget(self.groupBox_5) + self.gridLayout.addWidget(self.label_2, 1, 0, 1, 1) + self.dependencies_phashDatabaseStatusLabel = QLabel(self.groupBox_6) + self.dependencies_phashDatabaseStatusLabel.setObjectName(u"dependencies_phashDatabaseStatusLabel") + self.dependencies_phashDatabaseStatusLabel.setText(u"...") - self.verticalLayout_3.addLayout(self.horizontalLayout) + self.gridLayout.addWidget(self.dependencies_phashDatabaseStatusLabel, 2, 2, 1, 1) - self.horizontalLayout_3 = QHBoxLayout() - self.horizontalLayout_3.setObjectName(u"horizontalLayout_3") - self.groupBox_4 = QGroupBox(TabOcr_B30) - self.groupBox_4.setObjectName(u"groupBox_4") - self.verticalLayout_5 = QVBoxLayout(self.groupBox_4) - self.verticalLayout_5.setObjectName(u"verticalLayout_5") - self.phashDatabaseSelector = FileSelector(self.groupBox_4) - self.phashDatabaseSelector.setObjectName(u"phashDatabaseSelector") + self.line = QFrame(self.groupBox_6) + self.line.setObjectName(u"line") + self.line.setFrameShape(QFrame.VLine) + self.line.setFrameShadow(QFrame.Sunken) - self.verticalLayout_5.addWidget(self.phashDatabaseSelector) + self.gridLayout.addWidget(self.line, 0, 1, 3, 1) + self.label_3 = QLabel(self.groupBox_6) + self.label_3.setObjectName(u"label_3") + self.label_3.setAlignment(Qt.AlignRight|Qt.AlignTrailing|Qt.AlignVCenter) - self.horizontalLayout_3.addWidget(self.groupBox_4) + self.gridLayout.addWidget(self.label_3, 2, 0, 1, 1) - self.groupBox_2 = QGroupBox(TabOcr_B30) - self.groupBox_2.setObjectName(u"groupBox_2") - self.verticalLayout_2 = QVBoxLayout(self.groupBox_2) - self.verticalLayout_2.setObjectName(u"verticalLayout_2") - self.imageSelector = FileSelector(self.groupBox_2) - self.imageSelector.setObjectName(u"imageSelector") + self.dependencies_b30KnnModelStatusLabel = QLabel(self.groupBox_6) + self.dependencies_b30KnnModelStatusLabel.setObjectName(u"dependencies_b30KnnModelStatusLabel") + self.dependencies_b30KnnModelStatusLabel.setText(u"...") - self.verticalLayout_2.addWidget(self.imageSelector) + self.gridLayout.addWidget(self.dependencies_b30KnnModelStatusLabel, 1, 2, 1, 1) + self.dependencies_knnModelSelector = FileSelector(self.groupBox_6) + self.dependencies_knnModelSelector.setObjectName(u"dependencies_knnModelSelector") - self.horizontalLayout_3.addWidget(self.groupBox_2) + self.gridLayout.addWidget(self.dependencies_knnModelSelector, 0, 4, 1, 1) + self.dependencies_b30KnnModelSelector = FileSelector(self.groupBox_6) + self.dependencies_b30KnnModelSelector.setObjectName(u"dependencies_b30KnnModelSelector") - self.verticalLayout_3.addLayout(self.horizontalLayout_3) + self.gridLayout.addWidget(self.dependencies_b30KnnModelSelector, 1, 4, 1, 1) + + self.dependencies_phashDatabaseSelector = FileSelector(self.groupBox_6) + self.dependencies_phashDatabaseSelector.setObjectName(u"dependencies_phashDatabaseSelector") + + self.gridLayout.addWidget(self.dependencies_phashDatabaseSelector, 2, 4, 1, 1) + + self.gridLayout.setColumnStretch(4, 1) + + self.verticalLayout_2.addWidget(self.groupBox_6) self.ocrQueue = OcrQueue(TabOcr_B30) self.ocrQueue.setObjectName(u"ocrQueue") @@ -107,7 +123,7 @@ class Ui_TabOcr_B30(object): sizePolicy.setHeightForWidth(self.ocrQueue.sizePolicy().hasHeightForWidth()) self.ocrQueue.setSizePolicy(sizePolicy) - self.verticalLayout_3.addWidget(self.ocrQueue) + self.verticalLayout_2.addWidget(self.ocrQueue) self.retranslateUi(TabOcr_B30) @@ -117,10 +133,10 @@ class Ui_TabOcr_B30(object): def retranslateUi(self, TabOcr_B30): self.groupBox.setTitle(QCoreApplication.translate("TabOcr_B30", u"b30type", None)) - self.groupBox_3.setTitle(QCoreApplication.translate("TabOcr_B30", u"knnModelSelector.title", None)) - self.groupBox_5.setTitle(QCoreApplication.translate("TabOcr_B30", u"b30KnnModelSelector.title", None)) - self.groupBox_4.setTitle(QCoreApplication.translate("TabOcr_B30", u"phashDatabaseSelector.title", None)) - self.groupBox_2.setTitle(QCoreApplication.translate("TabOcr_B30", u"imageSelector.title", None)) + self.groupBox_6.setTitle(QCoreApplication.translate("TabOcr_B30", u"dependencies.title", None)) + self.label.setText(QCoreApplication.translate("TabOcr_B30", u"dependencies.knnModel", None)) + self.label_2.setText(QCoreApplication.translate("TabOcr_B30", u"dependencies.b30KnnModel", None)) + self.label_3.setText(QCoreApplication.translate("TabOcr_B30", u"dependencies.phashDatabase", None)) pass # retranslateUi diff --git a/ui/extends/ocr/dependencies.py b/ui/extends/ocr/dependencies.py new file mode 100644 index 0000000..6fffbbb --- /dev/null +++ b/ui/extends/ocr/dependencies.py @@ -0,0 +1,28 @@ +import cv2 +from arcaea_offline_ocr.phash_db import ImagePhashDatabase + + +def getCv2StatModelStatusText(model: cv2.ml.StatModel): + if not isinstance(model, cv2.ml.StatModel): + return 'ERROR' + + varCount = model.getVarCount() + if varCount != 81: + return f'WARN, varCount {varCount}' + else: + return f'OK, varCount {varCount}' + + +def getPhashDatabaseStatusText(db: ImagePhashDatabase): + if not isinstance(db, ImagePhashDatabase): + return 'ERROR' + + jacketCount = len(db.jacket_hashes) + partnerIconCount = len(db.partner_icon_hashes) + + statusText = f"J{jacketCount} PI{partnerIconCount}" + + if partnerIconCount <= 0: + return f'WARN, {statusText}' + else: + return f'OK, {statusText}' diff --git a/ui/extends/shared/cv2_utils.py b/ui/extends/shared/cv2_utils.py deleted file mode 100644 index c587a4d..0000000 --- a/ui/extends/shared/cv2_utils.py +++ /dev/null @@ -1,28 +0,0 @@ -import cv2 -import numpy as np -from PySide6.QtGui import QImage - - -def cv2BgrMatToQImage(mat) -> QImage: - arr = np.ascontiguousarray(mat) - return QImage( - arr.data, - arr.shape[1], - arr.shape[0], - arr.strides[0], - QImage.Format.Format_RGB888, - ).rgbSwapped() - - -def qImageToCvMatBgr(qImg: QImage): - # from Bing AI, references - # 1: https://stackoverflow.com/q/384759/16484891 | CC BY-SA 4.0 - # 2: https://stackoverflow.com/q/37552924/16484891 | CC BY-SA 3.0 - qImg = qImg.convertToFormat(QImage.Format.Format_RGB888) - qImg = qImg.copy().rgbSwapped() - return np.ndarray( - (qImg.height(), qImg.width(), 3), - buffer=qImg.constBits(), - strides=[qImg.bytesPerLine(), 3, 1], - dtype=np.uint8, - ) diff --git a/ui/implements/tabs/tabOcr/tabOcr_B30.py b/ui/implements/tabs/tabOcr/tabOcr_B30.py index 37ff12b..0473c7b 100644 --- a/ui/implements/tabs/tabOcr/tabOcr_B30.py +++ b/ui/implements/tabs/tabOcr/tabOcr_B30.py @@ -1,16 +1,20 @@ import logging import cv2 +import numpy as np from arcaea_offline_ocr.b30.chieri.v4.ocr import ChieriBotV4Ocr -from arcaea_offline_ocr.phash_db import ImagePHashDatabase -from arcaea_offline_ocr.sift_db import SIFTDatabase +from arcaea_offline_ocr.phash_db import ImagePhashDatabase from arcaea_offline_ocr.utils import imread_unicode +from PIL import Image from PySide6.QtCore import Signal, Slot -from PySide6.QtWidgets import QWidget +from PySide6.QtWidgets import QFileDialog, QMessageBox, 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.ocr.dependencies import ( + getCv2StatModelStatusText, + getPhashDatabaseStatusText, +) from ui.extends.shared.language import LanguageChangeEventFilter from ui.extends.shared.settings import ( B30_KNN_MODEL_FILE, @@ -36,97 +40,119 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget): self.b30TypeComboBox.setCurrentIndex(0) self.b30TypeComboBox.setEnabled(False) - self.imageSelector.filesSelected.connect(self.imageSelected) - self.knnModelSelector.filesSelected.connect(self.knnModelSelected) - self.b30KnnModelSelector.filesSelected.connect(self.b30KnnModelSelected) - self.phashDatabaseSelector.filesSelected.connect(self.phashDatabaseSelected) + self.dependencies_knnModelSelector.filesSelected.connect(self.knnModelSelected) + self.dependencies_b30KnnModelSelector.filesSelected.connect( + self.b30KnnModelSelected + ) + self.dependencies_phashDatabaseSelector.filesSelected.connect( + self.phashDatabaseSelected + ) - 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.phashDatabase = None self.ocr = None - self.tryPrepareOcr.connect(self.prepareOcr) - logger.info("Applying settings...") - self.knnModelSelector.connectSettings(KNN_MODEL_FILE) - self.b30KnnModelSelector.connectSettings(B30_KNN_MODEL_FILE) - self.phashDatabaseSelector.connectSettings(PHASH_DATABASE_FILE) + self.dependencies_knnModelSelector.connectSettings(KNN_MODEL_FILE) + self.dependencies_b30KnnModelSelector.connectSettings(B30_KNN_MODEL_FILE) + self.dependencies_phashDatabaseSelector.connectSettings(PHASH_DATABASE_FILE) self.ocrQueueModel = OcrQueueModel(self) self.ocrQueue.setModel(self.ocrQueueModel) - def imageSelected(self): - if selectedFiles := self.imageSelector.selectedFiles(): - imagePath = selectedFiles[0] - self.imagePath = imagePath - self.img = imread_unicode(imagePath) - self.tryPrepareOcr.emit() + # def imageSelected(self): + # if selectedFiles := self.imageSelector.selectedFiles(): + # imagePath = selectedFiles[0] + # self.imagePath = imagePath + # self.img = imread_unicode(imagePath) + # self.tryPrepareOcr.emit() def knnModelSelected(self): - if selectedFiles := self.knnModelSelector.selectedFiles(): - knnModelPath = selectedFiles[0] - self.knnModel = cv2.ml.KNearest.load(knnModelPath) - self.tryPrepareOcr.emit() + try: + filePath = self.dependencies_knnModelSelector.selectedFiles()[0] + self.knnModel = cv2.ml.KNearest.load(filePath) + except Exception: + self.knnModel = None + logger.exception("Error loading knn model:") + finally: + self.dependencies_knnModelStatusLabel.setText( + getCv2StatModelStatusText(self.knnModel) + ) def b30KnnModelSelected(self): - if selectedFiles := self.b30KnnModelSelector.selectedFiles(): - b30KnnModelPath = selectedFiles[0] - self.b30KnnModel = cv2.ml.KNearest.load(b30KnnModelPath) - self.tryPrepareOcr.emit() + try: + filePath = self.dependencies_b30KnnModelSelector.selectedFiles()[0] + self.b30KnnModel = cv2.ml.KNearest.load(filePath) + except Exception: + self.b30KnnModel = None + logger.exception("Error loading b30 knn model:") + finally: + self.dependencies_b30KnnModelStatusLabel.setText( + getCv2StatModelStatusText(self.b30KnnModel) + ) def phashDatabaseSelected(self): - if selectedFiles := self.phashDatabaseSelector.selectedFiles(): - phashDatabasePath = selectedFiles[0] - self.phashDatabase = ImagePHashDatabase(phashDatabasePath) - self.tryPrepareOcr.emit() + try: + filePath = self.dependencies_phashDatabaseSelector.selectedFiles()[0] + self.phashDatabase = ImagePhashDatabase(filePath) + except Exception: + self.phashDatabase = None + logger.exception("Error loading phash database:") + finally: + self.dependencies_phashDatabaseStatusLabel.setText( + getPhashDatabaseStatusText(self.phashDatabase) + ) - def prepareOcr(self): + def checkDependencies(self): b30Type = self.b30TypeComboBox.currentData() if not b30Type: + return False + elif b30Type == "chieri_v4": + return ( + self.knnModel is not None + and self.b30KnnModel is not None + and self.phashDatabase is not None + ) + else: + return False + + @Slot() + def on_ocr_addImageButton_clicked(self): + if not self.checkDependencies(): + QMessageBox.critical(self, None, "Dependencies not configured.") return - if b30Type == "chieri_v4": - if ( - not self.imagePath - or not self.knnModel - or not self.b30KnnModel - or not self.phashDatabase - ): - return + imagePath, _ = QFileDialog.getOpenFileName( + self, None, "", "Image Files (*.png *.jpg *.jpeg *.bmp *.webp);;*" + ) - self.ocrQueueModel.clear() + if not imagePath: + return - ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.phashDatabase) - ocr.set_factor(self.img) - self.ocr = ocr + self.ocrQueueModel.clear() - roi = ocr.rois - for component in roi.components(self.img): - qImage = cv2BgrMatToQImage(component.copy()) - self.ocrQueueModel.addItem(qImage) + img = imread_unicode(imagePath, cv2.IMREAD_COLOR) + ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.phashDatabase) + ocr.set_factor(img) + self.ocr = ocr + + roi = ocr.rois + for component in roi.components(img): + qImage = Image.fromarray(component.copy()).toqimage() + 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.phashDatabase - ): + if not self.ocr: return for row in range(self.ocrQueueModel.rowCount()): index = self.ocrQueueModel.index(row, 0) qImage = index.data(OcrQueueModel.ImageQImageRole) - cv2Mat = qImageToCvMatBgr(qImage) + cv2Mat = np.array(Image.fromqimage(qImage)) runnable = ChieriV4OcrRunnable(self.ocr, cv2Mat) self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole) self.ocrQueueModel.setData( diff --git a/ui/resources/lang/en_US.ts b/ui/resources/lang/en_US.ts index 60cd6b7..973c9a2 100644 --- a/ui/resources/lang/en_US.ts +++ b/ui/resources/lang/en_US.ts @@ -692,24 +692,24 @@ validation B30 Image Type - - knnModelSelector.title - Select KNearest Model + + dependencies.title + OCR Dependencies - - b30KnnModelSelector.title - Select B30 Specialized KNearest Model + + dependencies.knnModel + KNearest model - phashDatabaseSelector.title - Select Image PHash Database + dependencies.b30KnnModel + B30 KNearest model - - imageSelector.title - Select Image + + dependencies.phashDatabase + Image pHash database diff --git a/ui/resources/lang/zh_CN.ts b/ui/resources/lang/zh_CN.ts index 47eb270..0064c23 100644 --- a/ui/resources/lang/zh_CN.ts +++ b/ui/resources/lang/zh_CN.ts @@ -691,24 +691,24 @@ B30 图片类型 - - knnModelSelector.title - 选择 KNearest 模型 + + dependencies.title + OCR 依赖 - - b30KnnModelSelector.title - 选择 B30 特别版 KNearest 模型 + + dependencies.knnModel + KNearest 模型 - phashDatabaseSelector.title - 选择图像 PHash 数据库 + dependencies.b30KnnModel + B30 KNearest 模型 - - imageSelector.title - 选择图片 + + dependencies.phashDatabase + 图像 pHash 数据库