From 0126ce6b9a578211aae73f1c0817da3c46adde59 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 15 Aug 2023 14:23:06 +0800 Subject: [PATCH] wip: refactor --- ui/designer/components/ocrQueue.ui | 174 +++++++ ui/designer/components/ocrQueue_ui.py | 141 ++++++ ui/designer/tabs/tabOcr.ui | 126 +---- ui/designer/tabs/tabOcr_ui.py | 89 +--- ui/extends/components/ocrQueue/__init__.py | 394 +++++++++++++++ ui/extends/ocr.py | 14 +- ui/extends/shared/delegates/imageDelegate.py | 70 +++ ui/extends/shared/delegates/scoreDelegate.py | 4 +- ui/extends/tabs/tabOcr.py | 487 ++----------------- ui/implements/components/__init__.py | 1 + ui/implements/components/devicesComboBox.py | 4 +- ui/implements/components/ocrQueue.py | 133 +++++ ui/implements/tabs/tabOcr.py | 107 +--- 13 files changed, 1010 insertions(+), 734 deletions(-) create mode 100644 ui/designer/components/ocrQueue.ui create mode 100644 ui/designer/components/ocrQueue_ui.py create mode 100644 ui/extends/components/ocrQueue/__init__.py create mode 100644 ui/extends/shared/delegates/imageDelegate.py create mode 100644 ui/implements/components/ocrQueue.py diff --git a/ui/designer/components/ocrQueue.ui b/ui/designer/components/ocrQueue.ui new file mode 100644 index 0000000..7a0ce83 --- /dev/null +++ b/ui/designer/components/ocrQueue.ui @@ -0,0 +1,174 @@ + + + OcrQueue + + + + 0 + 0 + 741 + 372 + + + + OcrQueue + + + + + + ocr.title + + + + + + ocr.queue.title + + + + + + ocr.queue.addImageButton + + + + + + + true + + + ocr.queue.removeSelected + + + + + + + true + + + ocr.queue.removeAll + + + + + + + Qt::Vertical + + + + 20 + 40 + + + + + + + + ocr.queue.startOcrButton + + + + + + + + + + + + QAbstractItemView::DoubleClicked|QAbstractItemView::EditKeyPressed + + + QAbstractItemView::MultiSelection + + + QAbstractItemView::SelectRows + + + QAbstractItemView::ScrollPerPixel + + + QAbstractItemView::ScrollPerPixel + + + + + + + 0 + + + 0 + + + 0 + + + Qt::AlignCenter + + + %v/%m - %p% + + + + + + + + + ocr.results + + + + + + true + + + ocr.results.acceptSelectedButton + + + + + + + ocr.results.acceptAllButton + + + + + + + Qt::Vertical + + + + 20 + 40 + + + + + + + + ocr.results.ignoreValidate + + + + + + + + + + + + + + diff --git a/ui/designer/components/ocrQueue_ui.py b/ui/designer/components/ocrQueue_ui.py new file mode 100644 index 0000000..b6ae814 --- /dev/null +++ b/ui/designer/components/ocrQueue_ui.py @@ -0,0 +1,141 @@ +# -*- coding: utf-8 -*- + +################################################################################ +## Form generated from reading UI file 'ocrQueue.ui' +## +## Created by: Qt User Interface Compiler version 6.5.1 +## +## WARNING! All changes made in this file will be lost when recompiling UI file! +################################################################################ + +from PySide6.QtCore import (QCoreApplication, QDate, QDateTime, QLocale, + QMetaObject, QObject, QPoint, QRect, + QSize, QTime, QUrl, Qt) +from PySide6.QtGui import (QBrush, QColor, QConicalGradient, QCursor, + QFont, QFontDatabase, QGradient, QIcon, + QImage, QKeySequence, QLinearGradient, QPainter, + QPalette, QPixmap, QRadialGradient, QTransform) +from PySide6.QtWidgets import (QAbstractItemView, QApplication, QCheckBox, QGroupBox, + QHBoxLayout, QHeaderView, QProgressBar, QPushButton, + QSizePolicy, QSpacerItem, QTableView, QVBoxLayout, + QWidget) + +class Ui_OcrQueue(object): + def setupUi(self, OcrQueue): + if not OcrQueue.objectName(): + OcrQueue.setObjectName(u"OcrQueue") + OcrQueue.resize(741, 372) + self.verticalLayout = QVBoxLayout(OcrQueue) + self.verticalLayout.setObjectName(u"verticalLayout") + self.groupBox_2 = QGroupBox(OcrQueue) + self.groupBox_2.setObjectName(u"groupBox_2") + self.horizontalLayout = QHBoxLayout(self.groupBox_2) + self.horizontalLayout.setObjectName(u"horizontalLayout") + self.groupBox_3 = QGroupBox(self.groupBox_2) + self.groupBox_3.setObjectName(u"groupBox_3") + self.verticalLayout_2 = QVBoxLayout(self.groupBox_3) + self.verticalLayout_2.setObjectName(u"verticalLayout_2") + self.ocr_addImageButton = QPushButton(self.groupBox_3) + self.ocr_addImageButton.setObjectName(u"ocr_addImageButton") + + self.verticalLayout_2.addWidget(self.ocr_addImageButton) + + self.ocr_removeSelectedButton = QPushButton(self.groupBox_3) + self.ocr_removeSelectedButton.setObjectName(u"ocr_removeSelectedButton") + self.ocr_removeSelectedButton.setEnabled(True) + + self.verticalLayout_2.addWidget(self.ocr_removeSelectedButton) + + self.ocr_removeAllButton = QPushButton(self.groupBox_3) + self.ocr_removeAllButton.setObjectName(u"ocr_removeAllButton") + self.ocr_removeAllButton.setEnabled(True) + + self.verticalLayout_2.addWidget(self.ocr_removeAllButton) + + self.verticalSpacer = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) + + self.verticalLayout_2.addItem(self.verticalSpacer) + + self.ocr_startButton = QPushButton(self.groupBox_3) + self.ocr_startButton.setObjectName(u"ocr_startButton") + + self.verticalLayout_2.addWidget(self.ocr_startButton) + + + self.horizontalLayout.addWidget(self.groupBox_3) + + self.verticalLayout_3 = QVBoxLayout() + self.verticalLayout_3.setObjectName(u"verticalLayout_3") + self.tableView = QTableView(self.groupBox_2) + self.tableView.setObjectName(u"tableView") + self.tableView.setEditTriggers(QAbstractItemView.DoubleClicked|QAbstractItemView.EditKeyPressed) + self.tableView.setSelectionMode(QAbstractItemView.MultiSelection) + self.tableView.setSelectionBehavior(QAbstractItemView.SelectRows) + self.tableView.setVerticalScrollMode(QAbstractItemView.ScrollPerPixel) + self.tableView.setHorizontalScrollMode(QAbstractItemView.ScrollPerPixel) + + self.verticalLayout_3.addWidget(self.tableView) + + self.progressBar = QProgressBar(self.groupBox_2) + self.progressBar.setObjectName(u"progressBar") + self.progressBar.setMinimum(0) + self.progressBar.setMaximum(0) + self.progressBar.setValue(0) + self.progressBar.setAlignment(Qt.AlignCenter) + self.progressBar.setFormat(u"%v/%m - %p%") + + self.verticalLayout_3.addWidget(self.progressBar) + + + self.horizontalLayout.addLayout(self.verticalLayout_3) + + self.groupBox_5 = QGroupBox(self.groupBox_2) + self.groupBox_5.setObjectName(u"groupBox_5") + self.verticalLayout_4 = QVBoxLayout(self.groupBox_5) + self.verticalLayout_4.setObjectName(u"verticalLayout_4") + self.ocr_acceptSelectedButton = QPushButton(self.groupBox_5) + self.ocr_acceptSelectedButton.setObjectName(u"ocr_acceptSelectedButton") + self.ocr_acceptSelectedButton.setEnabled(True) + + self.verticalLayout_4.addWidget(self.ocr_acceptSelectedButton) + + self.ocr_acceptAllButton = QPushButton(self.groupBox_5) + self.ocr_acceptAllButton.setObjectName(u"ocr_acceptAllButton") + + self.verticalLayout_4.addWidget(self.ocr_acceptAllButton) + + self.verticalSpacer_2 = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) + + self.verticalLayout_4.addItem(self.verticalSpacer_2) + + self.ocr_ignoreValidateCheckBox = QCheckBox(self.groupBox_5) + self.ocr_ignoreValidateCheckBox.setObjectName(u"ocr_ignoreValidateCheckBox") + + self.verticalLayout_4.addWidget(self.ocr_ignoreValidateCheckBox) + + + self.horizontalLayout.addWidget(self.groupBox_5) + + + self.verticalLayout.addWidget(self.groupBox_2) + + + self.retranslateUi(OcrQueue) + + QMetaObject.connectSlotsByName(OcrQueue) + # setupUi + + def retranslateUi(self, OcrQueue): + OcrQueue.setWindowTitle(QCoreApplication.translate("OcrQueue", u"OcrQueue", None)) + self.groupBox_2.setTitle(QCoreApplication.translate("OcrQueue", u"ocr.title", None)) + self.groupBox_3.setTitle(QCoreApplication.translate("OcrQueue", u"ocr.queue.title", None)) + self.ocr_addImageButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.queue.addImageButton", None)) + self.ocr_removeSelectedButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.queue.removeSelected", None)) + self.ocr_removeAllButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.queue.removeAll", None)) + self.ocr_startButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.queue.startOcrButton", None)) + self.groupBox_5.setTitle(QCoreApplication.translate("OcrQueue", u"ocr.results", None)) + self.ocr_acceptSelectedButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.results.acceptSelectedButton", None)) + self.ocr_acceptAllButton.setText(QCoreApplication.translate("OcrQueue", u"ocr.results.acceptAllButton", None)) + self.ocr_ignoreValidateCheckBox.setText(QCoreApplication.translate("OcrQueue", u"ocr.results.ignoreValidate", None)) + # retranslateUi + diff --git a/ui/designer/tabs/tabOcr.ui b/ui/designer/tabs/tabOcr.ui index 4be5b93..b03a1f8 100644 --- a/ui/designer/tabs/tabOcr.ui +++ b/ui/designer/tabs/tabOcr.ui @@ -55,125 +55,7 @@ - - - ocr.queue.title - - - - - - ocr.queue.addImageButton - - - - - - - true - - - ocr.queue.removeSelected - - - - - - - true - - - ocr.queue.removeAll - - - - - - - Qt::Vertical - - - - 20 - 40 - - - - - - - - ocr.queue.startOcrButton - - - - - - - - - - QAbstractItemView::DoubleClicked|QAbstractItemView::EditKeyPressed - - - QAbstractItemView::MultiSelection - - - QAbstractItemView::SelectRows - - - QAbstractItemView::ScrollPerPixel - - - QAbstractItemView::ScrollPerPixel - - - - - - - ocr.results - - - - - - true - - - ocr.results.acceptSelectedButton - - - - - - - ocr.results.acceptAllButton - - - - - - - Qt::Vertical - - - - 20 - 40 - - - - - - - - ocr.results.ignoreValidate - - - - - + @@ -192,6 +74,12 @@ QComboBox
ui.implements.components
+ + OcrQueue + QWidget +
ui.implements.components
+ 1 + diff --git a/ui/designer/tabs/tabOcr_ui.py b/ui/designer/tabs/tabOcr_ui.py index 3b466ab..e29f627 100644 --- a/ui/designer/tabs/tabOcr_ui.py +++ b/ui/designer/tabs/tabOcr_ui.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'tabOcr.ui' ## -## Created by: Qt User Interface Compiler version 6.5.0 +## Created by: Qt User Interface Compiler version 6.5.1 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ @@ -15,11 +15,10 @@ from PySide6.QtGui import (QBrush, QColor, QConicalGradient, QCursor, QFont, QFontDatabase, QGradient, QIcon, QImage, QKeySequence, QLinearGradient, QPainter, QPalette, QPixmap, QRadialGradient, QTransform) -from PySide6.QtWidgets import (QAbstractItemView, QApplication, QCheckBox, QGroupBox, - QHBoxLayout, QHeaderView, QPushButton, QSizePolicy, - QSpacerItem, QTableView, QVBoxLayout, QWidget) +from PySide6.QtWidgets import (QApplication, QGroupBox, QHBoxLayout, QPushButton, + QSizePolicy, QVBoxLayout, QWidget) -from ui.implements.components import (DevicesComboBox, FileSelector) +from ui.implements.components import (DevicesComboBox, FileSelector, OcrQueue) class Ui_TabOcr(object): def setupUi(self, TabOcr): @@ -67,75 +66,10 @@ class Ui_TabOcr(object): self.groupBox_2.setObjectName(u"groupBox_2") self.horizontalLayout = QHBoxLayout(self.groupBox_2) self.horizontalLayout.setObjectName(u"horizontalLayout") - self.groupBox_3 = QGroupBox(self.groupBox_2) - self.groupBox_3.setObjectName(u"groupBox_3") - self.verticalLayout_2 = QVBoxLayout(self.groupBox_3) - self.verticalLayout_2.setObjectName(u"verticalLayout_2") - self.ocr_addImageButton = QPushButton(self.groupBox_3) - self.ocr_addImageButton.setObjectName(u"ocr_addImageButton") + self.ocrQueue = OcrQueue(self.groupBox_2) + self.ocrQueue.setObjectName(u"ocrQueue") - self.verticalLayout_2.addWidget(self.ocr_addImageButton) - - self.ocr_removeSelectedButton = QPushButton(self.groupBox_3) - self.ocr_removeSelectedButton.setObjectName(u"ocr_removeSelectedButton") - self.ocr_removeSelectedButton.setEnabled(True) - - self.verticalLayout_2.addWidget(self.ocr_removeSelectedButton) - - self.ocr_removeAllButton = QPushButton(self.groupBox_3) - self.ocr_removeAllButton.setObjectName(u"ocr_removeAllButton") - self.ocr_removeAllButton.setEnabled(True) - - self.verticalLayout_2.addWidget(self.ocr_removeAllButton) - - self.verticalSpacer = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) - - self.verticalLayout_2.addItem(self.verticalSpacer) - - self.ocr_startButton = QPushButton(self.groupBox_3) - self.ocr_startButton.setObjectName(u"ocr_startButton") - - self.verticalLayout_2.addWidget(self.ocr_startButton) - - - self.horizontalLayout.addWidget(self.groupBox_3) - - self.tableView = QTableView(self.groupBox_2) - self.tableView.setObjectName(u"tableView") - self.tableView.setEditTriggers(QAbstractItemView.DoubleClicked|QAbstractItemView.EditKeyPressed) - self.tableView.setSelectionMode(QAbstractItemView.MultiSelection) - self.tableView.setSelectionBehavior(QAbstractItemView.SelectRows) - self.tableView.setVerticalScrollMode(QAbstractItemView.ScrollPerPixel) - self.tableView.setHorizontalScrollMode(QAbstractItemView.ScrollPerPixel) - - self.horizontalLayout.addWidget(self.tableView) - - self.groupBox_5 = QGroupBox(self.groupBox_2) - self.groupBox_5.setObjectName(u"groupBox_5") - self.verticalLayout_4 = QVBoxLayout(self.groupBox_5) - self.verticalLayout_4.setObjectName(u"verticalLayout_4") - self.ocr_acceptSelectedButton = QPushButton(self.groupBox_5) - self.ocr_acceptSelectedButton.setObjectName(u"ocr_acceptSelectedButton") - self.ocr_acceptSelectedButton.setEnabled(True) - - self.verticalLayout_4.addWidget(self.ocr_acceptSelectedButton) - - self.ocr_acceptAllButton = QPushButton(self.groupBox_5) - self.ocr_acceptAllButton.setObjectName(u"ocr_acceptAllButton") - - self.verticalLayout_4.addWidget(self.ocr_acceptAllButton) - - self.verticalSpacer_2 = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) - - self.verticalLayout_4.addItem(self.verticalSpacer_2) - - self.ocr_ignoreValidateCheckBox = QCheckBox(self.groupBox_5) - self.ocr_ignoreValidateCheckBox.setObjectName(u"ocr_ignoreValidateCheckBox") - - self.verticalLayout_4.addWidget(self.ocr_ignoreValidateCheckBox) - - - self.horizontalLayout.addWidget(self.groupBox_5) + self.horizontalLayout.addWidget(self.ocrQueue) self.verticalLayout_3.addWidget(self.groupBox_2) @@ -151,15 +85,6 @@ class Ui_TabOcr(object): self.groupBox.setTitle(QCoreApplication.translate("TabOcr", u"deviceSelector.title", None)) self.groupBox_4.setTitle(QCoreApplication.translate("TabOcr", u"tesseractSelector.title", None)) self.groupBox_2.setTitle(QCoreApplication.translate("TabOcr", u"ocr.title", None)) - self.groupBox_3.setTitle(QCoreApplication.translate("TabOcr", u"ocr.queue.title", None)) - self.ocr_addImageButton.setText(QCoreApplication.translate("TabOcr", u"ocr.queue.addImageButton", None)) - self.ocr_removeSelectedButton.setText(QCoreApplication.translate("TabOcr", u"ocr.queue.removeSelected", None)) - self.ocr_removeAllButton.setText(QCoreApplication.translate("TabOcr", u"ocr.queue.removeAll", None)) - self.ocr_startButton.setText(QCoreApplication.translate("TabOcr", u"ocr.queue.startOcrButton", None)) - self.groupBox_5.setTitle(QCoreApplication.translate("TabOcr", u"ocr.results", None)) - self.ocr_acceptSelectedButton.setText(QCoreApplication.translate("TabOcr", u"ocr.results.acceptSelectedButton", None)) - self.ocr_acceptAllButton.setText(QCoreApplication.translate("TabOcr", u"ocr.results.acceptAllButton", None)) - self.ocr_ignoreValidateCheckBox.setText(QCoreApplication.translate("TabOcr", u"ocr.results.ignoreValidate", None)) pass # retranslateUi diff --git a/ui/extends/components/ocrQueue/__init__.py b/ui/extends/components/ocrQueue/__init__.py new file mode 100644 index 0000000..08060d8 --- /dev/null +++ b/ui/extends/components/ocrQueue/__init__.py @@ -0,0 +1,394 @@ +import logging +from typing import Any, Callable + +from arcaea_offline.calculate import calculate_score_range +from arcaea_offline.database import Database +from arcaea_offline.models import Chart, ScoreInsert +from arcaea_offline_ocr.device.shared import DeviceOcrResult +from PySide6.QtCore import ( + QAbstractListModel, + QAbstractTableModel, + QCoreApplication, + QFileInfo, + QModelIndex, + QObject, + QRunnable, + Qt, + QThreadPool, + Signal, + Slot, +) +from PySide6.QtGui import QImage, QPixmap + +from ui.extends.shared.delegates.chartDelegate import ChartDelegate +from ui.extends.shared.delegates.imageDelegate import ImageDelegate +from ui.extends.shared.delegates.scoreDelegate import ScoreDelegate + +logger = logging.getLogger(__name__) + + +class OcrRunnableSignals(QObject): + rowId: int = -1 + + resultReady = Signal(DeviceOcrResult) + finished = Signal() + + +class OcrRunnable(QRunnable): + def __init__(self): + super().__init__() + self.signals = OcrRunnableSignals() + + +class OcrQueueModel(QAbstractListModel): + ImagePathRole = Qt.ItemDataRole.UserRole + 1 + ImageQImageRole = Qt.ItemDataRole.UserRole + 2 + ImagePixmapRole = Qt.ItemDataRole.UserRole + 3 + + DeviceOcrResultRole = Qt.ItemDataRole.UserRole + 10 + ScoreInsertRole = Qt.ItemDataRole.UserRole + 11 + ChartRole = Qt.ItemDataRole.UserRole + 12 + ScoreValidateOkRole = Qt.ItemDataRole.UserRole + 13 + + OcrRunnableRole = Qt.ItemDataRole.UserRole + 20 + ProcessOcrResultFuncRole = ( + Qt.ItemDataRole.UserRole + 21 + ) # Callable[[imageStr, DeviceOcrResult], tuple[Chart, ScoreInsert]] + + started = Signal() + progress = Signal(int) + finished = Signal() + + def __init__(self, parent=None): + super().__init__(parent) + self.__db = Database() + self.__items: list[dict[int, Any]] = [] + + self.__taskFinishedNum = 0 + + @property + def imagePaths(self): + return [item.get(self.ImagePathRole) for item in self.__items] + + def clear(self): + self.beginResetModel() + self.beginRemoveRows(QModelIndex(), 0, self.rowCount() - 1) + self.__items.clear() + self.endRemoveRows() + self.__taskFinishedNum = 0 + self.endResetModel() + + def rowCount(self, *args): + return len(self.__items) + + def data(self, index, role): + if ( + index.isValid() + and 0 <= index.row() < self.rowCount() + and index.column() == 0 + ): + return self.__items[index.row()].get(role) + return None + + def setData(self, index: QModelIndex, value: Any, role: int): + if not 0 <= index.row() < self.rowCount(): + return False + + item = self.__items[index.row()] + updateRole = None + + if role == self.DeviceOcrResultRole and isinstance(value, DeviceOcrResult): + item[self.DeviceOcrResultRole] = value + self.updateRole = role + + if role == self.ChartRole and isinstance(value, Chart): + item[self.ChartRole] = value + self.updateScoreValidateOk(index.row()) + self.updateRole = role + + if role == self.ScoreInsertRole and isinstance(value, ScoreInsert): + item[self.ScoreInsertRole] = value + self.updateScoreValidateOk(index.row()) + self.updateRole = role + + if role == self.ScoreValidateOkRole and isinstance(value, bool): + item[self.ScoreValidateOkRole] = value + self.updateRole = role + + if role == self.OcrRunnableRole and isinstance(value, OcrRunnable): + item[self.OcrRunnableRole] = value + self.updateRole = role + + if role == self.ProcessOcrResultFuncRole and callable(value): + item[self.ProcessOcrResultFuncRole] = value + self.updateRole = role + + if updateRole is not None: + self.dataChanged.emit(index, index, [updateRole]) + return True + else: + return False + + def addItem( + self, + imagePath: str, + runnable: OcrRunnable = None, + process_func: Callable = None, + ): + if imagePath in self.imagePaths or not QFileInfo(imagePath).exists(): + logger.warning(f"Attempting to add an invalid file {imagePath}") + return + + self.beginInsertRows(QModelIndex(), self.rowCount(), self.rowCount()) + self.__items.append( + { + self.ImagePathRole: imagePath, + self.ImageQImageRole: QImage(imagePath), + self.ImagePixmapRole: QPixmap(imagePath), + self.DeviceOcrResultRole: None, + self.ScoreInsertRole: None, + self.ChartRole: None, + self.ScoreValidateOkRole: False, + self.OcrRunnableRole: runnable, + self.ProcessOcrResultFuncRole: process_func, + } + ) + self.endInsertRows() + + def updateOcrResult(self, row: int, result: DeviceOcrResult) -> bool: + if not 0 <= row < self.rowCount() or not isinstance(result, DeviceOcrResult): + return False + + index = self.index(row, 0) + imagePath: str = index.data(self.ImagePathRole) + processOcrResultFunc = index.data(self.ProcessOcrResultFuncRole) + + chart, scoreInsert = processOcrResultFunc(imagePath, result) + + # song_id = self.__db.fuzzy_search_song_id(result.title)[0][0] + + self.setData(index, result, self.DeviceOcrResultRole) + self.setData(index, chart, self.ChartRole) + self.setData(index, scoreInsert, self.ScoreInsertRole) + return True + + @Slot(DeviceOcrResult) + def ocrTaskReady(self, result: DeviceOcrResult): + row = self.sender().rowId + print(row) + self.updateOcrResult(row, result) + + @Slot() + def ocrTaskFinished(self): + self.__taskFinishedNum += 1 + self.progress.emit(self.__taskFinishedNum) + if self.__taskFinishedNum == self.__taskNum: + self.finished.emit() + print("model finished") + + def startQueue(self): + self.__taskNum = self.rowCount() + self.__taskFinishedNum = 0 + self.started.emit() + for row in range(self.rowCount()): + modelIndex = self.index(row, 0) + runnable: OcrRunnable = modelIndex.data(self.OcrRunnableRole) + runnable.signals.rowId = row + runnable.signals.resultReady.connect(self.ocrTaskReady) + runnable.signals.finished.connect(self.ocrTaskFinished) + QThreadPool.globalInstance().start(runnable) + + def updateScoreValidateOk(self, row: int): + if not 0 <= row < self.rowCount(): + return + + index = self.index(row, 0) + chart = index.data(self.ChartRole) + score = index.data(self.ScoreInsertRole) + if isinstance(chart, Chart) and isinstance(score, ScoreInsert): + scoreRange = calculate_score_range(chart, score.pure, score.far) + scoreValidateOk = scoreRange[0] <= score.score <= scoreRange[1] + self.setData(index, scoreValidateOk, self.ScoreValidateOkRole) + else: + self.setData(index, False, self.ScoreValidateOkRole) + + def acceptItem(self, row: int, ignoreValidate: bool = False): + if not 0 <= row < self.rowCount(): + return + + item = self.__items[row] + score = item[self.ScoreInsertRole] + if not isinstance(score, ScoreInsert) or ( + not item[self.ScoreValidateOkRole] and not ignoreValidate + ): + return + + try: + self.__db.insert_score(score) + self.beginRemoveRows(QModelIndex(), row, row) + self.__items.pop(row) + self.endRemoveRows() + return + except Exception as e: + logger.exception(f"Error accepting {repr(item)}") + return + + def acceptItems(self, __rows: list[int], ignoreValidate: bool = False): + items = sorted(__rows, reverse=True) + [self.acceptItem(item, ignoreValidate) for item in items] + + def acceptAllItems(self, ignoreValidate: bool = False): + self.acceptItems([*range(self.rowCount())], ignoreValidate) + + def removeItem(self, row: int): + if not 0 <= row < self.rowCount(): + return + + self.beginRemoveRows(QModelIndex(), row, row) + self.__items.pop(row) + self.endRemoveRows() + + def removeItems(self, __rows: list[int]): + rows = sorted(__rows, reverse=True) + [self.removeItem(row) for row in rows] + + +class OcrQueueTableProxyModel(QAbstractTableModel): + def __init__(self, parent=None): + super().__init__(parent) + self.retranslateHeaders() + self.__sourceModel = None + self.__columnRoleMapping = [ + [Qt.ItemDataRole.CheckStateRole], + [ + OcrQueueModel.ImagePathRole, + OcrQueueModel.ImageQImageRole, + OcrQueueModel.ImagePixmapRole, + ], + [ + OcrQueueModel.DeviceOcrResultRole, + OcrQueueModel.ChartRole, + ], + [ + OcrQueueModel.DeviceOcrResultRole, + OcrQueueModel.ScoreInsertRole, + OcrQueueModel.ChartRole, + OcrQueueModel.ScoreValidateOkRole, + ], + ] + + def retranslateHeaders(self): + self.__horizontalHeaders = [ + # fmt: off + QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.select"), + QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.imagePreview"), + QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.chart"), + QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.score"), + # fmt: on + ] + + def sourceModel(self) -> OcrQueueModel: + return self.__sourceModel + + def setSourceModel(self, sourceModel): + if not isinstance(sourceModel, OcrQueueModel): + return False + + # connect signals + sourceModel.rowsAboutToBeInserted.connect(self.rowsAboutToBeInserted) + sourceModel.rowsInserted.connect(self.rowsInserted) + sourceModel.rowsAboutToBeRemoved.connect(self.rowsAboutToBeRemoved) + sourceModel.rowsRemoved.connect(self.rowsRemoved) + sourceModel.dataChanged.connect(self.dataChanged) + sourceModel.layoutAboutToBeChanged.connect(self.layoutAboutToBeChanged) + sourceModel.layoutChanged.connect(self.layoutChanged) + + self.__sourceModel = sourceModel + return True + + def rowCount(self, *args): + return self.sourceModel().rowCount() + + def columnCount(self, *args): + return len(self.__horizontalHeaders) + + def headerData(self, section: int, orientation: Qt.Orientation, role: int): + if ( + orientation == Qt.Orientation.Horizontal + and 0 <= section < len(self.__horizontalHeaders) + and role == Qt.ItemDataRole.DisplayRole + ): + return self.__horizontalHeaders[section] + return None + + def data(self, index, role): + if ( + 0 <= index.row() < self.rowCount() + and 0 <= index.column() < self.columnCount() + and role in self.__columnRoleMapping[index.column()] + ): + srcIndex = self.sourceModel().index(index.row(), 0) + return srcIndex.data(role) + return None + + def setData(self, index, value, role): + if index.column() == 2 and role == OcrQueueModel.ChartRole: + return self.sourceModel().setItemChart(index.row(), value) + if index.column() == 3 and role == OcrQueueModel.ScoreInsertRole: + return self.sourceModel().setItemScore(index.row(), value) + return False + + def flags(self, index: QModelIndex) -> Qt.ItemFlag: + flags = ( + self.sourceModel().flags(index) + if isinstance(self.sourceModel(), OcrQueueModel) + else super().flags(index) + ) + flags = flags | Qt.ItemFlag.ItemIsEnabled + flags = flags | Qt.ItemFlag.ItemIsEditable + flags = flags | Qt.ItemFlag.ItemIsSelectable + if index.column() == 0: + flags = flags & ~Qt.ItemFlag.ItemIsEnabled & ~Qt.ItemFlag.ItemIsEditable + return flags + + +class OcrImageDelegate(ImageDelegate): + def getPixmap(self, index: QModelIndex): + return index.data(OcrQueueModel.ImagePixmapRole) + + def getImagePath(self, index: QModelIndex): + return index.data(OcrQueueModel.ImagePathRole) + + +class OcrChartDelegate(ChartDelegate): + def getChart(self, index: QModelIndex) -> Chart | None: + return index.data(OcrQueueModel.ChartRole) + + def paintWarningBackground(self, index: QModelIndex) -> bool: + return isinstance( + index.data(OcrQueueModel.DeviceOcrResultRole), DeviceOcrResult + ) + + def setModelData(self, editor, model: OcrQueueTableProxyModel, index): + if editor.validate(): + model.setData(index, editor.value(), OcrQueueModel.ChartRole) + + +class OcrScoreDelegate(ScoreDelegate): + def getScoreInsert(self, index: QModelIndex): + return index.data(OcrQueueModel.ScoreInsertRole) + + def getChart(self, index: QModelIndex): + return index.data(OcrQueueModel.ChartRole) + + def getScoreValidateOk(self, index: QModelIndex): + return index.data(OcrQueueModel.ScoreValidateOkRole) + + def paintWarningBackground(self, index: QModelIndex) -> bool: + return isinstance( + index.data(OcrQueueModel.DeviceOcrResultRole), DeviceOcrResult + ) + + def setModelData(self, editor, model: OcrQueueTableProxyModel, index): + if super().confirmSetModelData(editor): + model.setData(index, editor.value(), OcrQueueModel.ScoreInsertRole) diff --git a/ui/extends/ocr.py b/ui/extends/ocr.py index 56563e7..212ab91 100644 --- a/ui/extends/ocr.py +++ b/ui/extends/ocr.py @@ -1,16 +1,24 @@ try: import json - from arcaea_offline_ocr.device import Device + from arcaea_offline_ocr.device.v1.definition import DeviceV1 + from arcaea_offline_ocr.device.v2.definition import DeviceV2 - def load_devices_json(filepath: str) -> list[Device]: + def load_devices_json(filepath: str) -> list[DeviceV1]: with open(filepath, "r", encoding="utf-8") as f: file_content = f.read() if len(file_content) == 0: return [] content = json.loads(file_content) assert isinstance(content, list) - return [Device.from_json_object(item) for item in content] + devices = [] + for item in content: + version = item["version"] + if version == 1: + devices.append(DeviceV1(**item)) + elif version == 2: + devices.append(DeviceV2(**item)) + return devices except Exception: diff --git a/ui/extends/shared/delegates/imageDelegate.py b/ui/extends/shared/delegates/imageDelegate.py new file mode 100644 index 0000000..64694c6 --- /dev/null +++ b/ui/extends/shared/delegates/imageDelegate.py @@ -0,0 +1,70 @@ +from PySide6.QtCore import QFileInfo, QModelIndex, QRect, QSize, Qt +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import QLabel, QStyledItemDelegate, QWidget + + +class ImageDelegate(QStyledItemDelegate): + def getPixmap(self, index: QModelIndex): + raise NotImplementedError("getPixmap not implemented.") + + def getImagePath(self, index: QModelIndex): + raise NotImplementedError("getImagePath not implemented.") + + def scalePixmap(self, pixmap: QPixmap): + return pixmap.scaled( + 100, + 100, + Qt.AspectRatioMode.KeepAspectRatio, + Qt.TransformationMode.SmoothTransformation, + ) + + def paint(self, painter, option, index): + pixmap = self.getPixmap(index) + if not isinstance(pixmap, QPixmap): + imagePath = self.getImagePath(index) + option.text = imagePath + super().paint(painter, option, index) + else: + pixmap = self.scalePixmap(pixmap) + # https://stackoverflow.com/a/32047499/16484891 + # CC BY-SA 3.0 + x = option.rect.center().x() - pixmap.rect().width() / 2 + y = option.rect.center().y() - pixmap.rect().height() / 2 + + painter.drawPixmap( + QRect(x, y, pixmap.rect().width(), pixmap.rect().height()), pixmap + ) + + def sizeHint(self, option, index) -> QSize: + pixmap = self.getPixmap(index) + if isinstance(pixmap, QPixmap): + pixmap = self.scalePixmap(pixmap) + return pixmap.size() + else: + return QSize(100, 75) + + def createEditor(self, parent, option, index) -> QWidget: + pixmap = self.getPixmap(index) + if isinstance(pixmap, QPixmap): + label = QLabel(parent) + label.setWindowFlags(Qt.WindowType.Window) + label.setWindowFlag(Qt.WindowType.WindowMinimizeButtonHint, False) + label.setWindowFlag(Qt.WindowType.WindowMaximizeButtonHint, False) + label.setWindowFlag(Qt.WindowType.WindowCloseButtonHint, True) + label.setWindowTitle(QFileInfo(self.getImagePath(index)).fileName()) + pixmap = pixmap.scaled( + 800, + 800, + Qt.AspectRatioMode.KeepAspectRatio, + Qt.TransformationMode.SmoothTransformation, + ) + label.setMinimumSize(pixmap.size()) + label.setPixmap(pixmap) + label.move(parent.mapToGlobal(parent.pos())) + return label + + def setModelData(self, *args): + ... + + def updateEditorGeometry(self, *args): + ... diff --git a/ui/extends/shared/delegates/scoreDelegate.py b/ui/extends/shared/delegates/scoreDelegate.py index c2c19e0..381a65e 100644 --- a/ui/extends/shared/delegates/scoreDelegate.py +++ b/ui/extends/shared/delegates/scoreDelegate.py @@ -126,6 +126,8 @@ class ScoreDelegate(TextSegmentDelegate): ] ] + score_str = str(score.score).rjust(8, "0") + score_str = f"{score_str[:-6]}'{score_str[-6:-3]}'{score_str[-3:]}" score_font = QFont(option.font) score_font.setPointSize(12) score_grade_font = QFont(score_font) @@ -140,7 +142,7 @@ class ScoreDelegate(TextSegmentDelegate): self.FontRole: score_grade_font, }, {self.TextRole: " | "}, - {self.TextRole: str(score.score), self.FontRole: score_font}, + {self.TextRole: score_str, self.FontRole: score_font}, ], [ { diff --git a/ui/extends/tabs/tabOcr.py b/ui/extends/tabs/tabOcr.py index 38e9950..b0bebe4 100644 --- a/ui/extends/tabs/tabOcr.py +++ b/ui/extends/tabs/tabOcr.py @@ -1,476 +1,71 @@ import contextlib import logging -from typing import Any +from typing import Tuple -import exif -from arcaea_offline.calculate import calculate_score_range from arcaea_offline.database import Database -from arcaea_offline.models import Chart, ScoreInsert -from arcaea_offline_ocr.device import Device -from arcaea_offline_ocr.recognize import RecognizeResult, recognize -from PySide6.QtCore import ( - QAbstractListModel, - QAbstractTableModel, - QCoreApplication, - QDateTime, - QFileInfo, - QModelIndex, - QObject, - QRect, - QRunnable, - QSize, - Qt, - QThreadPool, - Signal, - Slot, -) -from PySide6.QtGui import QPixmap -from PySide6.QtWidgets import QLabel, QStyledItemDelegate, QWidget +from arcaea_offline.models import ScoreInsert, Chart +from arcaea_offline_ocr.device.shared import DeviceOcrResult +from arcaea_offline_ocr.device.v2.ocr import DeviceV2Ocr +from arcaea_offline_ocr.device.v2.rois import DeviceV2Rois +from arcaea_offline_ocr.utils import imread_unicode +from PySide6.QtCore import QDateTime, QFileInfo -from ui.extends.shared.delegates.chartDelegate import ChartDelegate -from ui.extends.shared.delegates.scoreDelegate import ScoreDelegate -from ui.implements.components.scoreEditor import ScoreEditor +from ui.extends.components.ocrQueue import OcrRunnable logger = logging.getLogger(__name__) - -class OcrTaskSignals(QObject): - resultReady = Signal(int, RecognizeResult) - finished = Signal(int) +import exif -class OcrTask(QRunnable): - def __init__(self, index: int, device: Device, imagePath: str): +class TabDeviceV2OcrRunnable(OcrRunnable): + def __init__(self, imagePath, device, knnModel, siftDb): super().__init__() - self.index = index - self.device = device self.imagePath = imagePath - self.signals = OcrTaskSignals() + self.device = device + self.knnModel = knnModel + self.siftDb = siftDb def run(self): try: - result = recognize(self.imagePath, self.device) - self.signals.resultReady.emit(self.index, result) - logger.info( - f"OcrTask {self.imagePath} with {repr(self.device)} got result {repr(result)}" - ) - except Exception as e: - logger.exception( - f"OcrTask {self.imagePath} with {repr(self.device)} failed" - ) + rois = DeviceV2Rois(self.device, imread_unicode(self.imagePath)) + ocr = DeviceV2Ocr(self.knnModel, self.siftDb) + result = ocr.ocr(rois) + self.signals.resultReady.emit(result) + except Exception: + logger.exception(f"DeviceV2 ocr {self.imagePath} error") finally: - self.signals.finished.emit(self.index) + self.signals.finished.emit() -class OcrQueueModel(QAbstractListModel): - ImagePathRole = Qt.ItemDataRole.UserRole + 1 - ImagePixmapRole = Qt.ItemDataRole.UserRole + 2 - RecognizeResultRole = Qt.ItemDataRole.UserRole + 10 - ScoreInsertRole = Qt.ItemDataRole.UserRole + 11 - ChartRole = Qt.ItemDataRole.UserRole + 12 - ScoreValidateOkRole = Qt.ItemDataRole.UserRole + 13 +def getImageDate(imagePath: str) -> QDateTime: + datetime = None + with contextlib.suppress(Exception): + with open(imagePath, "rb") as imgf: + exifImage = exif.Image(imgf.read()) + if exifImage.has_exif and exifImage.get("datetime_original"): + datetimeStr = exifImage.get("datetime_original") + datetime = QDateTime.fromString(datetimeStr, "yyyy:MM:dd hh:mm:ss") + if not isinstance(datetime, QDateTime): + datetime = QFileInfo(imagePath).birthTime() + return datetime - started = Signal() - finished = Signal() - def __init__(self, parent=None): - super().__init__(parent) - self.__db = Database() - self.__items: list[dict[int, Any]] = [] - - @property - def imagePaths(self): - return [item.get(self.ImagePathRole) for item in self.__items] - - def clear(self): - self.beginResetModel() - self.beginRemoveRows(QModelIndex(), 0, self.rowCount() - 1) - self.__items.clear() - self.endRemoveRows() - self.endResetModel() - - def rowCount(self, *args): - return len(self.__items) - - def data(self, index, role): - if ( - index.isValid() - and 0 <= index.row() < self.rowCount() - and index.column() == 0 - ): - return self.__items[index.row()].get(role) - return None - - def setData(self, *args): - return False - - def addItem(self, imagePath: str): - if imagePath in self.imagePaths or not QFileInfo(imagePath).exists(): - return - - self.beginInsertRows(QModelIndex(), self.rowCount(), self.rowCount()) - self.__items.append( - { - self.ImagePathRole: imagePath, - self.ImagePixmapRole: QPixmap(imagePath), - self.RecognizeResultRole: None, - self.ScoreInsertRole: None, - self.ChartRole: None, - self.ScoreValidateOkRole: False, - } - ) - self.endInsertRows() - - def updateOcrResult(self, row: int, result: RecognizeResult) -> bool: - if not 0 <= row < self.rowCount() or not isinstance(result, RecognizeResult): - return False - - item = self.__items[row] - - imagePath: str = item[self.ImagePathRole] - datetime = None - with contextlib.suppress(Exception): - with open(imagePath, "rb") as imgf: - exifImage = exif.Image(imgf.read()) - if exifImage.has_exif and exifImage.get("datetime_original"): - datetimeStr = exifImage.get("datetime_original") - datetime = QDateTime.fromString(datetimeStr, "yyyy:MM:dd hh:mm:ss") - if not isinstance(datetime, QDateTime): - datetime = QFileInfo(imagePath).birthTime() - - score = ScoreInsert( - song_id=self.__db.fuzzy_search_song_id(result.title)[0][0], +class ScoreInsertConverter: + @staticmethod + def deviceV2(imagePath: str, result: DeviceOcrResult) -> Tuple[Chart, ScoreInsert]: + db = Database() + scoreInsert = ScoreInsert( + song_id=result.song_id, rating_class=result.rating_class, score=result.score, pure=result.pure, far=result.far, lost=result.lost, - time=datetime.toSecsSinceEpoch(), + time=getImageDate(imagePath).toSecsSinceEpoch(), max_recall=result.max_recall, clear_type=None, ) chart = Chart.from_db_row( - self.__db.get_chart(score.song_id, score.rating_class) + db.get_chart(scoreInsert.song_id, scoreInsert.rating_class) ) - - item[self.RecognizeResultRole] = result - self.setItemChart(row, chart) - self.setItemScore(row, score) - modelIndex = self.index(row, 0) - self.dataChanged.emit( - modelIndex, - modelIndex, - [self.RecognizeResultRole, self.ScoreInsertRole, self.ChartRole], - ) - return True - - @Slot(int, RecognizeResult) - def ocrTaskReady(self, row: int, result: RecognizeResult): - self.updateOcrResult(row, result) - - @Slot(int) - def ocrTaskFinished(self, row: int): - self.__taskFinishedNum += 1 - if self.__taskFinishedNum == self.__taskNum: - self.finished.emit() - - def startQueue(self, device: Device): - self.__taskNum = self.rowCount() - self.__taskFinishedNum = 0 - self.started.emit() - for row in range(self.rowCount()): - modelIndex = self.index(row, 0) - imagePath: str = modelIndex.data(self.ImagePathRole) - task = OcrTask(row, device, imagePath) - task.signals.resultReady.connect(self.ocrTaskReady) - task.signals.finished.connect(self.ocrTaskFinished) - QThreadPool.globalInstance().start(task) - - def updateScoreValidateOk(self, row: int): - if not 0 <= row < self.rowCount(): - return - - item = self.__items[row] - chart = item[self.ChartRole] - score = item[self.ScoreInsertRole] - if isinstance(chart, Chart) and isinstance(score, ScoreInsert): - scoreRange = calculate_score_range(chart, score.pure, score.far) - scoreValidateOk = scoreRange[0] <= score.score <= scoreRange[1] - item[self.ScoreValidateOkRole] = scoreValidateOk - else: - item[self.ScoreValidateOkRole] = False - - modelIndex = self.index(row, 0) - self.dataChanged.emit(modelIndex, modelIndex, [self.ScoreValidateOkRole]) - - def setItemChart(self, row: int, chart: Chart): - if not 0 <= row < self.rowCount() or not isinstance(chart, Chart): - return False - - item = self.__items[row] - item[self.ChartRole] = chart - updatedRoles = [self.ChartRole] - - self.updateScoreValidateOk(row) - - modelIndex = self.index(row, 0) - self.dataChanged.emit(modelIndex, modelIndex, updatedRoles) - return True - - def setItemScore(self, row: int, score: ScoreInsert) -> bool: - if not 0 <= row < self.rowCount() or not isinstance(score, ScoreInsert): - return False - - item = self.__items[row] - item[self.ScoreInsertRole] = score - updatedRoles = [self.ScoreInsertRole] - - self.updateScoreValidateOk(row) - - modelIndex = self.index(row, 0) - self.dataChanged.emit(modelIndex, modelIndex, updatedRoles) - return True - - def acceptItem(self, row: int, ignoreValidate: bool = False): - if not 0 <= row < self.rowCount(): - return - - item = self.__items[row] - score = item[self.ScoreInsertRole] - if not isinstance(score, ScoreInsert) or ( - not item[self.ScoreValidateOkRole] and not ignoreValidate - ): - return - - try: - self.__db.insert_score(score) - self.beginRemoveRows(QModelIndex(), row, row) - self.__items.pop(row) - self.endRemoveRows() - return - except Exception as e: - logger.exception(f"Error accepting {repr(item)}") - return - - def acceptItems(self, __rows: list[int], ignoreValidate: bool = False): - items = sorted(__rows, reverse=True) - [self.acceptItem(item, ignoreValidate) for item in items] - - def acceptAllItems(self, ignoreValidate: bool = False): - self.acceptItems([*range(self.rowCount())], ignoreValidate) - - def removeItem(self, row: int): - if not 0 <= row < self.rowCount(): - return - - self.beginRemoveRows(QModelIndex(), row, row) - self.__items.pop(row) - self.endRemoveRows() - - def removeItems(self, __rows: list[int]): - rows = sorted(__rows, reverse=True) - [self.removeItem(row) for row in rows] - - -class OcrQueueTableProxyModel(QAbstractTableModel): - def __init__(self, parent=None): - super().__init__(parent) - self.retranslateHeaders() - self.__sourceModel = None - self.__columnRoleMapping = [ - [Qt.ItemDataRole.CheckStateRole], - [OcrQueueModel.ImagePathRole, OcrQueueModel.ImagePixmapRole], - [ - OcrQueueModel.RecognizeResultRole, - OcrQueueModel.ChartRole, - ], - [ - OcrQueueModel.RecognizeResultRole, - OcrQueueModel.ScoreInsertRole, - OcrQueueModel.ChartRole, - OcrQueueModel.ScoreValidateOkRole, - ], - ] - - def retranslateHeaders(self): - self.__horizontalHeaders = [ - # fmt: off - QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.select"), - QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.imagePreview"), - QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.chart"), - QCoreApplication.translate("OcrTableModel", "horizontalHeader.title.score"), - # fmt: on - ] - - def sourceModel(self) -> OcrQueueModel: - return self.__sourceModel - - def setSourceModel(self, sourceModel): - if not isinstance(sourceModel, OcrQueueModel): - return False - - # connect signals - sourceModel.rowsAboutToBeInserted.connect(self.rowsAboutToBeInserted) - sourceModel.rowsInserted.connect(self.rowsInserted) - sourceModel.rowsAboutToBeRemoved.connect(self.rowsAboutToBeRemoved) - sourceModel.rowsRemoved.connect(self.rowsRemoved) - sourceModel.dataChanged.connect(self.dataChanged) - sourceModel.layoutAboutToBeChanged.connect(self.layoutAboutToBeChanged) - sourceModel.layoutChanged.connect(self.layoutChanged) - - self.__sourceModel = sourceModel - return True - - def rowCount(self, *args): - return self.sourceModel().rowCount() - - def columnCount(self, *args): - return len(self.__horizontalHeaders) - - def headerData(self, section: int, orientation: Qt.Orientation, role: int): - if ( - orientation == Qt.Orientation.Horizontal - and 0 <= section < len(self.__horizontalHeaders) - and role == Qt.ItemDataRole.DisplayRole - ): - return self.__horizontalHeaders[section] - return None - - def data(self, index, role): - if ( - 0 <= index.row() < self.rowCount() - and 0 <= index.column() < self.columnCount() - and role in self.__columnRoleMapping[index.column()] - ): - srcIndex = self.sourceModel().index(index.row(), 0) - return srcIndex.data(role) - return None - - def setData(self, index, value, role): - if index.column() == 2 and role == OcrQueueModel.ChartRole: - return self.sourceModel().setItemChart(index.row(), value) - if index.column() == 3 and role == OcrQueueModel.ScoreInsertRole: - return self.sourceModel().setItemScore(index.row(), value) - return False - - def flags(self, index: QModelIndex) -> Qt.ItemFlag: - flags = ( - self.sourceModel().flags(index) - if isinstance(self.sourceModel(), OcrQueueModel) - else super().flags(index) - ) - flags = flags | Qt.ItemFlag.ItemIsEnabled - flags = flags | Qt.ItemFlag.ItemIsEditable - flags = flags | Qt.ItemFlag.ItemIsSelectable - if index.column() == 0: - flags = flags & ~Qt.ItemFlag.ItemIsEnabled & ~Qt.ItemFlag.ItemIsEditable - return flags - - -class ImageDelegate(QStyledItemDelegate): - def getPixmap(self, index: QModelIndex): - return index.data(OcrQueueModel.ImagePixmapRole) - - def getImagePath(self, index: QModelIndex): - return index.data(OcrQueueModel.ImagePathRole) - - def scalePixmap(self, pixmap: QPixmap): - return pixmap.scaled( - 100, - 100, - Qt.AspectRatioMode.KeepAspectRatio, - Qt.TransformationMode.SmoothTransformation, - ) - - def paint(self, painter, option, index): - pixmap = self.getPixmap(index) - if not isinstance(pixmap, QPixmap): - imagePath = self.getImagePath(index) - option.text = imagePath - super().paint(painter, option, index) - else: - pixmap = self.scalePixmap(pixmap) - # https://stackoverflow.com/a/32047499/16484891 - # CC BY-SA 3.0 - x = option.rect.center().x() - pixmap.rect().width() / 2 - y = option.rect.center().y() - pixmap.rect().height() / 2 - - painter.drawPixmap( - QRect(x, y, pixmap.rect().width(), pixmap.rect().height()), pixmap - ) - - def sizeHint(self, option, index) -> QSize: - pixmap = self.getPixmap(index) - if isinstance(pixmap, QPixmap): - pixmap = self.scalePixmap(pixmap) - return pixmap.size() - else: - return QSize(100, 75) - - def createEditor(self, parent, option, index) -> QWidget: - pixmap = self.getPixmap(index) - if isinstance(pixmap, QPixmap): - label = QLabel(parent) - label.setWindowFlags(Qt.WindowType.Window) - label.setWindowFlag(Qt.WindowType.WindowMinimizeButtonHint, False) - label.setWindowFlag(Qt.WindowType.WindowMaximizeButtonHint, False) - label.setWindowFlag(Qt.WindowType.WindowCloseButtonHint, True) - label.setWindowTitle(QFileInfo(self.getImagePath(index)).fileName()) - pixmap = pixmap.scaled( - 800, - 800, - Qt.AspectRatioMode.KeepAspectRatio, - Qt.TransformationMode.SmoothTransformation, - ) - label.setMinimumSize(pixmap.size()) - label.setPixmap(pixmap) - label.move(parent.mapToGlobal(parent.pos())) - return label - - def setModelData(self, *args): - ... - - def updateEditorGeometry(self, *args): - ... - - -class TableChartDelegate(ChartDelegate): - def getChart(self, index: QModelIndex) -> Chart | None: - return index.data(OcrQueueModel.ChartRole) - - def paintWarningBackground(self, index: QModelIndex) -> bool: - return isinstance( - index.data(OcrQueueModel.RecognizeResultRole), RecognizeResult - ) - - def setModelData(self, editor, model: OcrQueueTableProxyModel, index): - if editor.validate(): - model.setData(index, editor.value(), OcrQueueModel.ChartRole) - - -class TableScoreDelegate(ScoreDelegate): - def getScoreInsert(self, index: QModelIndex): - return index.data(OcrQueueModel.ScoreInsertRole) - - def getChart(self, index: QModelIndex): - return index.data(OcrQueueModel.ChartRole) - - def getScoreValidateOk(self, index: QModelIndex): - return index.data(OcrQueueModel.ScoreValidateOkRole) - - def paintWarningBackground(self, index: QModelIndex) -> bool: - return isinstance( - index.data(OcrQueueModel.RecognizeResultRole), RecognizeResult - ) - - # def createEditor(self, parent, option, index): - # editor = super().createEditor(parent, option, index) - # editor.setManualHandleCommit(True) - # return editor - - def setModelData(self, editor, model: OcrQueueTableProxyModel, index): - # userAcceptMessageBox = editor.triggerValidateMessageBox() - # if userAcceptMessageBox: - # model.setData(index, editor.value(), OcrQueueModel.ScoreInsertRole) - if super().confirmSetModelData(editor): - model.setData(index, editor.value(), OcrQueueModel.ScoreInsertRole) + return (chart, scoreInsert) diff --git a/ui/implements/components/__init__.py b/ui/implements/components/__init__.py index b6670d2..aa13166 100644 --- a/ui/implements/components/__init__.py +++ b/ui/implements/components/__init__.py @@ -4,3 +4,4 @@ from .elidedLabel import ElidedLabel from .fileSelector import FileSelector from .ratingClassRadioButton import RatingClassRadioButton from .scoreEditor import ScoreEditor +from .ocrQueue import OcrQueue diff --git a/ui/implements/components/devicesComboBox.py b/ui/implements/components/devicesComboBox.py index d8e8d0f..4156c9f 100644 --- a/ui/implements/components/devicesComboBox.py +++ b/ui/implements/components/devicesComboBox.py @@ -1,4 +1,4 @@ -from arcaea_offline_ocr.device import Device +from arcaea_offline_ocr.device.v1.definition import DeviceV1 from PySide6.QtCore import Qt from PySide6.QtWidgets import QComboBox @@ -13,7 +13,7 @@ class DevicesComboBox(QComboBox): super().__init__(parent) self.setItemDelegate(DescriptionDelegate(self)) - def setDevices(self, devices: list[Device]): + def setDevices(self, devices: list[DeviceV1]): self.clear() for device in devices: self.addItem(f"{device.name} ({device.uuid})", device) diff --git a/ui/implements/components/ocrQueue.py b/ui/implements/components/ocrQueue.py new file mode 100644 index 0000000..d0acd65 --- /dev/null +++ b/ui/implements/components/ocrQueue.py @@ -0,0 +1,133 @@ +from typing import Optional + +from PySide6.QtCore import Qt, QTimer, Slot +from PySide6.QtGui import QColor, QPalette +from PySide6.QtWidgets import QWidget + +from ui.designer.components.ocrQueue_ui import Ui_OcrQueue +from ui.extends.components.ocrQueue import ( + OcrChartDelegate, + OcrImageDelegate, + OcrQueueModel, + OcrQueueTableProxyModel, + OcrScoreDelegate, +) + + +class OcrQueue(Ui_OcrQueue, QWidget): + def __init__(self, parent=None): + super().__init__(parent) + self.setupUi(self) + self.__model: Optional[OcrQueueModel] = None + self.__tableProxyModel: Optional[OcrQueueTableProxyModel] = None + + self.__firstResizeDone = False + self.resizeTimer = QTimer(self) + self.resizeTimer.timeout.connect(self.tableView.resizeRowsToContents) + self.resizeTimer.timeout.connect(self.tableView.resizeColumnsToContents) + + self.tableView.setItemDelegateForColumn(1, OcrImageDelegate(self.tableView)) + self.tableView.setItemDelegateForColumn(2, OcrChartDelegate(self.tableView)) + self.tableView.setItemDelegateForColumn(3, OcrScoreDelegate(self.tableView)) + + tableViewPalette = QPalette(self.tableView.palette()) + highlightColor = QColor(tableViewPalette.color(QPalette.ColorRole.Highlight)) + highlightColor.setAlpha(25) + tableViewPalette.setColor(QPalette.ColorRole.Highlight, highlightColor) + self.tableView.setPalette(tableViewPalette) + + def model(self): + return self.__model + + def setModel(self, model: OcrQueueModel): + model.dataChanged.connect(self.resizeViewWhenDataChanged) + model.started.connect(self.ocrStarted) + model.progress.connect(self.ocrProgress) + model.finished.connect(self.ocrFinished) + model.rowsInserted.connect(self.updateProgressBarMaximum) + model.rowsRemoved.connect(self.updateProgressBarMaximum) + model.modelReset.connect(self.modelReseted) + proxyModel = OcrQueueTableProxyModel(self) + proxyModel.setSourceModel(model) + self.tableView.setModel(proxyModel) + + if self.__model: + self.__model.dataChanged.disconnect(self.resizeViewWhenDataChanged) + self.__model.started.disconnect(self.ocrStarted) + self.__model.progress.disconnect(self.ocrProgress) + self.__model.finished.disconnect(self.ocrFinished) + self.__model.rowsInserted.disconnect(self.updateProgressBarMaximum) + self.__model.rowsRemoved.disconnect(self.updateProgressBarMaximum) + self.__model.modelReset.disconnect(self.modelReseted) + if self.__tableProxyModel: + self.__tableProxyModel.deleteLater() + + self.__model = model + self.__tableProxyModel = proxyModel + + def tableProxyModel(self): + return self.__tableProxyModel + + def setOcrButtonsEnabled(self, __bool: bool): + self.ocr_addImageButton.setEnabled(__bool) + self.ocr_removeSelectedButton.setEnabled(__bool) + self.ocr_removeAllButton.setEnabled(__bool) + self.ocr_startButton.setEnabled(__bool) + self.ocr_acceptSelectedButton.setEnabled(__bool) + self.ocr_acceptAllButton.setEnabled(__bool) + self.ocr_ignoreValidateCheckBox.setEnabled(__bool) + + def resizeTableView(self): + self.tableView.resizeRowsToContents() + self.tableView.resizeColumnsToContents() + + def resizeViewWhenDataChanged(self): + if not self.__firstResizeDone: + self.resizeTableView() + self.__firstResizeDone = True + + def ocrStarted(self): + self.setOcrButtonsEnabled(False) + + def updateProgressBarMaximum(self): + self.progressBar.setMaximum(self.model().rowCount()) + + @Slot(int) + def ocrProgress(self, progress: int): + self.progressBar.setValue(progress) + + def ocrFinished(self): + self.resizeTableView() + self.setOcrButtonsEnabled(True) + + def modelReseted(self): + self.progressBar.setMaximum(0) + + @Slot() + def on_ocr_removeSelectedButton_clicked(self): + if self.model(): + rows = [ + modelIndex.row() + for modelIndex in self.tableView.selectionModel().selectedRows(0) + ] + self.model().removeItems(rows) + + @Slot() + def on_ocr_acceptSelectedButton_clicked(self): + if self.model(): + ignoreValidate = ( + self.ocr_ignoreValidateCheckBox.checkState() == Qt.CheckState.Checked + ) + rows = [ + modelIndex.row() + for modelIndex in self.tableView.selectionModel().selectedRows(0) + ] + self.model().acceptItems(rows, ignoreValidate) + + @Slot() + def on_ocr_acceptAllButton_clicked(self): + if self.model(): + ignoreValidate = ( + self.ocr_ignoreValidateCheckBox.checkState() == Qt.CheckState.Checked + ) + self.model().acceptAllItems(ignoreValidate) diff --git a/ui/implements/tabs/tabOcr.py b/ui/implements/tabs/tabOcr.py index d06a5b2..5238cfb 100644 --- a/ui/implements/tabs/tabOcr.py +++ b/ui/implements/tabs/tabOcr.py @@ -1,24 +1,21 @@ import pytesseract -from arcaea_offline_ocr_device_creation_wizard.implements.wizard import Wizard + +# from arcaea_offline_ocr_device_creation_wizard.implements.wizard import Wizard from PySide6.QtCore import QModelIndex, Qt, Slot -from PySide6.QtGui import QColor, QPalette -from PySide6.QtWidgets import QFileDialog, QWidget +from PySide6.QtGui import QColor +from PySide6.QtWidgets import QFileDialog, QHeaderView, 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 ( - ImageDelegate, - OcrQueueModel, - OcrQueueTableProxyModel, - TableChartDelegate, - TableScoreDelegate, -) +from ui.extends.tabs.tabOcr import TabDeviceV2OcrRunnable, ScoreInsertConverter class TabOcr(Ui_TabOcr, QWidget): def __init__(self, parent=None): super().__init__(parent) self.setupUi(self) + self.openWizardButton.setEnabled(False) self.deviceFileSelector.filesSelected.connect(self.deviceFileSelected) self.tesseractFileSelector.filesSelected.connect( @@ -31,36 +28,14 @@ class TabOcr(Ui_TabOcr, QWidget): self.deviceComboBox.selectDevice(settings.deviceUuid()) self.ocrQueueModel = OcrQueueModel(self) - self.ocrQueueModel.dataChanged.connect(self.resizeViewWhenScoreChanged) - self.ocrQueueModel.started.connect(self.ocrStarted) - self.ocrQueueModel.finished.connect(self.ocrFinished) - self.ocrQueueProxyModel = OcrQueueTableProxyModel(self) - self.ocrQueueProxyModel.setSourceModel(self.ocrQueueModel) - - self.tableView.setModel(self.ocrQueueProxyModel) - self.tableView.setItemDelegateForColumn(1, ImageDelegate(self.tableView)) - self.tableView.setItemDelegateForColumn(2, TableChartDelegate(self.tableView)) - self.tableView.setItemDelegateForColumn(3, TableScoreDelegate(self.tableView)) - - tableViewPalette = QPalette(self.tableView.palette()) - highlightColor = QColor(tableViewPalette.color(QPalette.ColorRole.Highlight)) - highlightColor.setAlpha(25) - tableViewPalette.setColor(QPalette.ColorRole.Highlight, highlightColor) - self.tableView.setPalette(tableViewPalette) - - @Slot(QModelIndex, QModelIndex, list) - def resizeViewWhenScoreChanged( - self, topleft: QModelIndex, bottomRight: QModelIndex, roles: list[int] - ): - if OcrQueueModel.ScoreInsertRole in roles: - rows = [*range(topleft.row(), bottomRight.row() + 1)] - [self.tableView.resizeRowToContents(row) for row in rows] - self.tableView.resizeColumnsToContents() + self.ocrQueue.setModel(self.ocrQueueModel) + self.ocrQueueProxyModel = self.ocrQueue.tableProxyModel() @Slot() def on_openWizardButton_clicked(self): - wizard = Wizard(self) - wizard.open() + # wizard = Wizard(self) + # wizard.open() + pass def deviceFileSelected(self): selectedFiles = self.deviceFileSelector.selectedFiles() @@ -73,15 +48,6 @@ class TabOcr(Ui_TabOcr, QWidget): if selectedFiles: pytesseract.pytesseract.tesseract_cmd = selectedFiles[0] - def setOcrButtonsEnabled(self, __bool: bool): - self.ocr_addImageButton.setEnabled(__bool) - self.ocr_removeSelectedButton.setEnabled(__bool) - self.ocr_removeAllButton.setEnabled(__bool) - self.ocr_startButton.setEnabled(__bool) - self.ocr_acceptSelectedButton.setEnabled(__bool) - self.ocr_acceptAllButton.setEnabled(__bool) - self.ocr_ignoreValidateCheckBox.setEnabled(__bool) - @Slot() def on_ocr_addImageButton_clicked(self): files, _filter = QFileDialog.getOpenFileNames( @@ -89,45 +55,24 @@ class TabOcr(Ui_TabOcr, QWidget): ) for file in files: self.ocrQueueModel.addItem(file) - self.tableView.resizeRowsToContents() - self.tableView.resizeColumnsToContents() + self.ocrQueue.resizeTableView() @Slot() def on_ocr_startButton_clicked(self): - self.ocrQueueModel.startQueue(self.deviceComboBox.currentData()) - - def ocrStarted(self): - self.setOcrButtonsEnabled(False) - - def ocrFinished(self): - self.setOcrButtonsEnabled(True) - - @Slot() - def on_ocr_removeSelectedButton_clicked(self): - rows = [ - modelIndex.row() - for modelIndex in self.tableView.selectionModel().selectedRows(0) - ] - self.ocrQueueModel.removeItems(rows) + for row in range(self.ocrQueueModel.rowCount()): + index = self.ocrQueueModel.index(row, 0) + imagePath = index.data(OcrQueueModel.ImagePathRole) + runnable = TabDeviceV2OcrRunnable( + imagePath, self.deviceComboBox.currentData(), self.knn, self.siftDb + ) + self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole) + self.ocrQueueModel.setData( + index, + ScoreInsertConverter.deviceV2, + OcrQueueModel.ProcessOcrResultFuncRole, + ) + self.ocrQueueModel.startQueue() @Slot() def on_ocr_removeAllButton_clicked(self): self.ocrQueueModel.clear() - - @Slot() - def on_ocr_acceptSelectedButton_clicked(self): - ignoreValidate = ( - self.ocr_ignoreValidateCheckBox.checkState() == Qt.CheckState.Checked - ) - rows = [ - modelIndex.row() - for modelIndex in self.tableView.selectionModel().selectedRows(0) - ] - self.ocrQueueModel.acceptItems(rows, ignoreValidate) - - @Slot() - def on_ocr_acceptAllButton_clicked(self): - ignoreValidate = ( - self.ocr_ignoreValidateCheckBox.checkState() == Qt.CheckState.Checked - ) - self.ocrQueueModel.acceptAllItems(ignoreValidate)