283375/arcaea-offline-pyside-ui
1
0
Fork 0
mirror of https://github.com/283375/arcaea-offline-pyside-ui.git synced 2025-07-01 04:16:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip: b30 ocr tab

Browse Source
This commit is contained in:
283375
2023-08-19 01:21:23 +08:00
parent 9eab7ffc63
commit 38a08a4b17
20 changed files with 726 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
ui/extends/components/ocrQueue/__init__.py
View File

@ -1,9 +1,10 @@
import logging
from typing import Any, Callable
from typing import Any, Callable, Optional, overload
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.b30.shared import B30OcrResultItem
from arcaea_offline_ocr.device.shared import DeviceOcrResult
from PySide6.QtCore import (
QAbstractListModel,
@ -30,7 +31,7 @@ logger = logging.getLogger(__name__)
class OcrRunnableSignals(QObject):
rowId: int = -1
resultReady = Signal(DeviceOcrResult)
resultReady = Signal("QVariant")
finished = Signal()
@ -45,7 +46,7 @@ class OcrQueueModel(QAbstractListModel):
ImageQImageRole = Qt.ItemDataRole.UserRole + 2
ImagePixmapRole = Qt.ItemDataRole.UserRole + 3
DeviceOcrResultRole = Qt.ItemDataRole.UserRole + 10
OcrResultRole = Qt.ItemDataRole.UserRole + 10
ScoreInsertRole = Qt.ItemDataRole.UserRole + 11
ChartRole = Qt.ItemDataRole.UserRole + 12
ScoreValidateOkRole = Qt.ItemDataRole.UserRole + 13
@ -97,55 +98,86 @@ class OcrQueueModel(QAbstractListModel):
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.OcrResultRole:
item[self.OcrResultRole] = value
updateRole = role
if role == self.ChartRole and isinstance(value, Chart):
item[self.ChartRole] = value
self.updateScoreValidateOk(index.row())
self.updateRole = role
updateRole = role
if role == self.ScoreInsertRole and isinstance(value, ScoreInsert):
item[self.ScoreInsertRole] = value
self.updateScoreValidateOk(index.row())
self.updateRole = role
updateRole = role
if role == self.ScoreValidateOkRole and isinstance(value, bool):
item[self.ScoreValidateOkRole] = value
self.updateRole = role
updateRole = role
if role == self.OcrRunnableRole and isinstance(value, OcrRunnable):
item[self.OcrRunnableRole] = value
self.updateRole = role
updateRole = role
if role == self.ProcessOcrResultFuncRole and callable(value):
item[self.ProcessOcrResultFuncRole] = value
self.updateRole = role
updateRole = role
if updateRole is not None:
self.dataChanged.emit(index, index, [updateRole])
return True
else:
logger.warning(
f"{repr(self)} setData at row {index.row()} with role {role} and value {value} rejected."
)
return False
@overload
def addItem(
self,
image: str,
runnable: OcrRunnable = None,
process_func: Callable[[Optional[str], QImage, Any], ScoreInsert] = None,
):
...
@overload
def addItem(
self,
image: QImage,
runnable: OcrRunnable = None,
process_func: Callable[[Optional[str], QImage, Any], ScoreInsert] = None,
):
...
def addItem(
self,
imagePath: str,
runnable: OcrRunnable = None,
process_func: Callable = None,
image,
runnable=None,
process_func=None,
):
if imagePath in self.imagePaths or not QFileInfo(imagePath).exists():
logger.warning(f"Attempting to add an invalid file {imagePath}")
return
if isinstance(image, str):
if image in self.imagePaths or not QFileInfo(image).exists():
logger.warning(f"Attempting to add an invalid file {image}")
return
imagePath = image
qImage = QImage(image)
qPixmap = QPixmap(image)
elif isinstance(image, QImage):
imagePath = None
qImage = image.copy()
qPixmap = QPixmap(qImage)
else:
raise ValueError("Unsupported type for `image`")
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.ImageQImageRole: qImage,
self.ImagePixmapRole: qPixmap,
self.OcrResultRole: None,
self.ScoreInsertRole: None,
self.ChartRole: None,
self.ScoreValidateOkRole: False,
@ -155,19 +187,19 @@ class OcrQueueModel(QAbstractListModel):
)
self.endInsertRows()
def updateOcrResult(self, row: int, result: DeviceOcrResult) -> bool:
if not 0 <= row < self.rowCount() or not isinstance(result, DeviceOcrResult):
def updateOcrResult(self, row: int, result: Any) -> bool:
if not 0 <= row < self.rowCount():
return False
index = self.index(row, 0)
imagePath: str = index.data(self.ImagePathRole)
qImage: QImage = index.data(self.ImageQImageRole)
print(row, result)
processOcrResultFunc = index.data(self.ProcessOcrResultFuncRole)
chart, scoreInsert = processOcrResultFunc(imagePath, result)
chart, scoreInsert = processOcrResultFunc(imagePath, qImage, result)
# song_id = self.__db.fuzzy_search_song_id(result.title)[0][0]
self.setData(index, result, self.DeviceOcrResultRole)
self.setData(index, result, self.OcrResultRole)
self.setData(index, chart, self.ChartRole)
self.setData(index, scoreInsert, self.ScoreInsertRole)
return True
@ -175,7 +207,6 @@ class OcrQueueModel(QAbstractListModel):
@Slot(DeviceOcrResult)
def ocrTaskReady(self, result: DeviceOcrResult):
row = self.sender().rowId
print(row)
self.updateOcrResult(row, result)
@Slot()
@ -184,7 +215,6 @@ class OcrQueueModel(QAbstractListModel):
self.progress.emit(self.__taskFinishedNum)
if self.__taskFinishedNum == self.__taskNum:
self.finished.emit()
print("model finished")
def startQueue(self):
self.__taskNum = self.rowCount()
@ -266,11 +296,11 @@ class OcrQueueTableProxyModel(QAbstractTableModel):
OcrQueueModel.ImagePixmapRole,
],
[
OcrQueueModel.DeviceOcrResultRole,
OcrQueueModel.OcrResultRole,
OcrQueueModel.ChartRole,
],
[
OcrQueueModel.DeviceOcrResultRole,
OcrQueueModel.OcrResultRole,
OcrQueueModel.ScoreInsertRole,
OcrQueueModel.ChartRole,
OcrQueueModel.ScoreValidateOkRole,
@ -333,9 +363,9 @@ class OcrQueueTableProxyModel(QAbstractTableModel):
def setData(self, index, value, role):
if index.column() == 2 and role == OcrQueueModel.ChartRole:
return self.sourceModel().setItemChart(index.row(), value)
return self.sourceModel().setData(index, value, role)
if index.column() == 3 and role == OcrQueueModel.ScoreInsertRole:
return self.sourceModel().setItemScore(index.row(), value)
return self.sourceModel().setData(index, value, role)
return False
def flags(self, index: QModelIndex) -> Qt.ItemFlag:
@ -365,9 +395,7 @@ class OcrChartDelegate(ChartDelegate):
return index.data(OcrQueueModel.ChartRole)
def paintWarningBackground(self, index: QModelIndex) -> bool:
return isinstance(
index.data(OcrQueueModel.DeviceOcrResultRole), DeviceOcrResult
)
return isinstance(index.data(OcrQueueModel.OcrResultRole), DeviceOcrResult)
def setModelData(self, editor, model: OcrQueueTableProxyModel, index):
if editor.validate():
@ -385,9 +413,13 @@ class OcrScoreDelegate(ScoreDelegate):
return index.data(OcrQueueModel.ScoreValidateOkRole)
def paintWarningBackground(self, index: QModelIndex) -> bool:
return isinstance(
index.data(OcrQueueModel.DeviceOcrResultRole), DeviceOcrResult
)
return True
# return isinstance(self.getChart(index), Chart) and isinstance(
# self.getScore(index), ScoreInsert
# )
# return isinstance(
# index.data(OcrQueueModel.OcrResultRole), (DeviceOcrResult, B30OcrResultItem)
# )
def setModelData(self, editor, model: OcrQueueTableProxyModel, index):
if super().confirmSetModelData(editor):

28
ui/extends/shared/cv2_utils.py Normal file
View File

@ -0,0 +1,28 @@
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,
)

