diff --git a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui new file mode 100644 index 0000000..ac24876 --- /dev/null +++ b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui @@ -0,0 +1,255 @@ + + + tabOcr_BuildPHashDatabase + + + + 0 + 0 + 632 + 551 + + + + tabOcr_BuildPHashDatabase + + + + + + Qt::Vertical + + + + 20 + 40 + + + + + + + + folders.title + + + + + + + + folders.songDir + + + + + + + + 0 + 0 + + + + + + + + + + + + folders.charIconDir + + + + + + + + 0 + 0 + + + + + + + + + + + + + options.title + + + + + + + + + + hash_size + + + Qt::AlignRight|Qt::AlignTrailing|Qt::AlignVCenter + + + + + + + + 0 + 0 + + + + 2 + + + 64 + + + 16 + + + + + + + + + + + highfreq_factor + + + Qt::AlignRight|Qt::AlignTrailing|Qt::AlignVCenter + + + + + + + + 0 + 0 + + + + 32 + + + 4 + + + + + + + + + + + Qt::Horizontal + + + + 40 + 20 + + + + + + + + resetButton + + + + + + + + + + 0 + + + 0 + + + Qt::AlignCenter + + + %v/%m - %p% + + + + + + + + + Qt::Horizontal + + + + 40 + 20 + + + + + + + + buildButton + + + + + + + Qt::Horizontal + + + + 40 + 20 + + + + + + + + + + Qt::Vertical + + + + 20 + 40 + + + + + + + + + FileSelector + QWidget +
ui.implements.components.fileSelector
+ 1 + + + + + diff --git a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py new file mode 100644 index 0000000..d7ff2bd --- /dev/null +++ b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- + +################################################################################ +## Form generated from reading UI file 'tabOcr_BuildPHashDatabase.ui' +## +## Created by: Qt User Interface Compiler version 6.5.2 +## +## 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 (QApplication, QGroupBox, QHBoxLayout, QLabel, + QProgressBar, QPushButton, QSizePolicy, QSpacerItem, + QSpinBox, QVBoxLayout, QWidget) + +from ui.implements.components.fileSelector import FileSelector + +class Ui_tabOcr_BuildPHashDatabase(object): + def setupUi(self, tabOcr_BuildPHashDatabase): + if not tabOcr_BuildPHashDatabase.objectName(): + tabOcr_BuildPHashDatabase.setObjectName(u"tabOcr_BuildPHashDatabase") + tabOcr_BuildPHashDatabase.resize(632, 551) + tabOcr_BuildPHashDatabase.setWindowTitle(u"tabOcr_BuildPHashDatabase") + self.verticalLayout_3 = QVBoxLayout(tabOcr_BuildPHashDatabase) + self.verticalLayout_3.setObjectName(u"verticalLayout_3") + self.verticalSpacer = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) + + self.verticalLayout_3.addItem(self.verticalSpacer) + + self.groupBox = QGroupBox(tabOcr_BuildPHashDatabase) + self.groupBox.setObjectName(u"groupBox") + self.verticalLayout = QVBoxLayout(self.groupBox) + self.verticalLayout.setObjectName(u"verticalLayout") + self.horizontalLayout = QHBoxLayout() + self.horizontalLayout.setObjectName(u"horizontalLayout") + self.label = QLabel(self.groupBox) + self.label.setObjectName(u"label") + + self.horizontalLayout.addWidget(self.label) + + self.songDirSelector = FileSelector(self.groupBox) + self.songDirSelector.setObjectName(u"songDirSelector") + sizePolicy = QSizePolicy(QSizePolicy.MinimumExpanding, QSizePolicy.Preferred) + sizePolicy.setHorizontalStretch(0) + sizePolicy.setVerticalStretch(0) + sizePolicy.setHeightForWidth(self.songDirSelector.sizePolicy().hasHeightForWidth()) + self.songDirSelector.setSizePolicy(sizePolicy) + + self.horizontalLayout.addWidget(self.songDirSelector) + + + self.verticalLayout.addLayout(self.horizontalLayout) + + self.horizontalLayout_2 = QHBoxLayout() + self.horizontalLayout_2.setObjectName(u"horizontalLayout_2") + self.label_2 = QLabel(self.groupBox) + self.label_2.setObjectName(u"label_2") + + self.horizontalLayout_2.addWidget(self.label_2) + + self.charIconDirSelector = FileSelector(self.groupBox) + self.charIconDirSelector.setObjectName(u"charIconDirSelector") + sizePolicy.setHeightForWidth(self.charIconDirSelector.sizePolicy().hasHeightForWidth()) + self.charIconDirSelector.setSizePolicy(sizePolicy) + + self.horizontalLayout_2.addWidget(self.charIconDirSelector) + + + self.verticalLayout.addLayout(self.horizontalLayout_2) + + + self.verticalLayout_3.addWidget(self.groupBox) + + self.groupBox_2 = QGroupBox(tabOcr_BuildPHashDatabase) + self.groupBox_2.setObjectName(u"groupBox_2") + self.horizontalLayout_3 = QHBoxLayout(self.groupBox_2) + self.horizontalLayout_3.setObjectName(u"horizontalLayout_3") + self.verticalLayout_2 = QVBoxLayout() + self.verticalLayout_2.setObjectName(u"verticalLayout_2") + self.horizontalLayout_6 = QHBoxLayout() + self.horizontalLayout_6.setObjectName(u"horizontalLayout_6") + self.label_3 = QLabel(self.groupBox_2) + self.label_3.setObjectName(u"label_3") + self.label_3.setText(u"hash_size") + self.label_3.setAlignment(Qt.AlignRight|Qt.AlignTrailing|Qt.AlignVCenter) + + self.horizontalLayout_6.addWidget(self.label_3) + + self.hashSizeSpinBox = QSpinBox(self.groupBox_2) + self.hashSizeSpinBox.setObjectName(u"hashSizeSpinBox") + sizePolicy1 = QSizePolicy(QSizePolicy.Maximum, QSizePolicy.Fixed) + sizePolicy1.setHorizontalStretch(0) + sizePolicy1.setVerticalStretch(0) + sizePolicy1.setHeightForWidth(self.hashSizeSpinBox.sizePolicy().hasHeightForWidth()) + self.hashSizeSpinBox.setSizePolicy(sizePolicy1) + self.hashSizeSpinBox.setMinimum(2) + self.hashSizeSpinBox.setMaximum(64) + self.hashSizeSpinBox.setValue(16) + + self.horizontalLayout_6.addWidget(self.hashSizeSpinBox) + + + self.verticalLayout_2.addLayout(self.horizontalLayout_6) + + self.horizontalLayout_7 = QHBoxLayout() + self.horizontalLayout_7.setObjectName(u"horizontalLayout_7") + self.label_4 = QLabel(self.groupBox_2) + self.label_4.setObjectName(u"label_4") + self.label_4.setText(u"highfreq_factor") + self.label_4.setAlignment(Qt.AlignRight|Qt.AlignTrailing|Qt.AlignVCenter) + + self.horizontalLayout_7.addWidget(self.label_4) + + self.highfreqFactorSpinBox = QSpinBox(self.groupBox_2) + self.highfreqFactorSpinBox.setObjectName(u"highfreqFactorSpinBox") + sizePolicy1.setHeightForWidth(self.highfreqFactorSpinBox.sizePolicy().hasHeightForWidth()) + self.highfreqFactorSpinBox.setSizePolicy(sizePolicy1) + self.highfreqFactorSpinBox.setMaximum(32) + self.highfreqFactorSpinBox.setValue(4) + + self.horizontalLayout_7.addWidget(self.highfreqFactorSpinBox) + + + self.verticalLayout_2.addLayout(self.horizontalLayout_7) + + + self.horizontalLayout_3.addLayout(self.verticalLayout_2) + + self.horizontalSpacer_3 = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum) + + self.horizontalLayout_3.addItem(self.horizontalSpacer_3) + + self.optionsResetButton = QPushButton(self.groupBox_2) + self.optionsResetButton.setObjectName(u"optionsResetButton") + + self.horizontalLayout_3.addWidget(self.optionsResetButton) + + + self.verticalLayout_3.addWidget(self.groupBox_2) + + self.progressBar = QProgressBar(tabOcr_BuildPHashDatabase) + self.progressBar.setObjectName(u"progressBar") + 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_5 = QHBoxLayout() + self.horizontalLayout_5.setObjectName(u"horizontalLayout_5") + self.horizontalSpacer = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum) + + self.horizontalLayout_5.addItem(self.horizontalSpacer) + + self.buildButton = QPushButton(tabOcr_BuildPHashDatabase) + self.buildButton.setObjectName(u"buildButton") + + self.horizontalLayout_5.addWidget(self.buildButton) + + self.horizontalSpacer_2 = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum) + + self.horizontalLayout_5.addItem(self.horizontalSpacer_2) + + + self.verticalLayout_3.addLayout(self.horizontalLayout_5) + + self.verticalSpacer_2 = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding) + + self.verticalLayout_3.addItem(self.verticalSpacer_2) + + + self.retranslateUi(tabOcr_BuildPHashDatabase) + + QMetaObject.connectSlotsByName(tabOcr_BuildPHashDatabase) + # setupUi + + def retranslateUi(self, tabOcr_BuildPHashDatabase): + self.groupBox.setTitle(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"folders.title", None)) + self.label.setText(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"folders.songDir", None)) + self.label_2.setText(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"folders.charIconDir", None)) + self.groupBox_2.setTitle(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"options.title", None)) + self.optionsResetButton.setText(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"resetButton", None)) + self.buildButton.setText(QCoreApplication.translate("tabOcr_BuildPHashDatabase", u"buildButton", None)) + pass + # retranslateUi + diff --git a/ui/designer/tabs/tabOcrEntry.ui b/ui/designer/tabs/tabOcrEntry.ui index 7314699..cba2b68 100644 --- a/ui/designer/tabs/tabOcrEntry.ui +++ b/ui/designer/tabs/tabOcrEntry.ui @@ -29,6 +29,11 @@ tab.b30 + + + tab.buildPHashDatabase + + @@ -46,6 +51,12 @@
ui.implements.tabs.tabOcr.tabOcr_B30
1 + + TabOcr_BuildPHashDatabase + QWidget +
ui.implements.tabs.tabOcr.tabOcr_BuildPHashDatabase
+ 1 + diff --git a/ui/designer/tabs/tabOcrEntry_ui.py b/ui/designer/tabs/tabOcrEntry_ui.py index 6144e6a..003a04a 100644 --- a/ui/designer/tabs/tabOcrEntry_ui.py +++ b/ui/designer/tabs/tabOcrEntry_ui.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'tabOcrEntry.ui' ## -## Created by: Qt User Interface Compiler version 6.5.1 +## Created by: Qt User Interface Compiler version 6.5.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ @@ -19,6 +19,7 @@ from PySide6.QtWidgets import (QApplication, QSizePolicy, QTabWidget, QVBoxLayou QWidget) from ui.implements.tabs.tabOcr.tabOcr_B30 import TabOcr_B30 +from ui.implements.tabs.tabOcr.tabOcr_BuildPHashDatabase import TabOcr_BuildPHashDatabase from ui.implements.tabs.tabOcr.tabOcr_Device import TabOcr_Device class Ui_TabOcrEntry(object): @@ -37,6 +38,9 @@ class Ui_TabOcrEntry(object): self.tab_2 = TabOcr_B30() self.tab_2.setObjectName(u"tab_2") self.tabWidget.addTab(self.tab_2, "") + self.tab_3 = TabOcr_BuildPHashDatabase() + self.tab_3.setObjectName(u"tab_3") + self.tabWidget.addTab(self.tab_3, "") self.verticalLayout.addWidget(self.tabWidget) @@ -52,6 +56,7 @@ class Ui_TabOcrEntry(object): def retranslateUi(self, TabOcrEntry): self.tabWidget.setTabText(self.tabWidget.indexOf(self.tab), QCoreApplication.translate("TabOcrEntry", u"tab.device", None)) self.tabWidget.setTabText(self.tabWidget.indexOf(self.tab_2), QCoreApplication.translate("TabOcrEntry", u"tab.b30", None)) + self.tabWidget.setTabText(self.tabWidget.indexOf(self.tab_3), QCoreApplication.translate("TabOcrEntry", u"tab.buildPHashDatabase", None)) pass # retranslateUi diff --git a/ui/extends/ocr.py b/ui/extends/ocr/__init__.py similarity index 94% rename from ui/extends/ocr.py rename to ui/extends/ocr/__init__.py index 212ab91..bda9d16 100644 --- a/ui/extends/ocr.py +++ b/ui/extends/ocr/__init__.py @@ -1,3 +1,5 @@ +from .build_phash import build_image_phash_database + try: import json diff --git a/ui/extends/ocr/build_phash.py b/ui/extends/ocr/build_phash.py new file mode 100644 index 0000000..8907738 --- /dev/null +++ b/ui/extends/ocr/build_phash.py @@ -0,0 +1,61 @@ +import sqlite3 +import time +from pathlib import Path +from typing import Any, Callable, Optional + +import cv2 +from arcaea_offline_ocr.phash_db import phash_opencv + + +def build_image_phash_database( + images: list[Path], + labels: list[str], + *, + hash_size: int = 16, + highfreq_factor: int = 4, + progress_func: Optional[Callable[[int, int], Any]] = None, +): + assert len(images) == len(labels) + + conn = sqlite3.connect(":memory:", check_same_thread=False) + + with conn: + cursor = conn.cursor() + + cursor.execute("CREATE TABLE properties (key TEXT, value TEXT)") + cursor.executemany( + "INSERT INTO properties VALUES (?, ?)", + [ + ("hash_size", hash_size), + ("highfreq_factor", highfreq_factor), + ], + ) + + image_num = len(images) + id_hashes = [] + for i, label, image_path in zip(range(image_num), labels, images): + image_hash = phash_opencv( + cv2.imread(str(image_path.resolve()), cv2.IMREAD_GRAYSCALE), + hash_size=hash_size, + highfreq_factor=highfreq_factor, + ) + image_hash_bytes = image_hash.flatten().tobytes() + + id_hashes.append([label, image_hash_bytes]) + if progress_func: + progress_func(i + 1, image_num) + + hash_length = len(id_hashes[0][1]) + cursor.execute(f"CREATE TABLE hashes (id TEXT, hash BLOB({hash_length}))") + + cursor.executemany( + "INSERT INTO hashes VALUES (?, ?)", + id_hashes, + ) + cursor.executemany( + "INSERT INTO properties VALUES (?, ?)", + [("built_timestamp", int(time.time()))], + ) + conn.commit() + + return conn diff --git a/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py b/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py new file mode 100644 index 0000000..e334f17 --- /dev/null +++ b/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py @@ -0,0 +1,138 @@ +import logging +import re +import sqlite3 +import time +from pathlib import Path + +from PySide6.QtCore import QThread, Signal, Slot +from PySide6.QtWidgets import QFileDialog, QMessageBox, QWidget + +from ui.designer.tabs.tabOcr.tabOcr_BuildPHashDatabase_ui import ( + Ui_tabOcr_BuildPHashDatabase, +) +from ui.extends.ocr import build_image_phash_database + +logger = logging.getLogger(__name__) + + +class BuildDatabaseThread(QThread): + conn: sqlite3.Connection + + progress = Signal(int, int) + success = Signal() + error = Signal(str) + finished = Signal() + + def __init__( + self, + images: list[Path], + labels: list[str], + *, + hashSize: int | None = None, + highfreqFactor: int | None = None, + ): + super().__init__() + self.images = images + self.labels = labels + self.hashSize = hashSize + self.highfreqFactor = highfreqFactor + + def run(self): + try: + progressFunc = lambda i, total: self.progress.emit(i, total) + + kwargsDict = {} + if self.hashSize is not None: + kwargsDict["hash_size"] = self.hashSize + if self.highfreqFactor is not None: + kwargsDict["highfreq_factor"] = self.highfreqFactor + self.conn = build_image_phash_database( + self.images, self.labels, progress_func=progressFunc, **kwargsDict + ) + self.success.emit() + except Exception as e: + logger.exception("Error during pHash database build") + self.error.emit(str(e)) + finally: + self.finished.emit() + + +class TabOcr_BuildPHashDatabase(Ui_tabOcr_BuildPHashDatabase, QWidget): + def __init__(self, parent=None): + super().__init__(parent) + self.setupUi(self) + + self.songDirSelector.setMode(self.songDirSelector.getExistingDirectory) + self.charIconDirSelector.setMode(self.charIconDirSelector.getExistingDirectory) + + self.buildButton.clicked.connect(self.databaseBuildStart) + + @Slot() + def on_optionsResetButton_clicked(self): + self.hashSizeSpinBox.setValue(16) + self.highfreqFactorSpinBox.setValue(4) + + def databaseFileName(self): + return f"image-phash-{int(time.time() * 1000)}.db" + + def databaseBuildStart(self): + if not self.songDirSelector.selectedFiles(): + QMessageBox.critical(self, None, "Song directory not selected.") + return + if not self.charIconDirSelector.selectedFiles(): + QMessageBox.critical(self, None, "Char icon directory not selected.") + return + + songDir = self.songDirSelector.selectedFiles()[0] + charIconDir = self.charIconDirSelector.selectedFiles()[0] + + acceptExts = [".jpg", ".png"] + songFilePaths = [ + p for p in Path(songDir).glob("**/*") if p.suffix in acceptExts + ] + charIconFilePaths = [ + p for p in Path(charIconDir).glob("**/*") if p.suffix in acceptExts + ] + songLabels = [re.sub(r"_.*$", "", p.stem) for p in songFilePaths] + charLabels = [f"character||{p.stem}" for p in charIconFilePaths] + + self.databaseBuildThread = BuildDatabaseThread( + songFilePaths + charIconFilePaths, songLabels + charLabels + ) + self.databaseBuildThread.progress.connect(self.databaseBuildProgress) + self.databaseBuildThread.success.connect(self.databaseBuildSuccess) + self.databaseBuildThread.error.connect(self.databaseBuildError) + self.buildButton.setEnabled(False) + self.databaseBuildThread.start() + + @Slot(int, int) + def databaseBuildProgress(self, i: int, total: int): + if i < 5: + self.progressBar.setMaximum(total) + self.progressBar.setValue(i) + + @Slot(str) + def databaseBuildError(self, msg: str): + QMessageBox.critical(self, "Error", msg) + self.databaseBuildCleanUp() + + @Slot() + def databaseBuildSuccess(self): + dbMemory = self.databaseBuildThread.conn + + dbFileName, _ = QFileDialog.getSaveFileName(self, None, self.databaseFileName()) + if not dbFileName: + self.databaseBuildCleanUp() + QMessageBox.information(self, None, "User canceled operation.") + return + + dbDisk = sqlite3.connect(dbFileName) + dbMemory.backup(dbDisk) + self.databaseBuildCleanUp() + + def databaseBuildCleanUp(self): + self.databaseBuildThread.deleteLater() + self.databaseBuildThread = None + self.progressBar.setMaximum(0) + self.progressBar.setValue(0) + self.buildButton.setEnabled(True) diff --git a/ui/resources/lang/en_US.ts b/ui/resources/lang/en_US.ts index 972e3b6..c572a5d 100644 --- a/ui/resources/lang/en_US.ts +++ b/ui/resources/lang/en_US.ts @@ -89,12 +89,12 @@ Continue - + dialog.tryInitExistingDatabase The existing database doesn't seem to be initialized properly, try initialize again? - + dialog.confirmNewDatabase Database file does not exist. Create now? @@ -696,6 +696,11 @@ validation tab.b30 B30 + + + tab.buildPHashDatabase + Build pHash Database + TabOcr_B30 @@ -1167,4 +1172,37 @@ validation Result (play rating) + + tabOcr_BuildPHashDatabase + + + folders.title + Data Folders + + + + folders.songDir + Song jackets + + + + folders.charIconDir + Character icons + + + + options.title + Options + + + + resetButton + Reset + + + + buildButton + Build + + diff --git a/ui/resources/lang/zh_CN.ts b/ui/resources/lang/zh_CN.ts index c95af67..120d025 100644 --- a/ui/resources/lang/zh_CN.ts +++ b/ui/resources/lang/zh_CN.ts @@ -89,12 +89,12 @@ 继续 - + dialog.tryInitExistingDatabase 现有的数据库似乎没有正确初始化，是否尝试再次初始化？ - + dialog.confirmNewDatabase 数据库文件不存在，是否创建？ @@ -695,6 +695,11 @@ tab.b30 B30 + + + tab.buildPHashDatabase + 构建 pHash 数据库 + TabOcr_B30 @@ -1166,4 +1171,37 @@ 结果（单曲 PTT） + + tabOcr_BuildPHashDatabase + + + folders.title + 数据文件夹 + + + + folders.songDir + 曲封 + + + + folders.charIconDir + 搭档头像 + + + + options.title + 选项 + + + + resetButton + 重置 + + + + buildButton + 构建 + +