2 Commits

Author SHA1 Message Date
ad5e5ec694 wip: arcaea-offline-ocr==0.1.0
settings
2023-10-12 17:37:55 +08:00
5c5c1a227d wip: arcaea-offline-ocr==0.1.0
API changes, modifier & clear_type support
2023-10-12 17:05:04 +08:00
10 changed files with 203 additions and 252 deletions

1
.gitignore vendored
View File

@ -3,6 +3,7 @@ __debug*
arcaea_offline.db
arcaea_offline.ini
/data
ui/resources/VERSION

View File

@ -6,7 +6,7 @@ from arcaea_offline.calculate import calculate_score_range
from arcaea_offline.database import Database
from arcaea_offline.models import Chart, Score
from arcaea_offline_ocr.b30.shared import B30OcrResultItem
from arcaea_offline_ocr.device.shared import DeviceOcrResult
from arcaea_offline_ocr.device.common import DeviceOcrResult
from arcaea_offline_ocr.utils import convert_to_srgb
from PIL import Image
from PIL.ImageQt import ImageQt

47
ui/extends/shared/data.py Normal file
View File

@ -0,0 +1,47 @@
import json
import sys
from functools import cached_property
from pathlib import Path
from typing import Literal
from PySide6.QtCore import QFile
from .singleton import Singleton
TPartnerModifier = dict[str, Literal[0, 1, 2]]
class Data(metaclass=Singleton):
def __init__(self):
root = Path(sys.argv[0]).parent
self.__dataPath = (root / "data").resolve()
@property
def dataPath(self):
return self.__dataPath
@cached_property
def partnerModifiers(self) -> TPartnerModifier:
data = {}
builtinFile = QFile(":/partnerModifiers.json")
builtinFile.open(QFile.OpenModeFlag.ReadOnly)
builtinData = json.loads(str(builtinFile.readAll(), encoding="utf-8"))
builtinFile.close()
data |= builtinData
customFile = self.dataPath / "partnerModifiers.json"
if customFile.exists():
with open(customFile, "r", encoding="utf-8") as f:
customData = json.loads(f.read())
data |= customData
return data
def expirePartnerModifiersCache(self):
# expire property caches
# https://stackoverflow.com/a/69367025/16484891, CC BY-SA 4.0
self.__dict__.pop("partnerModifiers", None)
@property
def arcaeaPath(self):
return self.dataPath / "Arcaea"

View File

