From 94e4d73a950c128e544ead0b74c33005d6836e95 Mon Sep 17 00:00:00 2001 From: 283375 Date: Tue, 10 Oct 2023 01:26:20 +0800 Subject: [PATCH] impr: TabOcr_BuildPHashDatabase --- .../tabs/tabOcr/tabOcr_BuildPHashDatabase.ui | 36 +++++++-- .../tabOcr/tabOcr_BuildPHashDatabase_ui.py | 74 ++++++++++------- ui/extends/ocr/build_phash.py | 23 +++++- .../tabs/tabOcr/tabOcr_BuildPHashDatabase.py | 44 +++++++--- ui/resources/lang/en_US.ts | 81 +++++++++++-------- ui/resources/lang/zh_CN.ts | 81 +++++++++++-------- 6 files changed, 226 insertions(+), 113 deletions(-) diff --git a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui index ac24876..706dcaa 100644 --- a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui +++ b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase.ui @@ -1,7 +1,7 @@ - tabOcr_BuildPHashDatabase - + TabOcr_BuildPHashDatabase + 0 @@ -11,7 +11,7 @@ - tabOcr_BuildPHashDatabase + TabOcr_BuildPHashDatabase @@ -151,6 +151,16 @@ + + + + options.preprocessCharIcon + + + true + + + @@ -175,7 +185,7 @@ - + 0 @@ -186,7 +196,23 @@ Qt::AlignCenter - %v/%m - %p% + [Reading images] %v/%m - %p% + + + + + + + 0 + + + 0 + + + Qt::AlignCenter + + + [Calculate hashes] %v/%m - %p% diff --git a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py index d7ff2bd..608e678 100644 --- a/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py +++ b/ui/designer/tabs/tabOcr/tabOcr_BuildPHashDatabase_ui.py @@ -15,25 +15,25 @@ 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 PySide6.QtWidgets import (QApplication, QCheckBox, 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) +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 = QGroupBox(TabOcr_BuildPHashDatabase) self.groupBox.setObjectName(u"groupBox") self.verticalLayout = QVBoxLayout(self.groupBox) self.verticalLayout.setObjectName(u"verticalLayout") @@ -77,7 +77,7 @@ class Ui_tabOcr_BuildPHashDatabase(object): self.verticalLayout_3.addWidget(self.groupBox) - self.groupBox_2 = QGroupBox(tabOcr_BuildPHashDatabase) + 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") @@ -132,6 +132,12 @@ class Ui_tabOcr_BuildPHashDatabase(object): self.horizontalLayout_3.addLayout(self.verticalLayout_2) + self.preprocessCharIconCheckBox = QCheckBox(self.groupBox_2) + self.preprocessCharIconCheckBox.setObjectName(u"preprocessCharIconCheckBox") + self.preprocessCharIconCheckBox.setChecked(True) + + self.horizontalLayout_3.addWidget(self.preprocessCharIconCheckBox) + self.horizontalSpacer_3 = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum) self.horizontalLayout_3.addItem(self.horizontalSpacer_3) @@ -144,14 +150,21 @@ class Ui_tabOcr_BuildPHashDatabase(object): 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.readImageProgressBar = QProgressBar(TabOcr_BuildPHashDatabase) + self.readImageProgressBar.setObjectName(u"readImageProgressBar") + self.readImageProgressBar.setMaximum(0) + self.readImageProgressBar.setValue(0) + self.readImageProgressBar.setAlignment(Qt.AlignCenter) - self.verticalLayout_3.addWidget(self.progressBar) + self.verticalLayout_3.addWidget(self.readImageProgressBar) + + self.calculateHashProgressBar = QProgressBar(TabOcr_BuildPHashDatabase) + self.calculateHashProgressBar.setObjectName(u"calculateHashProgressBar") + self.calculateHashProgressBar.setMaximum(0) + self.calculateHashProgressBar.setValue(0) + self.calculateHashProgressBar.setAlignment(Qt.AlignCenter) + + self.verticalLayout_3.addWidget(self.calculateHashProgressBar) self.horizontalLayout_5 = QHBoxLayout() self.horizontalLayout_5.setObjectName(u"horizontalLayout_5") @@ -159,7 +172,7 @@ class Ui_tabOcr_BuildPHashDatabase(object): self.horizontalLayout_5.addItem(self.horizontalSpacer) - self.buildButton = QPushButton(tabOcr_BuildPHashDatabase) + self.buildButton = QPushButton(TabOcr_BuildPHashDatabase) self.buildButton.setObjectName(u"buildButton") self.horizontalLayout_5.addWidget(self.buildButton) @@ -176,18 +189,21 @@ class Ui_tabOcr_BuildPHashDatabase(object): self.verticalLayout_3.addItem(self.verticalSpacer_2) - self.retranslateUi(tabOcr_BuildPHashDatabase) + self.retranslateUi(TabOcr_BuildPHashDatabase) - QMetaObject.connectSlotsByName(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)) + 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.preprocessCharIconCheckBox.setText(QCoreApplication.translate("TabOcr_BuildPHashDatabase", u"options.preprocessCharIcon", None)) + self.optionsResetButton.setText(QCoreApplication.translate("TabOcr_BuildPHashDatabase", u"resetButton", None)) + self.readImageProgressBar.setFormat(QCoreApplication.translate("TabOcr_BuildPHashDatabase", u"[Reading images] %v/%m - %p%", None)) + self.calculateHashProgressBar.setFormat(QCoreApplication.translate("TabOcr_BuildPHashDatabase", u"[Calculate hashes] %v/%m - %p%", None)) + self.buildButton.setText(QCoreApplication.translate("TabOcr_BuildPHashDatabase", u"buildButton", None)) pass # retranslateUi diff --git a/ui/extends/ocr/build_phash.py b/ui/extends/ocr/build_phash.py index 8907738..81f8f6c 100644 --- a/ui/extends/ocr/build_phash.py +++ b/ui/extends/ocr/build_phash.py @@ -1,14 +1,29 @@ import sqlite3 import time -from pathlib import Path from typing import Any, Callable, Optional import cv2 +import numpy as np from arcaea_offline_ocr.phash_db import phash_opencv +def preprocess_char_icon(img_gray: cv2.Mat): + h, w = img_gray.shape[:2] + img = cv2.fillPoly( + img_gray, + [ + np.array([[0, 0], [round(w / 2), 0], [0, round(h / 2)]], np.int32), + np.array([[w, 0], [round(w / 2), 0], [w, round(h / 2)]], np.int32), + np.array([[0, h], [round(w / 2), h], [0, round(h / 2)]], np.int32), + np.array([[w, h], [round(w / 2), h], [w, round(h / 2)]], np.int32), + ], + (128), + ) + return img + + def build_image_phash_database( - images: list[Path], + images: list[cv2.Mat], labels: list[str], *, hash_size: int = 16, @@ -33,9 +48,9 @@ def build_image_phash_database( image_num = len(images) id_hashes = [] - for i, label, image_path in zip(range(image_num), labels, images): + for i, label, image in zip(range(image_num), labels, images): image_hash = phash_opencv( - cv2.imread(str(image_path.resolve()), cv2.IMREAD_GRAYSCALE), + image, hash_size=hash_size, highfreq_factor=highfreq_factor, ) diff --git a/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py b/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py index e334f17..7f7953b 100644 --- a/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py +++ b/ui/implements/tabs/tabOcr/tabOcr_BuildPHashDatabase.py @@ -4,13 +4,15 @@ import sqlite3 import time from pathlib import Path +import cv2 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, + Ui_TabOcr_BuildPHashDatabase, ) -from ui.extends.ocr import build_image_phash_database +from ui.extends.ocr.build_phash import build_image_phash_database, preprocess_char_icon +from ui.extends.shared.language import LanguageChangeEventFilter logger = logging.getLogger(__name__) @@ -57,11 +59,14 @@ class BuildDatabaseThread(QThread): self.finished.emit() -class TabOcr_BuildPHashDatabase(Ui_tabOcr_BuildPHashDatabase, QWidget): +class TabOcr_BuildPHashDatabase(Ui_TabOcr_BuildPHashDatabase, QWidget): def __init__(self, parent=None): super().__init__(parent) self.setupUi(self) + self.languageChangeEventFilter = LanguageChangeEventFilter(self) + self.installEventFilter(self.languageChangeEventFilter) + self.songDirSelector.setMode(self.songDirSelector.getExistingDirectory) self.charIconDirSelector.setMode(self.charIconDirSelector.getExistingDirectory) @@ -93,11 +98,30 @@ class TabOcr_BuildPHashDatabase(Ui_tabOcr_BuildPHashDatabase, QWidget): charIconFilePaths = [ p for p in Path(charIconDir).glob("**/*") if p.suffix in acceptExts ] + + self.readImageProgressBar.setMaximum( + len(songFilePaths) + len(charIconFilePaths) + ) + i = 0 + songMats = [] + charIconMats = [] + for image_path in songFilePaths: + songMats.append(cv2.imread(str(image_path.resolve()), cv2.IMREAD_GRAYSCALE)) + i += 1 + self.readImageProgressBar.setValue(i) + for image_path in charIconFilePaths: + mat = cv2.imread(str(image_path.resolve()), cv2.IMREAD_GRAYSCALE) + if self.preprocessCharIconCheckBox.isChecked(): + mat = preprocess_char_icon(mat) + charIconMats.append(mat) + i += 1 + self.readImageProgressBar.setValue(i) + songLabels = [re.sub(r"_.*$", "", p.stem) for p in songFilePaths] - charLabels = [f"character||{p.stem}" for p in charIconFilePaths] + charLabels = [f"partner||{p.stem}" for p in charIconFilePaths] self.databaseBuildThread = BuildDatabaseThread( - songFilePaths + charIconFilePaths, songLabels + charLabels + songMats + charIconMats, songLabels + charLabels ) self.databaseBuildThread.progress.connect(self.databaseBuildProgress) self.databaseBuildThread.success.connect(self.databaseBuildSuccess) @@ -108,8 +132,8 @@ class TabOcr_BuildPHashDatabase(Ui_tabOcr_BuildPHashDatabase, QWidget): @Slot(int, int) def databaseBuildProgress(self, i: int, total: int): if i < 5: - self.progressBar.setMaximum(total) - self.progressBar.setValue(i) + self.calculateHashProgressBar.setMaximum(total) + self.calculateHashProgressBar.setValue(i) @Slot(str) def databaseBuildError(self, msg: str): @@ -133,6 +157,8 @@ class TabOcr_BuildPHashDatabase(Ui_tabOcr_BuildPHashDatabase, QWidget): def databaseBuildCleanUp(self): self.databaseBuildThread.deleteLater() self.databaseBuildThread = None - self.progressBar.setMaximum(0) - self.progressBar.setValue(0) + self.readImageProgressBar.setMaximum(0) + self.readImageProgressBar.setValue(0) + self.calculateHashProgressBar.setMaximum(0) + self.calculateHashProgressBar.setValue(0) self.buildButton.setEnabled(True) diff --git a/ui/resources/lang/en_US.ts b/ui/resources/lang/en_US.ts index c572a5d..9ba841d 100644 --- a/ui/resources/lang/en_US.ts +++ b/ui/resources/lang/en_US.ts @@ -730,6 +730,54 @@ validation Select Image + + TabOcr_BuildPHashDatabase + + + folders.title + Data Folders + + + + folders.songDir + Song jackets + + + + folders.charIconDir + Partner icons + + + + options.title + Options + + + + options.preprocessCharIcon + Preprocess partner icons + + + + resetButton + Reset + + + + [Reading images] %v/%m - %p% + + + + + [Calculate hashes] %v/%m - %p% + + + + + buildButton + Build + + TabOcr_Device @@ -1172,37 +1220,4 @@ 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 120d025..1d6e40f 100644 --- a/ui/resources/lang/zh_CN.ts +++ b/ui/resources/lang/zh_CN.ts @@ -729,6 +729,54 @@ 选择图片 + + TabOcr_BuildPHashDatabase + + + folders.title + 数据文件夹 + + + + folders.songDir + 曲封 + + + + folders.charIconDir + 搭档头像 + + + + options.title + 选项 + + + + options.preprocessCharIcon + 预处理搭档头像 + + + + resetButton + 重置 + + + + [Reading images] %v/%m - %p% + [读取图片] %v/%m - %p% + + + + [Calculate hashes] %v/%m - %p% + [计算哈希] %v/%m - %p% + + + + buildButton + 构建 + + TabOcr_Device @@ -1171,37 +1219,4 @@ 结果（单曲 PTT） - - tabOcr_BuildPHashDatabase - - - folders.title - 数据文件夹 - - - - folders.songDir - 曲封 - - - - folders.charIconDir - 搭档头像 - - - - options.title - 选项 - - - - resetButton - 重置 - - - - buildButton - 构建 - -