5
ui/extends/shared/delegates/imageDelegate.py
View File

@ -51,7 +51,10 @@ class ImageDelegate(QStyledItemDelegate):
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())
imagePath = self.getImagePath(index)
label.setWindowTitle(
QFileInfo(imagePath).fileName() if imagePath else "Preview"
)
pixmap = pixmap.scaled(
800,
800,

51
ui/extends/tabs/tabOcr/tabOcr_B30.py Normal file
View File

@ -0,0 +1,51 @@
import logging
from arcaea_offline.database import Database
from arcaea_offline.models import Chart, ScoreInsert
from arcaea_offline_ocr.b30.chieri.v4.ocr import ChieriBotV4Ocr
from arcaea_offline_ocr.b30.shared import B30OcrResultItem
from PySide6.QtGui import QImage
logger = logging.getLogger(__name__)
from ui.extends.components.ocrQueue import OcrRunnable
class ChieriV4OcrRunnable(OcrRunnable):
def __init__(self, ocr: ChieriBotV4Ocr, component):
super().__init__()
self.ocr = ocr
self.component = component.copy()
def run(self):
try:
result = self.ocr.ocr_component(self.component)
self.signals.resultReady.emit(result)
except Exception:
logger.exception("ChieriV4 ocr component error")
finally:
self.signals.finished.emit()
def b30ResultToScoreInsert(_, qImage: QImage, result: B30OcrResultItem):
if not result.song_id and not result.title:
raise ValueError("no title or song_id")
db = Database()
if not result.song_id:
song_id = db.fuzzy_search_song_id(result.title)[0][0]
else:
song_id = result.song_id
chart = Chart.from_db_row(db.get_chart(song_id, result.rating_class))
score = ScoreInsert(
song_id=song_id,
rating_class=result.rating_class,
score=result.score,
time=1485014400,
pure=result.pure,
far=result.far,
lost=result.lost,
)
return (chart, score)

6
ui/extends/tabs/tabOcr.py → ui/extends/tabs/tabOcr/tabOcr_Device.py
View File

@ -3,7 +3,7 @@ import logging
from typing import Tuple
from arcaea_offline.database import Database
from arcaea_offline.models import ScoreInsert, Chart
from arcaea_offline.models import Chart, ScoreInsert
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
@ -52,7 +52,9 @@ def getImageDate(imagePath: str) -> QDateTime:
class ScoreInsertConverter:
@staticmethod
def deviceV2(imagePath: str, result: DeviceOcrResult) -> Tuple[Chart, ScoreInsert]:
def deviceV2(
imagePath: str, _, result: DeviceOcrResult
) -> Tuple[Chart, ScoreInsert]:
db = Database()
scoreInsert = ScoreInsert(
song_id=result.song_id,