@ -5,12 +5,11 @@ from PySide6.QtCore import QFileInfo, QSettings, Signal
from .singleton import QObjectSingleton
__all__ = [
"LANGUAGE",
"DATABASE_URL",
"DEVICES_JSON_FILE",
"DEVICE_UUID",
"TESSERACT_FILE",
"KNN_MODEL_FILE",
"SIFT_DATABASE_FILE",
"B30_KNN_MODEL_FILE",
"PHASH_DATABASE_FILE",
"ANDREAL_FOLDER",
"ANDREAL_EXECUTABLE",
"Settings",
@ -21,12 +20,8 @@ __all__ = [
LANGUAGE = "Language"
DATABASE_URL = "DatabaseUrl"
DEVICES_JSON_FILE = "Ocr/DevicesJsonFile"
DEVICE_UUID = "Ocr/DeviceUuid"
TESSERACT_FILE = "Ocr/TesseractFile"
KNN_MODEL_FILE = "Ocr/KnnModelFile"
B30_KNN_MODEL_FILE = "Ocr/B30KnnModelFile"
SIFT_DATABASE_FILE = "Ocr/SiftDatabaseFile"
PHASH_DATABASE_FILE = "Ocr/PHashDatabaseFile"
ANDREAL_FOLDER = "Andreal/AndrealFolder"
@ -70,33 +65,6 @@ class Settings(QSettings, metaclass=QObjectSingleton):
def setDatabaseUrl(self, value: str):
self._setStrItem(DATABASE_URL, value)
def devicesJsonFile(self):
return self._strItem(DEVICES_JSON_FILE)
def setDevicesJsonFile(self, value: str):
self._setStrItem(DEVICES_JSON_FILE, value)
def resetDevicesJsonFile(self):
self._resetStrItem(DEVICES_JSON_FILE)
def deviceUuid(self):
return self._strItem(DEVICE_UUID)
def setDeviceUuid(self, value: str):
self._setStrItem(DEVICE_UUID, value)
def resetDeviceUuid(self):
self._resetStrItem(DEVICE_UUID)
def tesseractPath(self):
return self._strItem(TESSERACT_FILE)
def setTesseractPath(self, value: str):
self._setStrItem(TESSERACT_FILE, value)
def resetTesseractPath(self):
self._resetStrItem(TESSERACT_FILE)
def knnModelFile(self):
return self._strItem(KNN_MODEL_FILE)
@ -115,15 +83,6 @@ class Settings(QSettings, metaclass=QObjectSingleton):
def resetB30KnnModelFile(self):
self._resetStrItem(B30_KNN_MODEL_FILE)
def siftDatabaseFile(self):
return self._strItem(SIFT_DATABASE_FILE)
def setSiftDatabaseFile(self, value: str):
self._setStrItem(SIFT_DATABASE_FILE, value)
def resetSiftDatabaseFile(self):
self._resetStrItem(SIFT_DATABASE_FILE)
def phashDatabaseFile(self):
return self._strItem(PHASH_DATABASE_FILE)

View File

@ -1,69 +1,58 @@
import contextlib
import logging
from typing import Tuple
from typing import Tuple, Type
import cv2
import exif
from arcaea_offline.database import Database
from arcaea_offline.models import Chart, Score
from arcaea_offline_ocr.device.shared import DeviceOcrResult
from arcaea_offline_ocr.device.v2 import DeviceV2AutoRois, DeviceV2Ocr, DeviceV2Rois
from arcaea_offline_ocr.device.v2.sizes import SizesV1, SizesV2
from arcaea_offline.utils.partner import KanaeDayNight, kanae_day_night
from arcaea_offline_ocr.device import DeviceOcr, DeviceOcrResult
from arcaea_offline_ocr.device.rois import (
DeviceRois,
DeviceRoisAuto,
DeviceRoisExtractor,
DeviceRoisMasker,
)
from arcaea_offline_ocr.phash_db import ImagePhashDatabase
from arcaea_offline_ocr.utils import imread_unicode
from PySide6.QtCore import QDateTime, QFileInfo
from ui.extends.components.ocrQueue import OcrRunnable
from ui.extends.shared.data import Data
logger = logging.getLogger(__name__)
import exif
class TabDeviceV2OcrRunnable(OcrRunnable):
def __init__(self, imagePath, device, knnModel, phashDb, *, sizesV2: bool):
class TabDeviceOcrRunnable(OcrRunnable):
def __init__(
self,
imagePath: str,
rois: DeviceRois | Type[DeviceRoisAuto],
masker: DeviceRoisMasker,
knnModel: cv2.ml.KNearest,
phashDb: ImagePhashDatabase,
):
super().__init__()
self.imagePath = imagePath
self.device = device
self.rois = rois
self.masker = masker
self.knnModel = knnModel
self.phashDb = phashDb
self.sizesV2 = sizesV2
def run(self):
try:
rois = DeviceV2Rois(
self.device, imread_unicode(self.imagePath, cv2.IMREAD_COLOR)
)
rois.sizes = (
SizesV2(self.device.factor)
if self.sizesV2
else SizesV1(self.device.factor)
)
ocr = DeviceV2Ocr(self.knnModel, self.phashDb)
result = ocr.ocr(rois)
img = imread_unicode(self.imagePath, cv2.IMREAD_COLOR)
if isinstance(self.rois, type) and issubclass(self.rois, DeviceRoisAuto):
rois = self.rois(img.shape[1], img.shape[0])
else:
rois = self.rois
extractor = DeviceRoisExtractor(img, rois)
ocr = DeviceOcr(extractor, self.masker, self.knnModel, self.phashDb)
result = ocr.ocr()
self.signals.resultReady.emit(result)
except Exception:
logger.exception(f"DeviceV2 ocr {self.imagePath} error")
finally:
self.signals.finished.emit()
class TabDeviceV2AutoRoisOcrRunnable(OcrRunnable):
def __init__(self, imagePath, knnModel, phashDb, *, sizesV2: bool):
super().__init__()
self.imagePath = imagePath
self.knnModel = knnModel
self.phashDb = phashDb
self.sizesV2 = sizesV2
def run(self):
try:
rois = DeviceV2AutoRois(imread_unicode(self.imagePath, cv2.IMREAD_COLOR))
factor = rois.sizes.factor
rois.sizes = SizesV2(factor) if self.sizesV2 else SizesV1(factor)
ocr = DeviceV2Ocr(self.knnModel, self.phashDb)
result = ocr.ocr(rois)
self.signals.resultReady.emit(result)
except Exception:
logger.exception(f"DeviceV2AutoRois ocr {self.imagePath} error")
logger.exception("DeviceOcr error:")
finally:
self.signals.finished.emit()
@ -83,7 +72,24 @@ def getImageDate(imagePath: str) -> QDateTime:
class ScoreConverter:
@staticmethod
def deviceV2(imagePath: str, _, result: DeviceOcrResult) -> Tuple[Chart, Score]:
def device(imagePath: str, _, result: DeviceOcrResult) -> Tuple[Chart, Score]:
partnerModifiers = Data().partnerModifiers
imageDate = getImageDate(imagePath)
# calculate clear type
if result.partner_id == "50":
dayNight = kanae_day_night(imageDate)
modifier = 1 if dayNight == KanaeDayNight.Day else 2
else:
modifier = partnerModifiers.get(result.partner_id, 0)
if result.clear_status == 1 and modifier == 1:
clearType = 4
elif result.clear_status == 1 and modifier == 2:
clearType = 5
else:
clearType = result.clear_status
db = Database()
score = Score(
song_id=result.song_id,
@ -92,16 +98,16 @@ class ScoreConverter:
pure=result.pure,
far=result.far,
lost=result.lost,
date=getImageDate(imagePath).toSecsSinceEpoch(),
date=imageDate.toSecsSinceEpoch(),
max_recall=result.max_recall,
modifier=modifier,
clear_type=clearType,
comment=f"OCR {QFileInfo(imagePath).fileName()}",
)
chart = db.get_chart(score.song_id, score.rating_class)
if not chart:
chart = Chart(
song_id=result.song_id,
rating_class=result.rating_class,
title=result.song_id,
constant=0.0,
)
chart = db.get_chart(score.song_id, score.rating_class) or Chart(
song_id=result.song_id,
rating_class=result.rating_class,
title=result.song_id,
constant=0.0,
)
return (chart, score)

View File

@ -1,32 +0,0 @@
from arcaea_offline_ocr.device.v1.definition import DeviceV1
from PySide6.QtCore import Qt
from PySide6.QtWidgets import QComboBox
from ui.extends.ocr import load_devices_json
from ui.extends.shared.delegates.descriptionDelegate import DescriptionDelegate
class DevicesComboBox(QComboBox):
DeviceUuidRole = Qt.ItemDataRole.UserRole + 10
def __init__(self, parent=None):
super().__init__(parent)
self.setItemDelegate(DescriptionDelegate(self))
def setDevices(self, devices: list[DeviceV1]):
self.clear()
for device in devices:
self.addItem(f"{device.name} ({device.uuid})", device)
row = self.count() - 1
self.setItemData(row, device.uuid, self.DeviceUuidRole)
self.setItemData(row, device.name, DescriptionDelegate.MainTextRole)
self.setItemData(row, device.uuid, DescriptionDelegate.DescriptionTextRole)
self.setCurrentIndex(-1)
def loadDevicesJson(self, path: str):
devices = load_devices_json(path)
self.setDevices(devices)
def selectDevice(self, deviceUuid: str):
index = self.findData(deviceUuid, self.DeviceUuidRole)
self.setCurrentIndex(index)

View File

@ -1,7 +1,6 @@
from PySide6.QtCore import QCoreApplication
from PySide6.QtWidgets import QLabel, QPushButton
from ui.implements.components.devicesComboBox import DevicesComboBox
from ui.implements.components.fileSelector import FileSelector
from ui.implements.settings.settingsBaseWidget import SettingsBaseWidget
@ -12,28 +11,6 @@ class SettingsOcr(SettingsBaseWidget):
self.setupUi(self)
if self.settings.devicesJsonFile():
self.devicesJsonValueWidget.selectFile(self.settings.devicesJsonFile())
self.devicesJsonValueWidget.filesSelected.connect(self.setDevicesJson)
self.devicesJsonResetButton.clicked.connect(self.resetDevicesJson)
self.insertItem(
"devicesJson",
self.devicesJsonLabel,
self.devicesJsonValueWidget,
self.devicesJsonResetButton,
)
if self.settings.deviceUuid():
self.deviceUuidValueWidget.selectDevice(self.settings.deviceUuid())
self.deviceUuidValueWidget.activated.connect(self.setDeviceUuid)
self.deviceUuidResetButton.clicked.connect(self.resetDeviceUuid)
self.insertItem(
"deviceUuid",
self.deviceUuidLabel,
self.deviceUuidValueWidget,
self.deviceUuidResetButton,
)
if self.settings.knnModelFile():
self.knnModelFileValueWidget.selectFile(self.settings.knnModelFile())
self.knnModelFileValueWidget.filesSelected.connect(self.setKnnModelFile)
@ -56,19 +33,6 @@ class SettingsOcr(SettingsBaseWidget):
self.b30KnnModelFileResetButton,
)
if self.settings.siftDatabaseFile():
self.siftDatabaseFileValueWidget.selectFile(
self.settings.siftDatabaseFile()
)
self.siftDatabaseFileValueWidget.filesSelected.connect(self.setSiftDatabaseFile)
self.siftDatabaseFileResetButton.clicked.connect(self.resetSiftDatabaseFile)
self.insertItem(
"siftDatabaseFile",
self.siftDatabaseFileLabel,
self.siftDatabaseFileValueWidget,
self.siftDatabaseFileResetButton,
)
if self.settings.phashDatabaseFile():
self.phashDatabaseFileValueWidget.selectFile(
self.settings.phashDatabaseFile()
@ -84,34 +48,6 @@ class SettingsOcr(SettingsBaseWidget):
self.phashDatabaseFileResetButton,
)
def setDevicesJson(self):
selectedFile = self.devicesJsonValueWidget.selectedFiles()
if selectedFile and selectedFile[0]:
file = selectedFile[0]
self.settings.setDevicesJsonFile(file)
self.fillDeviceUuidComboBox()
def fillDeviceUuidComboBox(self):
devicesJsonPath = self.devicesJsonValueWidget.selectedFiles()[0]
self.deviceUuidValueWidget.loadDevicesJson(devicesJsonPath)
storedDeviceUuid = self.settings.deviceUuid()
self.deviceUuidValueWidget.selectDevice(storedDeviceUuid)
def resetDevicesJson(self):
self.deviceUuidValueWidget.clear()
self.devicesJsonValueWidget.reset()
self.settings.resetDeviceUuid()
self.settings.resetDevicesJsonFile()
def setDeviceUuid(self):
if device := self.deviceUuidValueWidget.currentData():
self.settings.setDeviceUuid(device.uuid)
def resetDeviceUuid(self):
self.deviceUuidValueWidget.setCurrentIndex(-1)
self.settings.resetDeviceUuid()
def setKnnModelFile(self):
selectedFile = self.knnModelFileValueWidget.selectedFiles()
if selectedFile and selectedFile[0]:
@ -132,16 +68,6 @@ class SettingsOcr(SettingsBaseWidget):
self.b30KnnModelFileValueWidget.reset()
self.settings.resetB30KnnModelFile()
def setSiftDatabaseFile(self):
selectedFile = self.siftDatabaseFileValueWidget.selectedFiles()
if selectedFile and selectedFile[0]:
file = selectedFile[0]
self.settings.setSiftDatabaseFile(file)
def resetSiftDatabaseFile(self):
self.siftDatabaseFileValueWidget.reset()
self.settings.resetSiftDatabaseFile()
def setPHashDatabaseFile(self):
selectedFile = self.phashDatabaseFileValueWidget.selectedFiles()
if selectedFile and selectedFile[0]:
@ -153,14 +79,6 @@ class SettingsOcr(SettingsBaseWidget):
self.settings.resetPHashDatabaseFile()
def setupUi(self, *args):
self.devicesJsonLabel = QLabel(self)
self.devicesJsonValueWidget = FileSelector(self)
self.devicesJsonResetButton = QPushButton(self)
self.deviceUuidLabel = QLabel(self)
self.deviceUuidValueWidget = DevicesComboBox(self)
self.deviceUuidResetButton = QPushButton(self)
self.knnModelFileLabel = QLabel(self)
self.knnModelFileValueWidget = FileSelector(self)
self.knnModelFileResetButton = QPushButton(self)
@ -169,10 +87,6 @@ class SettingsOcr(SettingsBaseWidget):
self.b30KnnModelFileValueWidget = FileSelector(self)
self.b30KnnModelFileResetButton = QPushButton(self)
self.siftDatabaseFileLabel = QLabel(self)
self.siftDatabaseFileValueWidget = FileSelector(self)
self.siftDatabaseFileResetButton = QPushButton(self)
self.phashDatabaseFileLabel = QLabel(self)
self.phashDatabaseFileValueWidget = FileSelector(self)
self.phashDatabaseFileResetButton = QPushButton(self)
@ -186,21 +100,12 @@ class SettingsOcr(SettingsBaseWidget):
# fmt: off
self.setTitle(QCoreApplication.translate("Settings", "ocr.title"))
self.devicesJsonLabel.setText(QCoreApplication.translate("Settings", "ocr.devicesJson.label"))
self.devicesJsonResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
self.deviceUuidLabel.setText(QCoreApplication.translate("Settings", "ocr.deviceUuid.label"))
self.deviceUuidResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
self.knnModelFileLabel.setText(QCoreApplication.translate("Settings", "ocr.knnModelFile.label"))
self.knnModelFileResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
self.b30KnnModelFileLabel.setText(QCoreApplication.translate("Settings", "ocr.b30KnnModelFile.label"))
self.b30KnnModelFileResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
self.siftDatabaseFileLabel.setText(QCoreApplication.translate("Settings", "ocr.siftDatabaseFile.label"))
self.siftDatabaseFileResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
self.phashDatabaseFileLabel.setText(QCoreApplication.translate("Settings", "ocr.phashDatabaseFile.label"))
self.phashDatabaseFileResetButton.setText(QCoreApplication.translate("Settings", "resetButton"))
# fmt: on

View File

@ -1,20 +1,21 @@
import logging
import cv2
from arcaea_offline_ocr.device.rois import (
DeviceRoisAutoT1,
DeviceRoisAutoT2,
DeviceRoisMaskerAutoT1,
DeviceRoisMaskerAutoT2,
)
from arcaea_offline_ocr.phash_db import ImagePhashDatabase
from PySide6.QtCore import Qt, Slot
from PySide6.QtWidgets import QApplication, QFileDialog, QWidget
from PySide6.QtCore import Slot
from PySide6.QtWidgets import QApplication, QFileDialog, QMessageBox, 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 KNN_MODEL_FILE, PHASH_DATABASE_FILE
from ui.extends.tabs.tabOcr.tabOcr_Device import (
ScoreConverter,
TabDeviceV2AutoRoisOcrRunnable,
TabDeviceV2OcrRunnable,
)
from ui.extends.tabs.tabOcr.tabOcr_Device import ScoreConverter, TabDeviceOcrRunnable
logger = logging.getLogger(__name__)
@ -106,9 +107,15 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
try:
knnModelFile = self.dependencies_knnModelSelector.selectedFiles()[0]
self.knnModel = cv2.ml.KNearest.load(knnModelFile)
self.dependencies_knnModelStatusLabel.setText(
f'<font color="green">OK</font>, varCount {self.knnModel.getVarCount()}'
)
varCount = self.knnModel.getVarCount()
if varCount != 81:
self.dependencies_knnModelStatusLabel.setText(
f'<font color="darkorange">WARN</font>, varCount {varCount}'
)
else:
self.dependencies_knnModelStatusLabel.setText(
f'<font color="green">OK</font>, varCount {varCount}'
)
except Exception:
logger.exception("Error loading knn model:")
self.dependencies_knnModelStatusLabel.setText(
@ -150,30 +157,51 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
QApplication.processEvents()
self.ocrQueue.resizeTableView()
def deviceRois(self):
if self.options_roisUseCustomCheckBox.isChecked():
...
else:
selectedPreset = self.options_roisComboBox.currentData()
if selectedPreset == "AutoT1":
return DeviceRoisAutoT1
elif selectedPreset == "AutoT2":
return DeviceRoisAutoT2
else:
QMessageBox.critical(self, None, "Select a Rois preset first.")
return None
def deviceRoisMasker(self):
if self.options_maskerUseCustomCheckBox.isChecked():
...
else:
selectedPreset = self.options_maskerComboBox.currentData()
if selectedPreset == "AutoT1":
return DeviceRoisMaskerAutoT1()
elif selectedPreset == "AutoT2":
return DeviceRoisMaskerAutoT2()
else:
QMessageBox.critical(self, None, "Select a Masker preset first.")
return None
@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.phashDatabase,
sizesV2=self.deviceSizesV2CheckBox.isChecked(),
)
else:
runnable = TabDeviceV2OcrRunnable(
imagePath,
self.deviceComboBox.currentData(),
self.knnModel,
self.phashDatabase,
sizesV2=self.deviceSizesV2CheckBox.isChecked(),
)
rois = self.deviceRois()
masker = self.deviceRoisMasker()
if rois is None or masker is None:
return
runnable = TabDeviceOcrRunnable(
imagePath, rois, masker, self.knnModel, self.phashDatabase
)
self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
self.ocrQueueModel.setData(
index,
ScoreConverter.deviceV2,
ScoreConverter.device,
OcrQueueModel.ProcessOcrResultFuncRole,
)
self.ocrQueueModel.startQueue()

View File

@ -0,0 +1,35 @@
{
"__COMMENT__": "1: EASY, 2: HARD",
"0": 1,
"0u": 1,
"7": 2,
"9": 1,
"10": 2,
"10u": 2,
"15": 1,
"16": 1,
"20": 1,
"28": 2,
"28u": 2,
"29": 2,
"29u": 2,
"35": 2,
"36": 2,
"36u": 2,
"37": 2,
"41": 2,
"42": 2,
"42u": 2,
"43": 2,
"43u": 2,
"54": 2,
"55": 2,
"57": 2,
"61": 2,
"64": 2,
"66": 2,
"66u": 2,
"67": 2,
"68": 1,
"70": 2
}

View File

@ -4,6 +4,8 @@
<file>VERSION</file>
<file>LICENSE</file>
<file>partnerModifiers.json</file>
<file>images/icon.png</file>
<file>images/logo.png</file>
<file>images/stepCalculator/stamina.png</file>