mirror of
https://github.com/283375/arcaea-offline-pyside-ui.git
synced 2025-07-01 12:26:26 +00:00
wip: b30 ocr tab
This commit is contained in:
161
ui/implements/tabs/tabOcr/tabOcr_B30.py
Normal file
161
ui/implements/tabs/tabOcr/tabOcr_B30.py
Normal file
@ -0,0 +1,161 @@
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
from arcaea_offline_ocr.b30.chieri.v4.ocr import ChieriBotV4Ocr
|
||||
from arcaea_offline_ocr.sift_db import SIFTDatabase
|
||||
from arcaea_offline_ocr.utils import imread_unicode
|
||||
|
||||
# from paddleocr import PaddleOCR
|
||||
from PySide6.QtCore import Signal, Slot
|
||||
from PySide6.QtWidgets import QWidget
|
||||
|
||||
from ui.designer.tabs.tabOcr.tabOcr_B30_ui import Ui_TabOcr_B30
|
||||
from ui.extends.components.ocrQueue import OcrQueueModel
|
||||
from ui.extends.shared.cv2_utils import cv2BgrMatToQImage, qImageToCvMatBgr
|
||||
from ui.extends.tabs.tabOcr.tabOcr_B30 import (
|
||||
ChieriV4OcrRunnable,
|
||||
b30ResultToScoreInsert,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TabOcr_B30(Ui_TabOcr_B30, QWidget):
|
||||
tryPrepareOcr = Signal()
|
||||
|
||||
def __init__(self, parent=None):
|
||||
super().__init__(parent)
|
||||
self.setupUi(self)
|
||||
|
||||
self.b30TypeComboBox.addItem("ChieriV4", "chieri_v4")
|
||||
self.b30TypeComboBox.setCurrentIndex(0)
|
||||
self.b30TypeComboBox.setEnabled(False)
|
||||
|
||||
# self.paddleFolderSelector.setMode(
|
||||
# self.paddleFolderSelector.getExistingDirectory
|
||||
# )
|
||||
|
||||
self.imageSelector.filesSelected.connect(self.imageSelected)
|
||||
self.knnModelSelector.filesSelected.connect(self.knnModelSelected)
|
||||
self.b30KnnModelSelector.filesSelected.connect(self.b30KnnModelSelected)
|
||||
# self.paddleFolderSelector.filesSelected.connect(self.paddleFolderSelected)
|
||||
self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseSelected)
|
||||
|
||||
self.imagePath = None # for checking only
|
||||
self.img = None
|
||||
self.paddleFolder = None
|
||||
self.paddle = None
|
||||
self.knnModel = None
|
||||
self.b30KnnModel = None
|
||||
self.siftDatabase = None
|
||||
|
||||
self.ocr = None
|
||||
|
||||
self.tryPrepareOcr.connect(self.prepareOcr)
|
||||
|
||||
self.ocrQueueModel = OcrQueueModel(self)
|
||||
self.ocrQueue.setModel(self.ocrQueueModel)
|
||||
|
||||
def imageSelected(self):
|
||||
selectedFiles = self.imageSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
imagePath = selectedFiles[0]
|
||||
self.imagePath = imagePath
|
||||
self.img = imread_unicode(imagePath)
|
||||
self.tryPrepareOcr.emit()
|
||||
|
||||
def knnModelSelected(self):
|
||||
selectedFiles = self.knnModelSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
knnModelPath = selectedFiles[0]
|
||||
self.knnModel = cv2.ml.KNearest.load(knnModelPath)
|
||||
self.tryPrepareOcr.emit()
|
||||
|
||||
def b30KnnModelSelected(self):
|
||||
selectedFiles = self.b30KnnModelSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
b30KnnModelPath = selectedFiles[0]
|
||||
self.b30KnnModel = cv2.ml.KNearest.load(b30KnnModelPath)
|
||||
self.tryPrepareOcr.emit()
|
||||
|
||||
def siftDatabaseSelected(self):
|
||||
selectedFiles = self.siftDatabaseSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
siftDatabasePath = selectedFiles[0]
|
||||
self.siftDatabase = SIFTDatabase(siftDatabasePath)
|
||||
self.tryPrepareOcr.emit()
|
||||
|
||||
def paddleFolderSelected(self):
|
||||
selectedFiles = self.paddleFolderSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
self.paddleFolder = selectedFiles[0]
|
||||
self.initPaddle()
|
||||
self.tryPrepareOcr.emit()
|
||||
|
||||
def initPaddle(self):
|
||||
paddleFolder = Path(self.paddleFolder)
|
||||
paddleDetFolder = paddleFolder / "det"
|
||||
paddleClsFolder = paddleFolder / "cls"
|
||||
paddleRecFolder = paddleFolder / "rec"
|
||||
|
||||
if not (paddleDetFolder.exists() and paddleRecFolder.exists()):
|
||||
logger.warning("paddleocr folder incomplete, aborting.")
|
||||
return
|
||||
|
||||
self.paddle = PaddleOCR(
|
||||
show_log=False,
|
||||
use_angle_cls=False,
|
||||
det_model_dir=str(paddleDetFolder),
|
||||
cls_model_dir=str(paddleClsFolder),
|
||||
rec_model_dir=str(paddleRecFolder),
|
||||
)
|
||||
|
||||
def prepareOcr(self):
|
||||
b30Type = self.b30TypeComboBox.currentData()
|
||||
if not b30Type:
|
||||
return
|
||||
|
||||
if b30Type == "chieri_v4":
|
||||
if (
|
||||
not self.imagePath
|
||||
or not self.knnModel
|
||||
or not self.b30KnnModel
|
||||
or not self.siftDatabase
|
||||
):
|
||||
return
|
||||
|
||||
self.ocrQueueModel.clear()
|
||||
|
||||
ocr = ChieriBotV4Ocr(self.knnModel, self.b30KnnModel, self.siftDatabase)
|
||||
ocr.set_factor(self.img)
|
||||
self.ocr = ocr
|
||||
|
||||
roi = ocr.rois
|
||||
for component in roi.components(self.img):
|
||||
qImage = cv2BgrMatToQImage(component.copy())
|
||||
self.ocrQueueModel.addItem(qImage)
|
||||
self.ocrQueue.resizeTableView()
|
||||
|
||||
@Slot()
|
||||
def on_ocr_startButton_clicked(self):
|
||||
if (
|
||||
not self.imagePath
|
||||
or not self.knnModel
|
||||
or not self.b30KnnModel
|
||||
or not self.siftDatabase
|
||||
):
|
||||
return
|
||||
|
||||
for row in range(self.ocrQueueModel.rowCount()):
|
||||
index = self.ocrQueueModel.index(row, 0)
|
||||
qImage = index.data(OcrQueueModel.ImageQImageRole)
|
||||
cv2Mat = qImageToCvMatBgr(qImage)
|
||||
runnable = ChieriV4OcrRunnable(self.ocr, cv2Mat)
|
||||
self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
|
||||
self.ocrQueueModel.setData(
|
||||
index,
|
||||
b30ResultToScoreInsert,
|
||||
OcrQueueModel.ProcessOcrResultFuncRole,
|
||||
)
|
||||
self.ocrQueueModel.startQueue()
|
||||
108
ui/implements/tabs/tabOcr/tabOcr_Device.py
Normal file
108
ui/implements/tabs/tabOcr/tabOcr_Device.py
Normal file
@ -0,0 +1,108 @@
|
||||
import cv2
|
||||
import pytesseract
|
||||
|
||||
# from arcaea_offline_ocr_device_creation_wizard.implements.wizard import Wizard
|
||||
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 PySide6.QtCore import Slot
|
||||
from PySide6.QtWidgets import QFileDialog, QWidget
|
||||
|
||||
from ui.designer.tabs.tabOcr.tabOcr_Device_ui import Ui_TabOcr_Device
|
||||
from ui.extends.components.ocrQueue import OcrQueueModel
|
||||
from ui.extends.settings import Settings
|
||||
from ui.extends.tabs.tabOcr.tabOcr_Device import (
|
||||
ScoreInsertConverter,
|
||||
TabDeviceV2OcrRunnable,
|
||||
)
|
||||
|
||||
|
||||
class TabOcr_Device(Ui_TabOcr_Device, QWidget):
|
||||
def __init__(self, parent=None):
|
||||
super().__init__(parent)
|
||||
self.setupUi(self)
|
||||
self.openWizardButton.setEnabled(False)
|
||||
|
||||
self.deviceComboBox.currentIndexChanged.connect(
|
||||
self.changeDeviceDepStackedWidget
|
||||
)
|
||||
|
||||
self.deviceFileSelector.filesSelected.connect(self.deviceFileSelected)
|
||||
self.knnModelSelector.filesSelected.connect(self.knnModelFileSelected)
|
||||
self.tesseractFileSelector.filesSelected.connect(
|
||||
self.tesseractFileSelectorFilesSelected
|
||||
)
|
||||
self.siftDatabaseSelector.filesSelected.connect(self.siftDatabaseFileSelected)
|
||||
|
||||
settings = Settings()
|
||||
self.deviceFileSelector.selectFile(settings.devicesJsonFile())
|
||||
self.tesseractFileSelector.selectFile(settings.tesseractPath())
|
||||
self.deviceComboBox.selectDevice(settings.deviceUuid())
|
||||
|
||||
self.ocrQueueModel = OcrQueueModel(self)
|
||||
self.ocrQueue.setModel(self.ocrQueueModel)
|
||||
self.ocrQueueProxyModel = self.ocrQueue.tableProxyModel()
|
||||
|
||||
@Slot()
|
||||
def on_openWizardButton_clicked(self):
|
||||
# wizard = Wizard(self)
|
||||
# wizard.open()
|
||||
pass
|
||||
|
||||
def changeDeviceDepStackedWidget(self):
|
||||
device = self.deviceComboBox.currentData()
|
||||
if isinstance(device, (DeviceV1, DeviceV2)):
|
||||
self.deviceDependenciesStackedWidget.setCurrentIndex(device.version - 1)
|
||||
|
||||
def deviceFileSelected(self):
|
||||
selectedFiles = self.deviceFileSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
file = selectedFiles[0]
|
||||
self.deviceComboBox.loadDevicesJson(file)
|
||||
|
||||
def knnModelFileSelected(self):
|
||||
selectedFiles = self.knnModelSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
self.knnModel = cv2.ml.KNearest.load(selectedFiles[0])
|
||||
|
||||
def tesseractFileSelectorFilesSelected(self):
|
||||
selectedFiles = self.tesseractFileSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
pytesseract.pytesseract.tesseract_cmd = selectedFiles[0]
|
||||
|
||||
def siftDatabaseFileSelected(self):
|
||||
selectedFiles = self.siftDatabaseSelector.selectedFiles()
|
||||
if selectedFiles:
|
||||
self.siftDatabase = SIFTDatabase(selectedFiles[0])
|
||||
|
||||
@Slot()
|
||||
def on_ocr_addImageButton_clicked(self):
|
||||
files, _filter = QFileDialog.getOpenFileNames(
|
||||
self, None, "", "Image Files (*.png *.jpg *.jpeg *.bmp *.webp);;*"
|
||||
)
|
||||
for file in files:
|
||||
self.ocrQueueModel.addItem(file)
|
||||
self.ocrQueue.resizeTableView()
|
||||
|
||||
@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)
|
||||
runnable = TabDeviceV2OcrRunnable(
|
||||
imagePath,
|
||||
self.deviceComboBox.currentData(),
|
||||
self.knnModel,
|
||||
self.siftDatabase,
|
||||
)
|
||||
self.ocrQueueModel.setData(index, runnable, OcrQueueModel.OcrRunnableRole)
|
||||
self.ocrQueueModel.setData(
|
||||
index,
|
||||
ScoreInsertConverter.deviceV2,
|
||||
OcrQueueModel.ProcessOcrResultFuncRole,
|
||||
)
|
||||
self.ocrQueueModel.startQueue()
|
||||
|
||||
@Slot()
|
||||
def on_ocr_removeAllButton_clicked(self):
|
||||
self.ocrQueueModel.clear()
|
||||
Reference in New Issue
Block a user