wip: ImagePHashDatabase & SizesV2

This commit is contained in:
2023-09-27 17:16:33 +08:00
parent d701055c74
commit 9c06c6d9f1
7 changed files with 363 additions and 315 deletions

View File

@ -2,6 +2,7 @@ import logging
import cv2
from arcaea_offline_ocr.b30.chieri.v4.ocr import ChieriBotV4Ocr
from arcaea_offline_ocr.phash_db import ImagePHashDatabase
from arcaea_offline_ocr.sift_db import SIFTDatabase
from arcaea_offline_ocr.utils import imread_unicode
from PySide6.QtCore import Signal, Slot
@ -34,10 +35,11 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
self.imageSelector.filesSelected.connect(self.imageSelected)
self.knnModelSelector.filesSelected.connect(self.knnModelSelected)
self.b30KnnModelSelector.filesSelected.connect(self.b30KnnModelSelected)
self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseSelected)
# self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseSelected)
self.phashDatabaseSelector.filesSelected.connect(self.phashDatabaseSelected)
self.knnModelSelector.connectSettings(KNN_MODEL_FILE)
self.siftDatabaseSelector.connectSettings(SIFT_DATABASE_FILE)
# self.siftDatabaseSelector.connectSettings(SIFT_DATABASE_FILE)
self.imagePath = None # for checking only
self.img = None
@ -45,7 +47,8 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
self.paddle = None
self.knnModel = None
self.b30KnnModel = None
self.siftDatabase = None
# self.siftDatabase = None
self.phashDatabase = None
self.ocr = None
@ -54,7 +57,7 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
settings = Settings()
logger.info("Applying default settings...")
self.knnModelSelector.selectFile(settings.knnModelFile())
self.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())
# self.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())
self.ocrQueueModel = OcrQueueModel(self)
self.ocrQueue.setModel(self.ocrQueueModel)
@ -84,6 +87,12 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
self.siftDatabase = SIFTDatabase(siftDatabasePath)
self.tryPrepareOcr.emit()
def phashDatabaseSelected(self):
if selectedFiles := self.phashDatabaseSelector.selectedFiles():
phashDatabasePath = selectedFiles[0]
self.phashDatabase = ImagePHashDatabase(phashDatabasePath)
self.tryPrepareOcr.emit()
def prepareOcr(self):
b30Type = self.b30TypeComboBox.currentData()
if not b30Type:
@ -94,13 +103,13 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
not self.imagePath
or not self.knnModel
or not self.b30KnnModel
or not self.siftDatabase
or not self.phashDatabase
):
return
self.ocrQueueModel.clear()
ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.siftDatabase)
ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.phashDatabase)
ocr.set_factor(self.img)
self.ocr = ocr
@ -116,7 +125,7 @@ class TabOcr_B30(Ui_TabOcr_B30, QWidget):
not self.imagePath
or not self.knnModel
or not self.b30KnnModel
or not self.siftDatabase
or not self.phashDatabase
):
return

View File

@ -6,6 +6,7 @@ import cv2
from arcaea_offline_ocr.device.v1.definition import DeviceV1
from arcaea_offline_ocr.device.v2.definition import DeviceV2
from arcaea_offline_ocr.sift_db import SIFTDatabase
from arcaea_offline_ocr.phash_db import ImagePHashDatabase
from PySide6.QtCore import Qt, Slot
from PySide6.QtWidgets import QApplication, QFileDialog, QWidget
@ -38,11 +39,12 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
self.deviceFileSelector.filesSelected.connect(self.deviceFileSelected)
self.knnModelSelector.filesSelected.connect(self.knnModelFileSelected)
self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseFileSelected)
# self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseFileSelected)
self.phashDatabaseSelector.filesSelected.connect(self.phashDatabaseFileSelected)
self.deviceFileSelector.connectSettings(DEVICES_JSON_FILE)
self.knnModelSelector.connectSettings(KNN_MODEL_FILE)
self.siftDatabaseSelector.connectSettings(SIFT_DATABASE_FILE)
# self.siftDatabaseSelector.connectSettings(SIFT_DATABASE_FILE)
settings = Settings()
logger.info("Applying default settings...")
@ -50,7 +52,7 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
self.tesseractFileSelector.selectFile(settings.tesseractPath())
self.deviceComboBox.selectDevice(settings.deviceUuid())
self.knnModelSelector.selectFile(settings.knnModelFile())
self.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())
# vself.siftDatabaseSelector.selectFile(settings.siftDatabaseFile())
self.ocrQueueModel = OcrQueueModel(self)
self.ocrQueue.setModel(self.ocrQueueModel)
@ -96,6 +98,10 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
if selectedFiles := self.siftDatabaseSelector.selectedFiles():
self.siftDatabase = SIFTDatabase(selectedFiles[0])
def phashDatabaseFileSelected(self):
if selectedFiles := self.phashDatabaseSelector.selectedFiles():
self.phashDatabase = ImagePHashDatabase(selectedFiles[0])
@Slot()
def on_ocr_addImageButton_clicked(self):
files, _filter = QFileDialog.getOpenFileNames(
@ -115,14 +121,16 @@ class TabOcr_Device(Ui_TabOcr_Device, QWidget):
runnable = TabDeviceV2AutoRoisOcrRunnable(
imagePath,
self.knnModel,
self.siftDatabase,
self.phashDatabase,
sizesV2=self.deviceSizesV2CheckBox.isChecked(),
)
else:
runnable = TabDeviceV2OcrRunnable(
imagePath,
self.deviceComboBox.currentData(),
self.knnModel,
self.siftDatabase,
self.phashDatabase,
sizesV2=self.deviceSizesV2CheckBox.isChecked(),
)
self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
self.ocrQueueModel.setData(