From 04b2985230ea7d3f5f54f5494cb9e339a5f6ee2e Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Wed, 18 Sep 2024 22:07:59 -0400 Subject: [PATCH 1/6] refactor: Add video settings functionality and UI --- camera_view.py | 30 ++++++-- main.py | 25 ++++++ mainwindow.ui | 10 +++ scoresight.spec | 2 + source_view.py | 16 ++-- ui_mainwindow.py | 7 ++ ui_video_settings.py | 118 +++++++++++++++++++++++++++++ video_settings.py | 108 ++++++++++++++++++++++++++ video_settings.ui | 177 +++++++++++++++++++++++++++++++++++++++++++ 9 files changed, 482 insertions(+), 11 deletions(-) create mode 100644 ui_video_settings.py create mode 100644 video_settings.py create mode 100644 video_settings.ui diff --git a/camera_view.py b/camera_view.py index 7fa6f7e..88bf820 100644 --- a/camera_view.py +++ b/camera_view.py @@ -190,9 +190,9 @@ def connect_video_capture(self) -> bool: # attempt to set the highest resolution # check if camera index is a OpenCV camera index - if self.camera_info.type == CameraInfo.CameraType.OPENCV: - # make sure to open the camera at the highest resolution - set_camera_highest_resolution(self.video_capture) + # if self.camera_info.type == CameraInfo.CameraType.OPENCV: + # # make sure to open the camera at the highest resolution + # set_camera_highest_resolution(self.video_capture) return True @@ -400,7 +400,11 @@ def toggleStabilization(self, state): class CameraView(QGraphicsView): first_frame_received_signal = Signal() - def __init__(self, camera_index, detectionTargetsStorage=None): + def __init__( + self, + camera_index: CameraInfo, + detectionTargetsStorage: TextDetectionTargetMemoryStorage | None = None, + ): super().__init__() self.scene = QGraphicsScene(self) self.setScene(self.scene) @@ -415,6 +419,14 @@ def __init__(self, camera_index, detectionTargetsStorage=None): self.fps_text = None self.error_text = None self.showOSD = True + self.camera_width = 0 + self.camera_height = 0 + + def getCameraCapture(self): + return self.timerThread.video_capture + + def getCameraInfo(self): + return self.timerThread.camera_info def setUpdateOnChange(self, updateOnChange): self.timerThread.updateOnChange = updateOnChange @@ -473,18 +485,24 @@ def update_pixmap(self, frame): if not self.firstFrameReceived: self.firstFrameReceived = True + self.camera_width = self.timerThread.video_capture.get( + cv2.CAP_PROP_FRAME_WIDTH + ) + self.camera_height = self.timerThread.video_capture.get( + cv2.CAP_PROP_FRAME_HEIGHT + ) self.first_frame_received_signal.emit() # update the fps text - fps_text = f"Frames/s: {self.timerThread.fps:.2f}\nUpdates/s: {self.timerThread.ups:.2f}\nPreviews/s: {self.timerThread.pps:.2f}" + fps_text = f"Frames/s: {self.timerThread.fps:.2f}\nUpdates/s: {self.timerThread.ups:.2f}\nPreviews/s: {self.timerThread.pps:.2f}\nResolution: {int(self.camera_width)}x{int(self.camera_height)}" if self.fps_text is None: self.fps_text = self.scene.addText(fps_text) self.fps_text.setPos(0, 0) self.fps_text.setZValue(2) self.fps_text.setDefaultTextColor(Qt.GlobalColor.white) - # scale the text according to the view size so its always the same size else: self.fps_text.setPlainText(fps_text) + # scale the text according to the view size so its always the same size self.fps_text.setScale(0.002 * self.scenePixmapItem.boundingRect().width()) def error_event(self, error): diff --git a/main.py b/main.py index c1bcb51..dfa9261 100644 --- a/main.py +++ b/main.py @@ -65,6 +65,7 @@ from update_check import check_for_updates from log_view import LogViewerDialog import file_output +from video_settings import VideoSettingsDialog from vmix_output import VMixAPI from ui_mainwindow import Ui_MainWindow from ui_about import Ui_Dialog as Ui_About @@ -201,6 +202,9 @@ def __init__(self, translator: QTranslator, parent: QObject): self.ui.toolButton_addBox.clicked.connect(self.addBox) self.ui.toolButton_removeBox.clicked.connect(self.removeCustomBox) + self.video_settings_dialog = None + self.ui.toolButton_videoSettings.clicked.connect(self.openVideoSettings) + self.obs_websocket_client = None ocr_models = [ @@ -379,6 +383,24 @@ def __init__(self, translator: QTranslator, parent: QObject): self.get_sources.connect(self.getSources) self.get_sources.emit() + def openVideoSettings(self): + # only allow opening the video settings for an OpenCV type source + if not self.image_viewer or not self.image_viewer.getCameraCapture(): + return + + if self.image_viewer.getCameraInfo().type != CameraInfo.CameraType.OPENCV: + return + + if self.video_settings_dialog is None: + # open the logs dialog + self.video_settings_dialog = VideoSettingsDialog() + self.video_settings_dialog.setWindowTitle("Video Settings") + + self.video_settings_dialog.init_ui(self.image_viewer.getCameraCapture()) + + # show the dialog, non modal + self.video_settings_dialog.show() + def rotateImage(self): # store the rotation in the scoresight.json rotation = fetch_data("scoresight.json", "rotation", 0) @@ -1209,6 +1231,9 @@ def sourceSelectionSucessful(self): self.detectionTargetsStorage, self.itemSelected, ) + self.ui.toolButton_videoSettings.setEnabled( + camera_info.type == CameraInfo.CameraType.OPENCV + ) self.ui.pushButton_fourCorner.setEnabled(True) self.ui.pushButton_binary.setEnabled(True) self.ui.pushButton_fourCorner.toggled.connect( diff --git a/mainwindow.ui b/mainwindow.ui index 0933569..293c227 100644 --- a/mainwindow.ui +++ b/mainwindow.ui @@ -1727,6 +1727,16 @@ + + + + false + + + Video Settings + + + diff --git a/scoresight.spec b/scoresight.spec index bf23955..0ff8ff9 100644 --- a/scoresight.spec +++ b/scoresight.spec @@ -66,6 +66,7 @@ sources = [ 'storage.py', 'tesseract.py', 'text_detection_target.py', + 'video_settings.py', 'ui_about.py', 'ui_connect_obs.py', 'ui_log_view.py', @@ -73,6 +74,7 @@ sources = [ 'ui_screen_capture.py', 'ui_update_available.py', 'ui_url_source.py', + 'ui_video_settings.py', 'update_check.py', 'vmix_output.py', ] diff --git a/source_view.py b/source_view.py index 9f1c2c7..102d64a 100644 --- a/source_view.py +++ b/source_view.py @@ -6,8 +6,14 @@ QGraphicsRectItem, ) +from camera_info import CameraInfo from camera_view import CameraView -from storage import fetch_data, remove_data, store_data +from storage import ( + TextDetectionTargetMemoryStorage, + fetch_data, + remove_data, + store_data, +) from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult from sc_logging import logger from resizable_rect import ResizableRectWithNameTypeAndResult @@ -37,10 +43,10 @@ def angle_from_center(point): class ImageViewer(CameraView): def __init__( self, - camera_index, - fourCornersAppliedCallback, - detectionTargetsStorage, - itemSelectedCallback, + camera_index: CameraInfo, + fourCornersAppliedCallback: callable, + detectionTargetsStorage: TextDetectionTargetMemoryStorage | None, + itemSelectedCallback: callable, ): super().__init__(camera_index, detectionTargetsStorage) self.setMouseTracking(True) diff --git a/ui_mainwindow.py b/ui_mainwindow.py index d19a552..13231fc 100644 --- a/ui_mainwindow.py +++ b/ui_mainwindow.py @@ -898,6 +898,12 @@ def setupUi(self, MainWindow): self.horizontalLayout_3.addWidget(self.comboBox_camera_source) + self.toolButton_videoSettings = QToolButton(self.widget_3) + self.toolButton_videoSettings.setObjectName(u"toolButton_videoSettings") + self.toolButton_videoSettings.setEnabled(False) + + self.horizontalLayout_3.addWidget(self.toolButton_videoSettings) + self.pushButton_refresh_sources = QToolButton(self.widget_3) self.pushButton_refresh_sources.setObjectName(u"pushButton_refresh_sources") @@ -1222,6 +1228,7 @@ def retranslateUi(self, MainWindow): self.comboBox_camera_source.setItemText(2, QCoreApplication.translate("MainWindow", u"URL Source (HTTP, RTSP)", None)) self.comboBox_camera_source.setItemText(3, QCoreApplication.translate("MainWindow", u"Screen Capture", None)) + self.toolButton_videoSettings.setText(QCoreApplication.translate("MainWindow", u"Video Settings", None)) #if QT_CONFIG(tooltip) self.pushButton_refresh_sources.setToolTip(QCoreApplication.translate("MainWindow", u"Refresh Sources", None)) #endif // QT_CONFIG(tooltip) diff --git a/ui_video_settings.py b/ui_video_settings.py new file mode 100644 index 0000000..ad28b3b --- /dev/null +++ b/ui_video_settings.py @@ -0,0 +1,118 @@ +# -*- coding: utf-8 -*- + +################################################################################ +## Form generated from reading UI file 'video_settings.ui' +## +## Created by: Qt User Interface Compiler version 6.7.1 +## +## 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 (QAbstractButton, QApplication, QComboBox, QDialog, + QDialogButtonBox, QFormLayout, QGridLayout, QLabel, + QSizePolicy, QSpinBox, QWidget) + +class Ui_Dialog(object): + def setupUi(self, Dialog): + if not Dialog.objectName(): + Dialog.setObjectName(u"Dialog") + Dialog.resize(400, 153) + self.gridLayout = QGridLayout(Dialog) + self.gridLayout.setObjectName(u"gridLayout") + self.buttonBox = QDialogButtonBox(Dialog) + self.buttonBox.setObjectName(u"buttonBox") + self.buttonBox.setOrientation(Qt.Horizontal) + self.buttonBox.setStandardButtons(QDialogButtonBox.Cancel|QDialogButtonBox.Ok) + + self.gridLayout.addWidget(self.buttonBox, 2, 0, 1, 1) + + self.widget = QWidget(Dialog) + self.widget.setObjectName(u"widget") + self.formLayout = QFormLayout(self.widget) + self.formLayout.setObjectName(u"formLayout") + self.label = QLabel(self.widget) + self.label.setObjectName(u"label") + + self.formLayout.setWidget(1, QFormLayout.LabelRole, self.label) + + self.comboBox_fourcc = QComboBox(self.widget) + self.comboBox_fourcc.addItem("") + self.comboBox_fourcc.setObjectName(u"comboBox_fourcc") + + self.formLayout.setWidget(1, QFormLayout.FieldRole, self.comboBox_fourcc) + + self.label_2 = QLabel(self.widget) + self.label_2.setObjectName(u"label_2") + + self.formLayout.setWidget(2, QFormLayout.LabelRole, self.label_2) + + self.spinBox_fps = QSpinBox(self.widget) + self.spinBox_fps.setObjectName(u"spinBox_fps") + self.spinBox_fps.setMinimum(1) + self.spinBox_fps.setMaximum(60) + self.spinBox_fps.setValue(30) + + self.formLayout.setWidget(2, QFormLayout.FieldRole, self.spinBox_fps) + + self.label_3 = QLabel(self.widget) + self.label_3.setObjectName(u"label_3") + + self.formLayout.setWidget(3, QFormLayout.LabelRole, self.label_3) + + self.comboBox_resolution = QComboBox(self.widget) + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.addItem("") + self.comboBox_resolution.setObjectName(u"comboBox_resolution") + + self.formLayout.setWidget(3, QFormLayout.FieldRole, self.comboBox_resolution) + + + self.gridLayout.addWidget(self.widget, 1, 0, 1, 1) + + + self.retranslateUi(Dialog) + self.buttonBox.accepted.connect(Dialog.accept) + self.buttonBox.rejected.connect(Dialog.reject) + + QMetaObject.connectSlotsByName(Dialog) + # setupUi + + def retranslateUi(self, Dialog): + Dialog.setWindowTitle(QCoreApplication.translate("Dialog", u"Dialog", None)) + self.label.setText(QCoreApplication.translate("Dialog", u"Video Format", None)) + self.comboBox_fourcc.setItemText(0, QCoreApplication.translate("Dialog", u"Default", None)) + + self.label_2.setText(QCoreApplication.translate("Dialog", u"FPS", None)) + self.label_3.setText(QCoreApplication.translate("Dialog", u"Resolution", None)) + self.comboBox_resolution.setItemText(0, QCoreApplication.translate("Dialog", u"Default", None)) + self.comboBox_resolution.setItemText(1, QCoreApplication.translate("Dialog", u"320x240", None)) + self.comboBox_resolution.setItemText(2, QCoreApplication.translate("Dialog", u"640x480", None)) + self.comboBox_resolution.setItemText(3, QCoreApplication.translate("Dialog", u"800x600", None)) + self.comboBox_resolution.setItemText(4, QCoreApplication.translate("Dialog", u"1280x720", None)) + self.comboBox_resolution.setItemText(5, QCoreApplication.translate("Dialog", u"1280x960", None)) + self.comboBox_resolution.setItemText(6, QCoreApplication.translate("Dialog", u"1600x900", None)) + self.comboBox_resolution.setItemText(7, QCoreApplication.translate("Dialog", u"1980x1080", None)) + self.comboBox_resolution.setItemText(8, QCoreApplication.translate("Dialog", u"3840x2160", None)) + self.comboBox_resolution.setItemText(9, QCoreApplication.translate("Dialog", u"2560x1440", None)) + self.comboBox_resolution.setItemText(10, QCoreApplication.translate("Dialog", u"2048x1080", None)) + self.comboBox_resolution.setItemText(11, QCoreApplication.translate("Dialog", u"4096x2160", None)) + + # retranslateUi + diff --git a/video_settings.py b/video_settings.py new file mode 100644 index 0000000..d0b82ba --- /dev/null +++ b/video_settings.py @@ -0,0 +1,108 @@ +import cv2 + +from PySide6.QtWidgets import QDialog +from sc_logging import logger +from ui_video_settings import Ui_Dialog as Ui_VideoSettingsDialog +from storage import store_data + + +def fourcc_to_str(fourcc): + if fourcc == 0: + return "NULL" + if type(fourcc) == str: + return fourcc + if type(fourcc) == bytes: + return fourcc.decode() + if type(fourcc) != int: + fourcc = int(fourcc) + return "".join([chr((fourcc >> 8 * i) & 0xFF) for i in range(4)]) + + +def get_supported_fourcc(cap): + fourcc_list = [ + cv2.VideoWriter_fourcc(*"MJPG"), + cv2.VideoWriter_fourcc(*"YUYV"), + cv2.VideoWriter_fourcc(*"YUY2"), + cv2.VideoWriter_fourcc(*"YU12"), + cv2.VideoWriter_fourcc(*"YV12"), + cv2.VideoWriter_fourcc(*"RGB3"), + cv2.VideoWriter_fourcc(*"H264"), + cv2.VideoWriter_fourcc(*"X264"), + cv2.VideoWriter_fourcc(*"XVID"), + cv2.VideoWriter_fourcc(*"MPEG"), + cv2.VideoWriter_fourcc(*"NV12"), + cv2.VideoWriter_fourcc(*"I420"), + ] + + original_fourcc = int(cap.get(cv2.CAP_PROP_FOURCC)) & 0xFFFFFFFF + supported_fourccs = [fourcc_to_str(original_fourcc)] + + for fourcc in fourcc_list: + cap.set(cv2.CAP_PROP_FOURCC, fourcc) + actual_fourcc = int(cap.get(cv2.CAP_PROP_FOURCC)) + if actual_fourcc == fourcc: + fourcc_str = fourcc_to_str(fourcc) + if fourcc_str not in supported_fourccs: + supported_fourccs.append(fourcc_str) + + cap.set(cv2.CAP_PROP_FOURCC, original_fourcc) + + return supported_fourccs + + +class VideoSettingsDialog(QDialog): + def __init__(self): + super().__init__() + self.ui = Ui_VideoSettingsDialog() + self.ui.setupUi(self) + + def init_ui(self, cap): + self.cap = cap + self.ui.buttonBox.accepted.connect(self.save_settings) + self.ui.buttonBox.rejected.connect(self.close) + self.ui.spinBox_fps.setValue(int(cap.get(cv2.CAP_PROP_FPS))) + self.ui.comboBox_resolution.setCurrentText( + f"{cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}" + ) + self.populate_supported_fourcc() + + def populate_supported_fourcc(self): + supported_fourccs = get_supported_fourcc(self.cap) + self.ui.comboBox_fourcc.clear() + self.ui.comboBox_fourcc.addItems(supported_fourccs) + + def save_settings(self): + fps = self.ui.spinBox_fps.value() + resolution = self.ui.comboBox_resolution.currentText() + if "x" in resolution: + resolution = resolution.split("x") + width = int(resolution[0]) + height = int(resolution[1]) + else: + width = -1 + height = -1 + fourcc = self.ui.comboBox_fourcc.currentText() + store_data( + "scoresight.json", + "video_settings", + { + "fps": fps, + "width": width, + "height": height, + "fourcc": fourcc, + }, + ) + self.cap.set(cv2.CAP_PROP_FPS, fps) + if width > 0: + self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) + if height > 0: + self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) + fourcc_int = int.from_bytes(fourcc.encode(), "little") + self.cap.set(cv2.CAP_PROP_FOURCC, fourcc_int) + print(f"FPS: {fps}, Resolution: {width}x{height}, FourCC: {fourcc}") + logger.info(f"FPS: {fps}, Resolution: {width}x{height}, FourCC: {fourcc}") + print( + f"FPS: {self.cap.get(cv2.CAP_PROP_FPS)}, Resolution: {self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}, FourCC: {fourcc_to_str(self.cap.get(cv2.CAP_PROP_FOURCC))}" + ) + + self.close() diff --git a/video_settings.ui b/video_settings.ui new file mode 100644 index 0000000..e5a3213 --- /dev/null +++ b/video_settings.ui @@ -0,0 +1,177 @@ + + + Dialog + + + + 0 + 0 + 400 + 153 + + + + Dialog + + + + + + Qt::Horizontal + + + QDialogButtonBox::Cancel|QDialogButtonBox::Ok + + + + + + + + + + Video Format + + + + + + + + Default + + + + + + + + FPS + + + + + + + 1 + + + 60 + + + 30 + + + + + + + Resolution + + + + + + + + Default + + + + + 320x240 + + + + + 640x480 + + + + + 800x600 + + + + + 1280x720 + + + + + 1280x960 + + + + + 1600x900 + + + + + 1980x1080 + + + + + 3840x2160 + + + + + 2560x1440 + + + + + 2048x1080 + + + + + 4096x2160 + + + + + + + + + + + + + buttonBox + accepted() + Dialog + accept() + + + 248 + 254 + + + 157 + 274 + + + + + buttonBox + rejected() + Dialog + reject() + + + 316 + 260 + + + 286 + 274 + + + + + From 75774849b02c0e3377f47cfe6bf36a293be017de Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Thu, 19 Sep 2024 11:02:59 -0400 Subject: [PATCH 2/6] refactor: Update opencv-python and pyinstaller versions --- camera_thread.py | 452 +++++++++++ camera_view.py | 393 +-------- main.py | 1640 ++------------------------------------ mainwindow.py | 1549 +++++++++++++++++++++++++++++++++++ requirements.txt | 4 +- scoresight.spec | 2 + screen_capture_source.py | 22 + storage.py | 2 +- ui_about.py | 2 +- ui_connect_obs.py | 2 +- ui_log_view.py | 2 +- ui_mainwindow.py | 2 +- ui_screen_capture.py | 2 +- ui_update_available.py | 2 +- ui_url_source.py | 2 +- ui_video_settings.py | 51 +- video_settings.py | 146 +++- video_settings.ui | 53 +- 18 files changed, 2285 insertions(+), 2043 deletions(-) create mode 100644 camera_thread.py create mode 100644 mainwindow.py diff --git a/camera_thread.py b/camera_thread.py new file mode 100644 index 0000000..39a5b22 --- /dev/null +++ b/camera_thread.py @@ -0,0 +1,452 @@ +from PySide6.QtCore import QThread, Signal + +import platform +import time +import cv2 +from datetime import datetime +import threading + +from camera_info import CameraInfo +from ndi import NDICapture +from screen_capture_source import ScreenCapture, ScreenCaptureType +from storage import TextDetectionTargetMemoryStorage, subscribe_to_data, fetch_data +from tesseract import TextDetector +from text_detection_target import TextDetectionTargetWithResult +from sc_logging import logger +from frame_stabilizer import FrameStabilizer + + +# Function to set the resolution +def set_resolution(cap, width, height): + cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) + + +# Function to get the resolution +def get_resolution(cap): + width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) + height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) + return width, height + + +# Function to set the camera to the highest resolution +def set_camera_highest_resolution(cap): + # List of common resolutions to try + resolutions = [(1920, 1080), (1280, 720), (1024, 768), (800, 600), (640, 480)] + + # grab one frame to make sure the camera is initialized + ret, _ = cap.read() + if not ret: + logger.warn("Error: camera not initialized") + return + + # grab the current resolution + width, height = get_resolution(cap) + + # If the current resolution is already the highest, return + if width >= resolutions[0][0] and height >= resolutions[0][1]: + logger.info( + "Camera is already at the highest resolution: %d x %d", width, height + ) + return + + # Try each resolution and select the highest one that works + highest_res = resolutions[0] + for resolution in resolutions: + logger.debug("Trying camera resolution: %d x %d", resolution[0], resolution[1]) + set_resolution(cap, *resolution) + test_width, test_height = get_resolution(cap) + logger.debug("Found camera resolution: %d x %d", test_width, test_height) + # if found resolution is within close range of the target resolution, use it + if ( + abs(test_width - resolution[0]) < 100 + and abs(test_height - resolution[1]) < 100 + ): + logger.debug( + "Camera highest resolution set to: %d x %d", test_width, test_height + ) + highest_res = (test_width, test_height) + break + + # Set the highest resolution + logger.info("Setting camera resolution to: %d x %d", highest_res[0], highest_res[1]) + set_resolution(cap, *highest_res) + + +class FrameCropAndRotation: + def __init__(self): + self.isCropSet = fetch_data("scoresight.json", "crop_mode", False) + self.cropTop = fetch_data("scoresight.json", "top_crop", 0) + self.cropBottom = fetch_data("scoresight.json", "bottom_crop", 0) + self.cropLeft = fetch_data("scoresight.json", "left_crop", 0) + self.cropRight = fetch_data("scoresight.json", "right_crop", 0) + self.rotation = fetch_data("scoresight.json", "rotation", 0) + subscribe_to_data("scoresight.json", "crop_mode", self.setCropMode) + subscribe_to_data("scoresight.json", "top_crop", self.setCropTop) + subscribe_to_data("scoresight.json", "bottom_crop", self.setCropBottom) + subscribe_to_data("scoresight.json", "left_crop", self.setCropLeft) + subscribe_to_data("scoresight.json", "right_crop", self.setCropRight) + subscribe_to_data("scoresight.json", "rotation", self.setRotation) + + def setCropMode(self, crop_mode): + self.isCropSet = crop_mode + + def setCropTop(self, crop_top): + self.cropTop = crop_top + + def setCropBottom(self, crop_bottom): + self.cropBottom = crop_bottom + + def setCropLeft(self, crop_left): + self.cropLeft = crop_left + + def setCropRight(self, crop_right): + self.cropRight = crop_right + + def setRotation(self, rotation): + self.rotation = rotation + + +class OpenCVVideoCaptureWithSettings: + def __init__(self, capture_id: int | str, capture_backend: int = cv2.CAP_ANY): + self.video_capture = cv2.VideoCapture(capture_id, capture_backend) + self.video_capture_mutex = threading.Lock() + self.capture_id = capture_id + self.fps = fetch_data("scoresight.json", "fps", 30) + self.width = fetch_data("scoresight.json", "width", 0) + self.height = fetch_data("scoresight.json", "height", 0) + self.fourcc = fetch_data("scoresight.json", "fourcc", "MJPG") + self.backend = fetch_data("scoresight.json", "backend", cv2.CAP_ANY) + subscribe_to_data("scoresight.json", "fps", self.setFps) + subscribe_to_data("scoresight.json", "width", self.setWidth) + subscribe_to_data("scoresight.json", "height", self.setHeight) + subscribe_to_data("scoresight.json", "fourcc", self.setFourcc) + subscribe_to_data("scoresight.json", "backend", self.setBackend) + + def setFps(self, fps): + self.fps = fps + self.set(cv2.CAP_PROP_FPS, fps) + + def setWidth(self, width): + self.width = width + self.set(cv2.CAP_PROP_FRAME_WIDTH, width) + + def setHeight(self, height): + self.height = height + self.set(cv2.CAP_PROP_FRAME_HEIGHT, height) + + def setFourcc(self, fourcc: str): + self.fourcc = fourcc + self.set(cv2.CAP_PROP_FOURCC, int.from_bytes(fourcc.encode(), "little")) + + def setBackend(self, backend): + self.backend = backend + self.release() + with self.video_capture_mutex: + print("Setting backend to", backend) + self.video_capture = cv2.VideoCapture(self.capture_id, self.backend) + + def isOpened(self): + with self.video_capture_mutex: + return self.video_capture.isOpened() + + def read(self): + with self.video_capture_mutex: + return self.video_capture.read() + + def release(self): + with self.video_capture_mutex: + self.video_capture.release() + + def set(self, propId, value): + with self.video_capture_mutex: + self.video_capture.set(propId, value) + + def get(self, propId): + with self.video_capture_mutex: + return self.video_capture.get(propId) + + +class TimerThread(QThread): + update_signal = Signal(object) + update_error = Signal(object) + ocr_result_signal = Signal(list) + + def __init__( + self, + camera_info: CameraInfo, + detectionTargetsStorage: TextDetectionTargetMemoryStorage, + ): + super().__init__() + self.camera_info = camera_info + self.homography = None + self.detectionTargetsStorage = detectionTargetsStorage + self.textDetector = TextDetector() # initialize tesseract + self.show_binary = False + self.retry_count = 0 + self.retry_count_max = 50 + self.retry_high_water_mark = 25 + self.stabilizationEnabled = False + self.framestabilizer = FrameStabilizer() + self.video_capture: ( + OpenCVVideoCaptureWithSettings | NDICapture | ScreenCaptureType | None + ) = None + self.should_stop = False + self.frame_interval = 30 + self.update_frame_interval = 1000 / fetch_data( + "scoresight.json", "detection_cadence", 5 + ) + subscribe_to_data( + "scoresight.json", "detection_cadence", self.setUpdateFrameInterval + ) + self.preview_frame_interval = 1000 + self.fps = 1000 / self.frame_interval # frames per second + self.pps = 1000 / self.preview_frame_interval # previews per second + self.ups = 1000 / self.update_frame_interval # updates per second + self.fps_alpha = 0.1 # Smoothing factor + self.updateOnChange = True + self.crop = FrameCropAndRotation() + + def setUpdateFrameInterval(self, cadence): + self.update_frame_interval = 1000 / cadence + + def connect_video_capture(self) -> bool: + if self.camera_info.type == CameraInfo.CameraType.NDI: + self.video_capture = NDICapture(self.camera_info.uuid) + elif self.camera_info.type == CameraInfo.CameraType.SCREEN_CAPTURE: + self.video_capture = ScreenCapture(self.camera_info.id) + else: + os_name = platform.system() + if ( + os_name == "Windows" + and self.camera_info.type != CameraInfo.CameraType.FILE + and self.camera_info.type != CameraInfo.CameraType.URL + ): + # on windows use the dshow backend by default + self.video_capture = OpenCVVideoCaptureWithSettings( + self.camera_info.id, cv2.CAP_DSHOW + ) + else: + # for files/urls, mac and linux use the default backend + self.video_capture = OpenCVVideoCaptureWithSettings(self.camera_info.id) + + if self.retry_count != self.retry_high_water_mark: + # at the high water mark this is a reconnect + self.retry_count = 0 + + if not self.video_capture.isOpened(): + logger.warn( + "Error: unable to open camera. Check if the camera is connected." + ) + self.update_error.emit("Error: Unable to play video stream") + logger.info("Camera thread stopped") + return False + + # attempt to set the highest resolution + # check if camera index is a OpenCV camera index + # if self.camera_info.type == CameraInfo.CameraType.OPENCV: + # # make sure to open the camera at the highest resolution + # set_camera_highest_resolution(self.video_capture) + + return True + + def run(self): + description_ascii = ( + self.camera_info.description.encode("ascii", errors="ignore").decode() + if type(self.camera_info.description) == str + else str(self.camera_info.description) + ) + logger.info("Starting camera thread for: '%s'", description_ascii) + + if not self.connect_video_capture(): + self.should_stop = True + return + + self.last_update_timestamp = datetime.now() + self.last_frame_timestamp = datetime.now() + self.last_emit_time = datetime.now() + + while not self.should_stop: + if self.video_capture is None: + logger.warn("Error: video capture is None") + break + + if self.retry_count == self.retry_high_water_mark: + logger.warn("Error: retry high water mark exceeded") + # reconnect the video cap + if self.video_capture is not None: + self.video_capture.release() + self.video_capture = None + if not self.connect_video_capture(): + self.should_stop = True + break + self.sleep_fps_target() + self.retry_count += 1 + continue + if self.retry_count > self.retry_count_max: + logger.warn("Error: retry count exceeded") + self.should_stop = True + self.update_error.emit("Error: Unable to play video stream") + break + + # Read the frame from the camera + ret, frame_rgb = None, None + try: + ret, frame_rgb = self.video_capture.read() + except Exception as e: + self.retry_count += 1 + logger.exception( + "Error: unable to read frame from camera (retry count: %d), exception %s", + self.retry_count, + e, + ) + self.sleep_fps_target() + continue + + if not ret: + self.retry_count += 1 + if self.camera_info.type == CameraInfo.CameraType.FILE: + logger.debug("Restarting video file") + with self.video_capture_mutex: + self.video_capture.set(cv2.CAP_PROP_POS_FRAMES, 0) + self.sleep_fps_target() + continue + logger.warn( + "Error: unable to read frame from camera, return value False (retry count: %d)", + self.retry_count, + ) + self.sleep_fps_target() + continue + + self.retry_count = 0 # good frame, reset the retry count + current_time = datetime.now() + + # calculate the frame rate + time_diff_ms = ( + current_time - self.last_frame_timestamp + ).microseconds / 1000 + if time_diff_ms > 0: + self.fps = ( + self.fps_alpha * (1000 / time_diff_ms) + + (1.0 - self.fps_alpha) * self.fps + ) + self.last_frame_timestamp = current_time + + # check that enough time has passed since last update + time_diff_ms = ( + current_time - self.last_update_timestamp + ).microseconds / 1000 + if time_diff_ms < self.update_frame_interval: + # dump this frame since not enough time has passed + self.sleep_fps_target() + continue + # process this frame + self.last_update_timestamp = current_time + self.ups = ( + self.fps_alpha * (1000 / time_diff_ms) + + (1.0 - self.fps_alpha) * self.ups + ) + + # apply rotation if set + if self.crop.rotation != 0: + # use cv2.rotate to rotate the frame + rotateCode = ( + cv2.ROTATE_90_CLOCKWISE + if self.crop.rotation == 90 + else ( + cv2.ROTATE_90_COUNTERCLOCKWISE + if self.crop.rotation == 270 + else cv2.ROTATE_180 + ) + ) + frame_rgb = cv2.rotate(frame_rgb, rotateCode) + + # apply top-level crop if set + if self.crop.isCropSet: + frame_rgb = frame_rgb[ + self.crop.cropTop : frame_rgb.shape[0] - self.crop.cropBottom, + self.crop.cropLeft : frame_rgb.shape[1] - self.crop.cropRight, + ] + + # Stabilize the frame + if self.stabilizationEnabled: + frame_rgb = self.framestabilizer.stabilize_frame(frame_rgb) + + # Apply the homography to the frame + if self.homography is not None: + frame_rgb = cv2.warpPerspective( + frame_rgb, self.homography, (frame_rgb.shape[1], frame_rgb.shape[0]) + ) + + gray = cv2.cvtColor(frame_rgb, cv2.COLOR_BGR2GRAY) + _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU) + + # Detect text in the target + if not self.detectionTargetsStorage.is_empty(): + detectionTargets = self.detectionTargetsStorage.get_data() + texts = self.textDetector.detect_multi_text( + binary, gray, detectionTargets + ) + if len(texts) > 0 and len(detectionTargets) == len(texts): + # augment the text detection targets with the results + results = [] + for i, result in enumerate(texts): + if self.updateOnChange: + if ( + detectionTargets[i].last_text is not None + and detectionTargets[i].last_text == result.text + ): + result.state = ( + TextDetectionTargetWithResult.ResultState.SameNoChange + ) + + results.append( + TextDetectionTargetWithResult( + detectionTargets[i], + result.text, + result.state, + result.rect, + result.extra, + ) + ) + detectionTargets[i].last_text = result.text + + # emit the results + self.ocr_result_signal.emit(results) + + # Emit the signal to update the pixmap once per second + time_diff_prev = (current_time - self.last_emit_time).total_seconds() * 1000 + if time_diff_prev >= self.preview_frame_interval: + if self.show_binary: + self.update_signal.emit(binary) + else: + self.update_signal.emit(frame_rgb) + self.last_emit_time = current_time + self.pps = ( + self.fps_alpha * (1000 / time_diff_prev) + + (1.0 - self.fps_alpha) * self.pps + ) + + self.sleep_fps_target() + + if self.video_capture is not None: + self.video_capture.release() + self.video_capture = None + + logger.info("Camera thread stopped") + + def sleep_fps_target(self): + time_diff_ms = (datetime.now() - self.last_frame_timestamp).microseconds / 1000 + if time_diff_ms < self.frame_interval: + time.sleep((self.frame_interval - time_diff_ms) / 1000.0) + + # on destroy, stop the timer + def __del__(self): + logger.info("Stopping camera") + self.should_stop = True + self.wait() + + def toggleStabilization(self, state): + self.stabilizationEnabled = state + if not state: + self.framestabilizer.reset() diff --git a/camera_view.py b/camera_view.py index 88bf820..4aaa19f 100644 --- a/camera_view.py +++ b/camera_view.py @@ -5,396 +5,15 @@ ) from PySide6.QtCore import Qt from PySide6.QtGui import QImage, QPixmap, QPainter -from PySide6.QtCore import QThread, Signal +from PySide6.QtCore import Signal -import platform -import time import numpy as np import cv2 -import datetime -from datetime import datetime from camera_info import CameraInfo -from ndi import NDICapture -from screen_capture_source import ScreenCapture -from storage import TextDetectionTargetMemoryStorage, subscribe_to_data, fetch_data -from tesseract import TextDetector -from text_detection_target import TextDetectionTargetWithResult +from storage import TextDetectionTargetMemoryStorage, subscribe_to_data from sc_logging import logger -from frame_stabilizer import FrameStabilizer - - -# Function to set the resolution -def set_resolution(cap, width, height): - cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) - cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) - - -# Function to get the resolution -def get_resolution(cap): - width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) - height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) - return width, height - - -# Function to set the camera to the highest resolution -def set_camera_highest_resolution(cap): - # List of common resolutions to try - resolutions = [(1920, 1080), (1280, 720), (1024, 768), (800, 600), (640, 480)] - - # grab one frame to make sure the camera is initialized - ret, _ = cap.read() - if not ret: - logger.warn("Error: camera not initialized") - return - - # grab the current resolution - width, height = get_resolution(cap) - - # If the current resolution is already the highest, return - if width >= resolutions[0][0] and height >= resolutions[0][1]: - logger.info( - "Camera is already at the highest resolution: %d x %d", width, height - ) - return - - # Try each resolution and select the highest one that works - highest_res = resolutions[0] - for resolution in resolutions: - logger.debug("Trying camera resolution: %d x %d", resolution[0], resolution[1]) - set_resolution(cap, *resolution) - test_width, test_height = get_resolution(cap) - logger.debug("Found camera resolution: %d x %d", test_width, test_height) - # if found resolution is within close range of the target resolution, use it - if ( - abs(test_width - resolution[0]) < 100 - and abs(test_height - resolution[1]) < 100 - ): - logger.debug( - "Camera highest resolution set to: %d x %d", test_width, test_height - ) - highest_res = (test_width, test_height) - break - - # Set the highest resolution - logger.info("Setting camera resolution to: %d x %d", highest_res[0], highest_res[1]) - set_resolution(cap, *highest_res) - - -class FrameCropAndRotation: - def __init__(self): - self.isCropSet = fetch_data("scoresight.json", "crop_mode", False) - self.cropTop = fetch_data("scoresight.json", "top_crop", 0) - self.cropBottom = fetch_data("scoresight.json", "bottom_crop", 0) - self.cropLeft = fetch_data("scoresight.json", "left_crop", 0) - self.cropRight = fetch_data("scoresight.json", "right_crop", 0) - self.rotation = fetch_data("scoresight.json", "rotation", 0) - subscribe_to_data("scoresight.json", "crop_mode", self.setCropMode) - subscribe_to_data("scoresight.json", "top_crop", self.setCropTop) - subscribe_to_data("scoresight.json", "bottom_crop", self.setCropBottom) - subscribe_to_data("scoresight.json", "left_crop", self.setCropLeft) - subscribe_to_data("scoresight.json", "right_crop", self.setCropRight) - subscribe_to_data("scoresight.json", "rotation", self.setRotation) - - def setCropMode(self, crop_mode): - self.isCropSet = crop_mode - - def setCropTop(self, crop_top): - self.cropTop = crop_top - - def setCropBottom(self, crop_bottom): - self.cropBottom = crop_bottom - - def setCropLeft(self, crop_left): - self.cropLeft = crop_left - - def setCropRight(self, crop_right): - self.cropRight = crop_right - - def setRotation(self, rotation): - self.rotation = rotation - - -class TimerThread(QThread): - update_signal = Signal(object) - update_error = Signal(object) - ocr_result_signal = Signal(list) - - def __init__( - self, - camera_info: CameraInfo, - detectionTargetsStorage: TextDetectionTargetMemoryStorage, - ): - super().__init__() - self.camera_info = camera_info - self.homography = None - self.detectionTargetsStorage = detectionTargetsStorage - self.textDetector = TextDetector() # initialize tesseract - self.show_binary = False - self.retry_count = 0 - self.retry_count_max = 50 - self.retry_high_water_mark = 25 - self.stabilizationEnabled = False - self.framestabilizer = FrameStabilizer() - self.video_capture = None - self.should_stop = False - self.frame_interval = 30 - self.update_frame_interval = 1000 / fetch_data( - "scoresight.json", "detection_cadence", 5 - ) - subscribe_to_data( - "scoresight.json", "detection_cadence", self.setUpdateFrameInterval - ) - self.preview_frame_interval = 1000 - self.fps = 1000 / self.frame_interval # frames per second - self.pps = 1000 / self.preview_frame_interval # previews per second - self.ups = 1000 / self.update_frame_interval # updates per second - self.fps_alpha = 0.1 # Smoothing factor - self.updateOnChange = True - self.crop = FrameCropAndRotation() - - def setUpdateFrameInterval(self, cadence): - self.update_frame_interval = 1000 / cadence - - def connect_video_capture(self) -> bool: - if self.camera_info.type == CameraInfo.CameraType.NDI: - self.video_capture = NDICapture(self.camera_info.uuid) - elif self.camera_info.type == CameraInfo.CameraType.SCREEN_CAPTURE: - self.video_capture = ScreenCapture(self.camera_info.id) - else: - os_name = platform.system() - if ( - os_name == "Windows" - and self.camera_info.type != CameraInfo.CameraType.FILE - and self.camera_info.type != CameraInfo.CameraType.URL - ): - # on windows use the dshow backend - self.video_capture = cv2.VideoCapture( - self.camera_info.id, cv2.CAP_DSHOW - ) - else: - # for files/urls, mac and linux use the default backend - self.video_capture = cv2.VideoCapture(self.camera_info.id) - - if self.retry_count != self.retry_high_water_mark: - # at the high water mark this is a reconnect - self.retry_count = 0 - - if not self.video_capture.isOpened(): - logger.warn( - "Error: unable to open camera. Check if the camera is connected." - ) - self.update_error.emit("Error: Unable to play video stream") - logger.info("Camera thread stopped") - return False - - # attempt to set the highest resolution - # check if camera index is a OpenCV camera index - # if self.camera_info.type == CameraInfo.CameraType.OPENCV: - # # make sure to open the camera at the highest resolution - # set_camera_highest_resolution(self.video_capture) - - return True - - def run(self): - description_ascii = ( - self.camera_info.description.encode("ascii", errors="ignore").decode() - if type(self.camera_info.description) == str - else str(self.camera_info.description) - ) - logger.info("Starting camera thread for: '%s'", description_ascii) - - if not self.connect_video_capture(): - self.should_stop = True - return - - self.last_update_timestamp = datetime.now() - self.last_frame_timestamp = datetime.now() - self.last_emit_time = datetime.now() - - while not self.should_stop: - if self.video_capture is None: - logger.warn("Error: video capture is None") - break - - if self.retry_count == self.retry_high_water_mark: - logger.warn("Error: retry high water mark exceeded") - # reconnect the video cap - if self.video_capture is not None: - self.video_capture.release() - self.video_capture = None - if not self.connect_video_capture(): - self.should_stop = True - break - self.sleep_fps_target() - self.retry_count += 1 - continue - if self.retry_count > self.retry_count_max: - logger.warn("Error: retry count exceeded") - self.should_stop = True - self.update_error.emit("Error: Unable to play video stream") - break - - # Read the frame from the camera - ret, frame_rgb = None, None - try: - ret, frame_rgb = self.video_capture.read() - except Exception as e: - self.retry_count += 1 - logger.exception( - "Error: unable to read frame from camera (retry count: %d), exception %s", - self.retry_count, - e, - ) - self.sleep_fps_target() - continue - - if not ret: - self.retry_count += 1 - if self.camera_info.type == CameraInfo.CameraType.FILE: - logger.debug("Restarting video file") - self.video_capture.set(cv2.CAP_PROP_POS_FRAMES, 0) - self.sleep_fps_target() - continue - logger.warn( - "Error: unable to read frame from camera, return value False (retry count: %d)", - self.retry_count, - ) - self.sleep_fps_target() - continue - - self.retry_count = 0 # good frame, reset the retry count - current_time = datetime.now() - - # calculate the frame rate - time_diff_ms = ( - current_time - self.last_frame_timestamp - ).microseconds / 1000 - if time_diff_ms > 0: - self.fps = ( - self.fps_alpha * (1000 / time_diff_ms) - + (1.0 - self.fps_alpha) * self.fps - ) - self.last_frame_timestamp = current_time - - # check that enough time has passed since last update - time_diff_ms = ( - current_time - self.last_update_timestamp - ).microseconds / 1000 - if time_diff_ms < self.update_frame_interval: - # dump this frame since not enough time has passed - self.sleep_fps_target() - continue - # process this frame - self.last_update_timestamp = current_time - self.ups = ( - self.fps_alpha * (1000 / time_diff_ms) - + (1.0 - self.fps_alpha) * self.ups - ) - - # apply rotation if set - if self.crop.rotation != 0: - # use cv2.rotate to rotate the frame - rotateCode = ( - cv2.ROTATE_90_CLOCKWISE - if self.crop.rotation == 90 - else ( - cv2.ROTATE_90_COUNTERCLOCKWISE - if self.crop.rotation == 270 - else cv2.ROTATE_180 - ) - ) - frame_rgb = cv2.rotate(frame_rgb, rotateCode) - - # apply top-level crop if set - if self.crop.isCropSet: - frame_rgb = frame_rgb[ - self.crop.cropTop : frame_rgb.shape[0] - self.crop.cropBottom, - self.crop.cropLeft : frame_rgb.shape[1] - self.crop.cropRight, - ] - - # Stabilize the frame - if self.stabilizationEnabled: - frame_rgb = self.framestabilizer.stabilize_frame(frame_rgb) - - # Apply the homography to the frame - if self.homography is not None: - frame_rgb = cv2.warpPerspective( - frame_rgb, self.homography, (frame_rgb.shape[1], frame_rgb.shape[0]) - ) - - gray = cv2.cvtColor(frame_rgb, cv2.COLOR_BGR2GRAY) - _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU) - - # Detect text in the target - if not self.detectionTargetsStorage.is_empty(): - detectionTargets = self.detectionTargetsStorage.get_data() - texts = self.textDetector.detect_multi_text( - binary, gray, detectionTargets - ) - if len(texts) > 0 and len(detectionTargets) == len(texts): - # augment the text detection targets with the results - results = [] - for i, result in enumerate(texts): - if self.updateOnChange: - if ( - detectionTargets[i].last_text is not None - and detectionTargets[i].last_text == result.text - ): - result.state = ( - TextDetectionTargetWithResult.ResultState.SameNoChange - ) - - results.append( - TextDetectionTargetWithResult( - detectionTargets[i], - result.text, - result.state, - result.rect, - result.extra, - ) - ) - detectionTargets[i].last_text = result.text - - # emit the results - self.ocr_result_signal.emit(results) - - # Emit the signal to update the pixmap once per second - time_diff_prev = (current_time - self.last_emit_time).total_seconds() * 1000 - if time_diff_prev >= self.preview_frame_interval: - if self.show_binary: - self.update_signal.emit(binary) - else: - self.update_signal.emit(frame_rgb) - self.last_emit_time = current_time - self.pps = ( - self.fps_alpha * (1000 / time_diff_prev) - + (1.0 - self.fps_alpha) * self.pps - ) - - self.sleep_fps_target() - - if self.video_capture is not None: - self.video_capture.release() - self.video_capture = None - - logger.info("Camera thread stopped") - - def sleep_fps_target(self): - time_diff_ms = (datetime.now() - self.last_frame_timestamp).microseconds / 1000 - if time_diff_ms < self.frame_interval: - time.sleep((self.frame_interval - time_diff_ms) / 1000.0) - - # on destroy, stop the timer - def __del__(self): - logger.info("Stopping camera") - self.should_stop = True - self.wait() - - def toggleStabilization(self, state): - self.stabilizationEnabled = state - if not state: - self.framestabilizer.reset() +from camera_thread import TimerThread class CameraView(QGraphicsView): @@ -421,6 +40,10 @@ def __init__( self.showOSD = True self.camera_width = 0 self.camera_height = 0 + subscribe_to_data("scoresight.json", "video_settings", self.resetFrame) + + def resetFrame(self, data): + self.firstFrameReceived = False def getCameraCapture(self): return self.timerThread.video_capture @@ -503,7 +126,7 @@ def update_pixmap(self, frame): else: self.fps_text.setPlainText(fps_text) # scale the text according to the view size so its always the same size - self.fps_text.setScale(0.002 * self.scenePixmapItem.boundingRect().width()) + self.fps_text.setScale(0.0015 * self.scenePixmapItem.boundingRect().width()) def error_event(self, error): if self.error_text is not None: diff --git a/main.py b/main.py index dfa9261..86924f5 100644 --- a/main.py +++ b/main.py @@ -1,1590 +1,50 @@ -from functools import partial -import os -import platform -import sys -import datetime -from PySide6.QtWidgets import ( - QApplication, - QDialog, - QFileDialog, - QInputDialog, - QLabel, - QMainWindow, - QMenu, - QMessageBox, - QTableWidgetItem, -) -from PySide6.QtGui import QIcon, QStandardItemModel, QStandardItem, QDesktopServices -from PySide6.QtCore import ( - Qt, - Signal, - Slot, - QTranslator, - QLocale, - QObject, - QCoreApplication, - QEvent, - QMetaMethod, - QUrl, -) -from dotenv import load_dotenv -from os import path -from platformdirs import user_data_dir - -from camera_info import CameraInfo -from get_camera_info import get_camera_info -from http_server import start_http_server, update_http_server, stop_http_server -from screen_capture_source import ScreenCapture -from source_view import ImageViewer -from defaults import ( - default_boxes, - default_info_for_box_name, - normalize_settings_dict, - format_prefixes, -) - -from storage import ( - TextDetectionTargetMemoryStorage, - fetch_data, - remove_data, - store_data, - store_custom_box_name, - rename_custom_box_name_in_storage, - remove_custom_box_name_in_storage, - fetch_custom_box_names, -) -from obs_websocket import ( - create_obs_scene_from_export, - open_obs_websocket, - update_text_source, -) - -from template_fields import evaluate_template_field -from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult -from sc_logging import logger -from update_check import check_for_updates -from log_view import LogViewerDialog -import file_output -from video_settings import VideoSettingsDialog -from vmix_output import VMixAPI -from ui_mainwindow import Ui_MainWindow -from ui_about import Ui_Dialog as Ui_About -from ui_connect_obs import Ui_Dialog as Ui_ConnectObs -from ui_url_source import Ui_Dialog as Ui_UrlSource -from ui_screen_capture import Ui_Dialog as Ui_ScreenCapture - - -def clear_layout(layout): - while layout.count(): - item = layout.takeAt(0) - widget = item.widget() - if widget is not None: - widget.deleteLater() - widget = None - else: - clear_layout(item.layout()) - - -class MainWindow(QMainWindow): - # add a signal to update sources - update_sources = Signal(list) - get_sources = Signal() - - def __init__(self, translator: QTranslator, parent: QObject): - super(MainWindow, self).__init__() - self.parent_object = parent - self.ui = Ui_MainWindow() - logger.info("Starting ScoreSight") - self.ui.setupUi(self) - self.translator = translator - # load env variables - load_dotenv() - self.setWindowTitle(f"ScoreSight - v{os.getenv('LOCAL_RELEASE_TAG')}") - if platform.system() == "Windows": - # set the icon - self.setWindowIcon( - QIcon( - path.abspath( - path.join(path.dirname(__file__), "icons/Windows-icon-open.ico") - ) - ) - ) - - self.menubar = self.menuBar() - file_menu = self.menubar.addMenu("File") - - # check for updates - check_for_updates(False) - file_menu.addAction("Check for Updates", lambda: check_for_updates(True)) - file_menu.addAction("About", self.openAboutDialog) - file_menu.addAction("View Current Log", self.openLogsDialog) - file_menu.addAction("Import Configuration", self.importConfiguration) - file_menu.addAction("Export Configuration", self.exportConfiguration) - file_menu.addAction("Open Configuration Folder", self.openConfigurationFolder) - - # Add "Language" menu - languageMenu = file_menu.addMenu("Language") - - # Add language options - self.addLanguageOption(languageMenu, "English (US)", "en_US") - self.addLanguageOption(languageMenu, "French (France)", "fr_FR") - self.addLanguageOption(languageMenu, "Spanish (Spain)", "es_ES") - self.addLanguageOption(languageMenu, "German", "de_DE") - self.addLanguageOption(languageMenu, "Italian", "it_IT") - self.addLanguageOption(languageMenu, "Japanese", "ja_JP") - self.addLanguageOption(languageMenu, "Korean", "ko_KR") - self.addLanguageOption(languageMenu, "Dutch", "nl_NL") - self.addLanguageOption(languageMenu, "Polish", "pl_PL") - self.addLanguageOption(languageMenu, "Portuguese (Brazil)", "pt_BR") - self.addLanguageOption(languageMenu, "Portuguese (Portugal)", "pt_PT") - self.addLanguageOption(languageMenu, "Russian", "ru_RU") - self.addLanguageOption(languageMenu, "Chinese (Simplified)", "zh_CN") - - # Hide the menu bar by default - self.menubar.setVisible(False) - - # Show the menu bar when the Alt key is pressed - self.installEventFilter(self) - - self.ui.pushButton_connectObs.clicked.connect(self.openOBSConnectModal) - - self.vmixUiSetup() - - start_http_server() - - self.ui.pushButton_stabilize.setEnabled(True) - self.ui.pushButton_stabilize.clicked.connect(self.toggleStabilize) - - self.ui.toolButton_topCrop.clicked.connect(self.cropMode) - # check configuation if crop is enabled - self.ui.toolButton_topCrop.setChecked( - fetch_data("scoresight.json", "crop_mode", False) - ) - self.ui.widget_cropPanel.setVisible(self.ui.toolButton_topCrop.isChecked()) - self.ui.widget_cropPanel.setEnabled(self.ui.toolButton_topCrop.isChecked()) - self.ui.spinBox_leftCrop.valueChanged.connect( - partial(self.globalSettingsChanged, "left_crop") - ) - self.ui.spinBox_leftCrop.setValue(fetch_data("scoresight.json", "left_crop", 0)) - self.ui.spinBox_rightCrop.valueChanged.connect( - partial(self.globalSettingsChanged, "right_crop") - ) - self.ui.spinBox_rightCrop.setValue( - fetch_data("scoresight.json", "right_crop", 0) - ) - self.ui.spinBox_topCrop.valueChanged.connect( - partial(self.globalSettingsChanged, "top_crop") - ) - self.ui.spinBox_topCrop.setValue(fetch_data("scoresight.json", "top_crop", 0)) - self.ui.spinBox_bottomCrop.valueChanged.connect( - partial(self.globalSettingsChanged, "bottom_crop") - ) - self.ui.spinBox_bottomCrop.setValue( - fetch_data("scoresight.json", "bottom_crop", 0) - ) - self.ui.checkBox_enableOutAPI.toggled.connect( - partial(self.globalSettingsChanged, "enable_out_api") - ) - self.ui.checkBox_enableOutAPI.setChecked( - fetch_data("scoresight.json", "enable_out_api", False) - ) - - # connect toolButton_rotate - self.ui.toolButton_rotate.clicked.connect(self.rotateImage) - - self.ui.widget_detectionCadence.setVisible(True) - self.ui.horizontalSlider_detectionCadence.setValue( - fetch_data("scoresight.json", "detection_cadence", 5) - ) - self.ui.horizontalSlider_detectionCadence.valueChanged.connect( - self.detectionCadenceChanged - ) - self.ui.toolButton_addBox.clicked.connect(self.addBox) - self.ui.toolButton_removeBox.clicked.connect(self.removeCustomBox) - - self.video_settings_dialog = None - self.ui.toolButton_videoSettings.clicked.connect(self.openVideoSettings) - - self.obs_websocket_client = None - - ocr_models = [ - "Daktronics", - "General Scoreboard", - "General Fonts (English)", - "General Scoreboard Large", - ] - self.ui.comboBox_ocrModel.addItems(ocr_models) - self.ui.comboBox_ocrModel.setCurrentIndex( - fetch_data("scoresight.json", "ocr_model", 1) - ) # default to General Scoreboard - - self.ui.frame_source_view.setEnabled(False) - self.ui.groupBox_target_settings.setEnabled(False) - self.ui.pushButton_makeBox.clicked.connect(self.makeBox) - self.ui.pushButton_removeBox.clicked.connect(self.removeBox) - self.ui.tableWidget_boxes.itemClicked.connect(self.listItemClicked) - # connect the edit triggers - self.ui.tableWidget_boxes.itemDoubleClicked.connect(self.editBoxName) - self.ui.pushButton_refresh_sources.clicked.connect( - lambda: self.get_sources.emit() - ) - self.detectionTargetsStorage = TextDetectionTargetMemoryStorage() - self.detectionTargetsStorage.data_changed.connect(self.detectionTargetsChanged) - self.ui.pushButton_createOBSScene.clicked.connect(self.createOBSScene) - self.ui.pushButton_selectFolder.clicked.connect(self.selectOutputFolder) - self.ui.toolButton_trashFolder.clicked.connect(self.clearOutputFolder) - self.ui.pushButton_stopUpdates.toggled.connect(self.toggleStopUpdates) - self.ui.comboBox_ocrModel.currentIndexChanged.connect(self.ocrModelChanged) - self.ui.pushButton_restoreDefaults.clicked.connect(self.restoreDefaults) - self.ui.toolButton_zoomReset.clicked.connect(self.resetZoom) - self.ui.toolButton_osd.toggled.connect(self.toggleOSD) - self.ui.toolButton_showOCRrects.toggled.connect(self.toggleOCRRects) - self.ui.checkBox_smoothing.toggled.connect( - partial(self.genericSettingsChanged, "smoothing") - ) - self.ui.checkBox_skip_empty.toggled.connect( - partial(self.genericSettingsChanged, "skip_empty") - ) - self.ui.horizontalSlider_conf_thresh.valueChanged.connect( - self.confThreshChanged - ) - self.ui.lineEdit_format.textChanged.connect( - partial(self.genericSettingsChanged, "format_regex") - ) - self.ui.comboBox_fieldType.currentIndexChanged.connect( - partial(self.genericSettingsChanged, "type") - ) - self.ui.checkBox_skip_similar_image.toggled.connect( - partial(self.genericSettingsChanged, "skip_similar_image") - ) - self.ui.checkBox_autocrop.toggled.connect( - partial(self.genericSettingsChanged, "autocrop") - ) - self.ui.horizontalSlider_cleanup.valueChanged.connect(self.cleanupThreshChanged) - self.ui.horizontalSlider_dilate.valueChanged.connect( - partial(self.genericSettingsChanged, "dilate") - ) - self.ui.horizontalSlider_skew.valueChanged.connect( - partial(self.genericSettingsChanged, "skew") - ) - self.ui.horizontalSlider_vscale.valueChanged.connect( - partial(self.genericSettingsChanged, "vscale") - ) - self.ui.checkBox_removeLeadingZeros.toggled.connect( - partial(self.genericSettingsChanged, "remove_leading_zeros") - ) - self.ui.checkBox_rescalePatch.toggled.connect( - partial(self.genericSettingsChanged, "rescale_patch") - ) - self.ui.checkBox_normWHRatio.toggled.connect( - partial(self.genericSettingsChanged, "normalize_wh_ratio") - ) - self.ui.checkBox_invertPatch.toggled.connect( - partial(self.genericSettingsChanged, "invert_patch") - ) - self.ui.checkBox_dotDetector.toggled.connect( - partial(self.genericSettingsChanged, "dot_detector") - ) - self.ui.checkBox_ordinalIndicator.toggled.connect( - partial(self.genericSettingsChanged, "ordinal_indicator") - ) - self.ui.comboBox_binarizationMethod.currentIndexChanged.connect( - partial(self.genericSettingsChanged, "binarization_method") - ) - self.ui.checkBox_templatefield.toggled.connect(self.makeTemplateField) - self.ui.lineEdit_templatefield.textChanged.connect( - partial(self.genericSettingsChanged, "templatefield_text") - ) - self.ui.comboBox_formatPrefix.currentIndexChanged.connect( - self.formatPrefixChanged - ) - self.ui.checkBox_updateOnchange.toggled.connect(self.toggleUpdateOnChange) - - # populate the tableWidget_boxes with the default and custom boxes - custom_boxes_names = fetch_custom_box_names() - - for box_name in [box["name"] for box in default_boxes] + custom_boxes_names: - item = QTableWidgetItem( - QIcon( - path.abspath( - path.join(path.dirname(__file__), "icons/circle-x.svg") - ) - ), - box_name, - ) - item.setData(Qt.ItemDataRole.UserRole, "unchecked") - self.ui.tableWidget_boxes.insertRow(self.ui.tableWidget_boxes.rowCount()) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 0, item - ) - disabledItem = QTableWidgetItem() - disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem - ) - - self.image_viewer = None - self.obs_connect_modal = None - self.source_name = None - self.updateOCRResults = True - self.log_dialog = None - if fetch_data("scoresight.json", "open_on_startup", False): - logger.info("Opening log dialog on startup") - self.openLogsDialog() - - if fetch_data("scoresight.json", "obs"): - self.connectObs() - - self.out_folder = fetch_data("scoresight.json", "output_folder") - if self.out_folder: - if not path.exists(self.out_folder): - self.out_folder = None - remove_data("scoresight.json", "output_folder") - else: - self.ui.lineEdit_folder.setText(self.out_folder) - - self.first_csv_append = True - self.last_aggregate_save = datetime.datetime.now() - self.ui.checkBox_saveCsv.toggled.connect( - partial(self.globalSettingsChanged, "save_csv") - ) - self.ui.checkBox_saveCsv.setChecked( - fetch_data("scoresight.json", "save_csv", False) - ) - self.ui.checkBox_saveXML.toggled.connect( - partial(self.globalSettingsChanged, "save_xml") - ) - self.ui.checkBox_saveXML.setChecked( - fetch_data("scoresight.json", "save_xml", False) - ) - self.ui.comboBox_appendMethod.currentIndexChanged.connect( - partial(self.globalSettingsChanged, "append_method") - ) - self.ui.horizontalSlider_aggsPerSecond.valueChanged.connect( - partial(self.globalSettingsChanged, "aggs_per_second") - ) - self.ui.comboBox_appendMethod.setCurrentIndex( - fetch_data("scoresight.json", "append_method", 3) - ) - self.ui.horizontalSlider_aggsPerSecond.setValue( - fetch_data("scoresight.json", "aggs_per_second", 5) - ) - self.ui.checkBox_updateOnchange.setChecked( - fetch_data("scoresight.json", "update_on_change", True) - ) - self.ui.checkBox_vmix_send_same.setChecked( - fetch_data("scoresight.json", "vmix_send_same", False) - ) - self.ui.checkBox_vmix_send_same.toggled.connect( - partial(self.globalSettingsChanged, "vmix_send_same") - ) - - self.update_sources.connect(self.updateSources) - self.get_sources.connect(self.getSources) - self.get_sources.emit() - - def openVideoSettings(self): - # only allow opening the video settings for an OpenCV type source - if not self.image_viewer or not self.image_viewer.getCameraCapture(): - return - - if self.image_viewer.getCameraInfo().type != CameraInfo.CameraType.OPENCV: - return - - if self.video_settings_dialog is None: - # open the logs dialog - self.video_settings_dialog = VideoSettingsDialog() - self.video_settings_dialog.setWindowTitle("Video Settings") - - self.video_settings_dialog.init_ui(self.image_viewer.getCameraCapture()) - - # show the dialog, non modal - self.video_settings_dialog.show() - - def rotateImage(self): - # store the rotation in the scoresight.json - rotation = fetch_data("scoresight.json", "rotation", 0) - rotation += 90 - if rotation >= 360: - rotation = 0 - self.globalSettingsChanged("rotation", rotation) - - def cropMode(self): - # if the toolButton_topCrop is unchecked, go to crop mode - if self.ui.toolButton_topCrop.isChecked(): - self.ui.widget_cropPanel.setVisible(True) - self.ui.widget_cropPanel.setEnabled(True) - self.globalSettingsChanged("crop_mode", True) - else: - self.ui.widget_cropPanel.setVisible(False) - self.ui.widget_cropPanel.setEnabled(False) - self.globalSettingsChanged("crop_mode", False) - - def globalSettingsChanged(self, settingName, value): - store_data("scoresight.json", settingName, value) - - def eventFilter(self, obj, event): - if event.type() == QEvent.KeyPress and event.key() == Qt.Key_Alt: - self.menubar.setVisible(True) - elif event.type() == QEvent.FocusOut and self.menubar.isVisible(): - self.menubar.setVisible(False) - elif event.type() == QEvent.WindowDeactivate and self.menubar.isVisible(): - self.menubar.setVisible(False) - return super().eventFilter(obj, event) - - def focusOutEvent(self, event): - if self.menubar.isVisible(): - self.menubar.setVisible(False) - super().focusOutEvent(event) - - def changeEvent(self, event): - if event.type() == QEvent.WindowDeactivate and self.menubar.isVisible(): - self.menubar.setVisible(False) - super().changeEvent(event) - - def changeLanguage(self, locale): - locale_file = path.abspath( - path.join(path.dirname(__file__), "translations", f"scoresight_{locale}.qm") - ) - logger.info(f"Changing language to {locale_file}") - if not self.translator.load(locale_file): - logger.error(f"Could not load translation for {locale_file}") - return - appInstance = QApplication.instance() - if appInstance: - logger.info(f"installing translator for {locale}") - appInstance.installTranslator(self.translator) - self.ui.retranslateUi(self) - - def addLanguageOption(self, menu: QMenu, language_name: str, locale: str): - menu.addAction(language_name, lambda: self.changeLanguage(locale)) - - def toggleUpdateOnChange(self, value): - self.globalSettingsChanged("update_on_change", value) - if self.image_viewer: - self.image_viewer.setUpdateOnChange(value) - - def formatPrefixChanged(self, index): - if index == 12: - return # do nothing if "Select Preset" is selected - # based on the selected index, set the format prefix - # change lineEdit_format to the selected format prefix - self.ui.lineEdit_format.setText(format_prefixes[index]) - - def importConfiguration(self): - # open a file dialog to select a configuration file - file, _ = QFileDialog.getOpenFileName( - self, "Open Configuration File", "", "Configuration Files (*.json)" - ) - if not file: - return - # load the configuration from the file - if not self.detectionTargetsStorage.loadBoxesFromFile(file): - # show an error qmessagebox - logger.error("Error loading configuration file") - QMessageBox.critical( - self, - "Error", - "Error loading configuration file", - QMessageBox.StandardButton.Ok, - ) - return - - def exportConfiguration(self): - # open a file dialog to select the output file - file, _ = QFileDialog.getSaveFileName( - self, "Save Configuration File", "", "Configuration Files (*.json)" - ) - if not file: - return - # save the configuration to the file - self.detectionTargetsStorage.saveBoxesToFile(file) - - def openConfigurationFolder(self): - # open the configuration folder in the file explorer - QDesktopServices.openUrl( - QUrl( - "file:///" + user_data_dir("scoresight"), QUrl.ParsingMode.TolerantMode - ) - ) - - def toggleOSD(self, value): - if self.image_viewer: - self.image_viewer.toggleOSD(value) - - def toggleOCRRects(self, value): - if self.image_viewer: - self.image_viewer.toggleOCRRects(value) - - def resetZoom(self): - if self.image_viewer: - self.image_viewer.resetZoom() - - def detectionCadenceChanged(self, detections_per_second): - self.globalSettingsChanged("detection_cadence", detections_per_second) - if self.image_viewer and self.image_viewer.timerThread: - # convert the detections_per_second to milliseconds - self.image_viewer.timerThread.update_frame_interval = ( - 1000 / detections_per_second - ) - - def ocrModelChanged(self, index): - self.globalSettingsChanged("ocr_model", index) - # update the ocr model in the text detector - if ( - self.image_viewer - and self.image_viewer.timerThread - and self.image_viewer.timerThread.textDetector - ): - self.image_viewer.timerThread.textDetector.setOcrModel(index) - - def openLogsDialog(self): - if self.log_dialog is None: - # open the logs dialog - self.log_dialog = LogViewerDialog() - self.log_dialog.setWindowTitle("Logs") - - # show the dialog, non modal - self.log_dialog.show() - - def openAboutDialog(self): - # open the about dialog - about_dialog = QDialog() - about_dialog_ui = Ui_About() - about_dialog_ui.setupUi(about_dialog) - about_dialog.setWindowTitle("About ScoreSight") - about_dialog.exec() - - def toggleStabilize(self): - if not self.image_viewer: - return - # start or stop the stabilization - self.image_viewer.toggleStabilization(self.ui.pushButton_stabilize.isChecked()) - - def toggleStopUpdates(self, value): - self.updateOCRResults = not value - # change the text on the button - self.ui.pushButton_stopUpdates.setText( - self.translator.translate("main", "Resume updates") - if value - else self.translator.translate("main", "Stop updates") - ) - - def selectOutputFolder(self): - # open a Qt dialog to select the output folder - folder = QFileDialog.getExistingDirectory( - self, - "Select Output Folder", - fetch_data("scoresight.json", "output_folder", ""), - options=QFileDialog.Option.ShowDirsOnly, - ) - if folder and len(folder) > 0: - self.ui.lineEdit_folder.setText(folder) - self.out_folder = folder - self.globalSettingsChanged("output_folder", folder) - - def clearOutputFolder(self): - # clear the output folder - self.ui.lineEdit_folder.setText("") - self.out_folder = None - remove_data("scoresight.json", "output_folder") - - def editSettings(self, settingsMutatorCallback): - # update the selected item's settings in the detectionTargetsStorage - item = self.ui.tableWidget_boxes.currentItem() - if item is None: - logger.info("no item selected") - return - item_name = item.text() - item_obj = self.detectionTargetsStorage.find_item_by_name(item_name) - if item_obj is None: - logger.info("item not found: %s", item_name) - return - item_obj = settingsMutatorCallback(item_obj) - self.detectionTargetsStorage.edit_item(item_name, item_obj) - - def restoreDefaults(self): - # restore the default settings for the selected item - def restoreDefaultsSettings(item_obj): - info = default_info_for_box_name(item_obj.name) - item_obj.settings = normalize_settings_dict({}, info) - return item_obj - - self.editSettings(restoreDefaultsSettings) - self.populateSettings(self.ui.tableWidget_boxes.currentItem().text()) - - def confThreshChanged(self): - def editConfThreshSettings(item_obj): - item_obj.settings["conf_thresh"] = ( - float(self.ui.horizontalSlider_conf_thresh.value()) / 100.0 - ) - return item_obj - - self.editSettings(editConfThreshSettings) - - def cleanupThreshChanged(self): - def editCleanupThreshSettings(item_obj): - item_obj.settings["cleanup_thresh"] = ( - float(self.ui.horizontalSlider_cleanup.value()) / 100.0 - ) - return item_obj - - self.editSettings(editCleanupThreshSettings) - - def genericSettingsChanged(self, settingName, value): - def editGenericSettings(item_obj): - item_obj.settings[settingName] = value - return item_obj - - self.editSettings(editGenericSettings) - - def vmixConnectionChanged(self): - self.vmixUpdater = VMixAPI( - self.ui.lineEdit_vmixHost.text(), - self.ui.lineEdit_vmixPort.text(), - self.ui.inputLineEdit_vmix.text(), - {}, - ) - self.globalSettingsChanged("vmix_host", self.ui.lineEdit_vmixHost.text()) - self.globalSettingsChanged("vmix_port", self.ui.lineEdit_vmixPort.text()) - self.globalSettingsChanged("vmix_input", self.ui.inputLineEdit_vmix.text()) - - def vmixMappingChanged(self, _): - # store entire mapping data in scoresight.json - mapping = {} - model = self.ui.tableView_vmixMapping.model() - if isinstance(model, QStandardItemModel): - for i in range(model.rowCount()): - item = model.item(i, 0) - value = model.item(i, 1) - if item and value: - mapping[item.text()] = value.text() - self.globalSettingsChanged("vmix_mapping", mapping) - self.vmixUpdater.set_field_mapping(mapping) - else: - logger.error("vmixMappingChanged: model is not a QStandardItemModel") - - def vmixUiSetup(self): - # populate the vmix connection from storage - self.ui.lineEdit_vmixHost.setText( - fetch_data("scoresight.json", "vmix_host", "localhost") - ) - self.ui.lineEdit_vmixPort.setText( - fetch_data("scoresight.json", "vmix_port", "8099") - ) - self.ui.inputLineEdit_vmix.setText( - fetch_data("scoresight.json", "vmix_input", "1") - ) - # connect the lineEdits to vmixConnectionChanged - self.ui.lineEdit_vmixHost.textChanged.connect(self.vmixConnectionChanged) - self.ui.lineEdit_vmixPort.textChanged.connect(self.vmixConnectionChanged) - self.ui.inputLineEdit_vmix.textChanged.connect(self.vmixConnectionChanged) - - # create the vmixUpdater - self.vmixUpdater = VMixAPI( - self.ui.lineEdit_vmixHost.text(), - self.ui.lineEdit_vmixPort.text(), - self.ui.inputLineEdit_vmix.text(), - {}, - ) - # add standard item model to the tableView_vmixMapping - self.ui.tableView_vmixMapping.setModel(QStandardItemModel()) - mapping = fetch_data("scoresight.json", "vmix_mapping", {}) - if mapping: - self.vmixUpdater.set_field_mapping(mapping) - - self.ui.tableView_vmixMapping.model().dataChanged.connect( - self.vmixMappingChanged - ) - - self.ui.pushButton_startvmix.toggled.connect(self.togglevMix) - - def togglevMix(self, value): - if not self.vmixUpdater: - return - if value: - self.ui.pushButton_startvmix.setText("🛑 Stop vMix") - self.vmixUpdater.running = True - else: - self.ui.pushButton_startvmix.setText("▶️ Start vMix") - self.vmixUpdater.running = False - - def updatevMixTable(self, detectionTargets): - mapping_storage = fetch_data("scoresight.json", "vmix_mapping") - model = QStandardItemModel() - model.blockSignals(True) - - for box in detectionTargets: - # add the detection to the vmix output mapping: tableView_vmixMapping - # check if the table already has the detectionTarget - items = model.findItems(box.name, Qt.MatchFlag.MatchExactly) - if len(items) == 0: - # add the item to the list - row = model.rowCount() - model.insertRow(row) - model.setItem(row, 0, QStandardItem(box.name)) - # the first item shouldn't be editable - model.item(row, 0).setFlags(Qt.ItemFlag.NoItemFlags) - else: - # update the item in the list - item = items[0] - row = item.row() - - # get value from storage or use the box name - if mapping_storage and box.name in mapping_storage: - model.setItem(row, 1, QStandardItem(mapping_storage[box.name])) - else: - model.setItem(row, 1, QStandardItem(box.name)) - # remove the items that are not in the detectionTargets - for i in range(model.rowCount()): - item = model.item(i, 0) - if not any([box.name == item.text() for box in detectionTargets]): - model.removeRow(i) - - model.blockSignals(False) - self.ui.tableView_vmixMapping.setModel(model) - self.ui.tableView_vmixMapping.model().dataChanged.connect( - self.vmixMappingChanged - ) - - def detectionTargetsChanged(self, detectionTargets): - for box in detectionTargets: - logger.debug(f"Change: Detection target: {box.name}") - # change the list icon to green checkmark - items = self.ui.tableWidget_boxes.findItems( - box.name, Qt.MatchFlag.MatchExactly - ) - if len(items) == 0: - logger.warning(f"Item not found: {box.name}. Adding it to the list.") - # add the item to the list - self.ui.tableWidget_boxes.insertRow( - self.ui.tableWidget_boxes.rowCount() - ) - item = QTableWidgetItem(box.name) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 0, item - ) - disabledItem = QTableWidgetItem() - disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem - ) - else: - item = items[0] - - if not box.settings["templatefield"]: - # this is a detection target - item.setIcon( - QIcon( - path.abspath( - path.join(path.dirname(__file__), "icons/circle-check.svg") - ) - ) - ) - item.setData(Qt.ItemDataRole.UserRole, "checked") - else: - # this is a template field - item.setIcon( - QIcon( - path.abspath( - path.join( - path.dirname(__file__), "icons/template-field.svg" - ) - ) - ) - ) - item.setData(Qt.ItemDataRole.UserRole, "templatefield") - - self.updatevMixTable(detectionTargets) - - # if save_csv is enabled, truncate the aggregate file - if self.ui.checkBox_saveCsv.isChecked() and self.out_folder: - csv_output_file_path = path.abspath( - path.join(self.out_folder, "results.csv") - ) - try: - with open(csv_output_file_path, "w") as f: - f.write("") - self.first_csv_append = True - self.last_aggregate_save = datetime.datetime.now() - except Exception as e: - logger.error(f"Error truncating aggregate file: {e}") - - def populateSettings(self, name): - self.ui.lineEdit_format.blockSignals(True) - self.ui.comboBox_fieldType.blockSignals(True) - self.ui.checkBox_smoothing.blockSignals(True) - self.ui.checkBox_skip_empty.blockSignals(True) - self.ui.horizontalSlider_conf_thresh.blockSignals(True) - self.ui.checkBox_autocrop.blockSignals(True) - self.ui.checkBox_skip_similar_image.blockSignals(True) - self.ui.horizontalSlider_cleanup.blockSignals(True) - self.ui.horizontalSlider_dilate.blockSignals(True) - self.ui.horizontalSlider_skew.blockSignals(True) - self.ui.horizontalSlider_vscale.blockSignals(True) - self.ui.checkBox_removeLeadingZeros.blockSignals(True) - self.ui.checkBox_rescalePatch.blockSignals(True) - self.ui.checkBox_normWHRatio.blockSignals(True) - self.ui.checkBox_invertPatch.blockSignals(True) - self.ui.checkBox_ordinalIndicator.blockSignals(True) - self.ui.comboBox_binarizationMethod.blockSignals(True) - self.ui.comboBox_formatPrefix.blockSignals(True) - self.ui.checkBox_templatefield.blockSignals(True) - self.ui.lineEdit_templatefield.blockSignals(True) - - # populate the settings from the detectionTargetsStorage - item_obj = self.detectionTargetsStorage.find_item_by_name(name) - if item_obj is None: - self.ui.lineEdit_format.setText("") - self.ui.comboBox_fieldType.setCurrentIndex(0) - self.ui.checkBox_smoothing.setChecked(True) - self.ui.checkBox_skip_empty.setChecked(True) - self.ui.horizontalSlider_conf_thresh.setValue(50) - self.ui.checkBox_autocrop.setChecked(False) - self.ui.checkBox_skip_similar_image.setChecked(False) - self.ui.horizontalSlider_cleanup.setValue(0) - self.ui.horizontalSlider_dilate.setValue(1) - self.ui.horizontalSlider_skew.setValue(0) - self.ui.horizontalSlider_vscale.setValue(10) - self.ui.label_selectedInfo.setText("") - self.ui.checkBox_removeLeadingZeros.setChecked(False) - self.ui.checkBox_rescalePatch.setChecked(False) - self.ui.checkBox_normWHRatio.setChecked(False) - self.ui.checkBox_invertPatch.setChecked(False) - self.ui.checkBox_ordinalIndicator.setChecked(False) - self.ui.comboBox_binarizationMethod.setCurrentIndex(0) - self.ui.checkBox_templatefield.setChecked(False) - self.ui.lineEdit_templatefield.setText("") - else: - item_obj.settings = normalize_settings_dict( - item_obj.settings, default_info_for_box_name(item_obj.name) - ) - self.ui.label_selectedInfo.setText(f"{item_obj.name}") - self.ui.lineEdit_format.setText(item_obj.settings["format_regex"]) - self.ui.comboBox_fieldType.setCurrentIndex(item_obj.settings["type"]) - self.ui.checkBox_smoothing.setChecked(item_obj.settings["smoothing"]) - self.ui.checkBox_skip_empty.setChecked(item_obj.settings["skip_empty"]) - self.ui.horizontalSlider_conf_thresh.setValue( - int(item_obj.settings["conf_thresh"] * 100) - ) - self.ui.checkBox_autocrop.setChecked(item_obj.settings["autocrop"]) - self.ui.checkBox_skip_similar_image.setChecked( - item_obj.settings["skip_similar_image"] - ) - self.ui.horizontalSlider_cleanup.setValue( - int(item_obj.settings["cleanup_thresh"] * 100) - ) - self.ui.horizontalSlider_dilate.setValue(item_obj.settings["dilate"]) - self.ui.horizontalSlider_skew.setValue(item_obj.settings["skew"]) - self.ui.horizontalSlider_vscale.setValue(item_obj.settings["vscale"]) - self.ui.checkBox_removeLeadingZeros.setChecked( - item_obj.settings["remove_leading_zeros"] - ) - self.ui.checkBox_rescalePatch.setChecked(item_obj.settings["rescale_patch"]) - self.ui.checkBox_normWHRatio.setChecked( - item_obj.settings["normalize_wh_ratio"] - ) - self.ui.checkBox_invertPatch.setChecked(item_obj.settings["invert_patch"]) - self.ui.checkBox_dotDetector.setChecked(item_obj.settings["dot_detector"]) - self.ui.checkBox_ordinalIndicator.setChecked( - item_obj.settings["ordinal_indicator"] - ) - self.ui.comboBox_binarizationMethod.setCurrentIndex( - item_obj.settings["binarization_method"] - ) - self.ui.checkBox_templatefield.setChecked( - item_obj.settings["templatefield"] - ) - self.ui.lineEdit_templatefield.setText( - item_obj.settings["templatefield_text"] - ) - - self.ui.comboBox_formatPrefix.setCurrentIndex(12) - - self.ui.lineEdit_format.blockSignals(False) - self.ui.comboBox_fieldType.blockSignals(False) - self.ui.checkBox_smoothing.blockSignals(False) - self.ui.checkBox_skip_empty.blockSignals(False) - self.ui.horizontalSlider_conf_thresh.blockSignals(False) - self.ui.checkBox_autocrop.blockSignals(False) - self.ui.checkBox_skip_similar_image.blockSignals(False) - self.ui.horizontalSlider_cleanup.blockSignals(False) - self.ui.horizontalSlider_dilate.blockSignals(False) - self.ui.horizontalSlider_skew.blockSignals(False) - self.ui.horizontalSlider_vscale.blockSignals(False) - self.ui.checkBox_removeLeadingZeros.blockSignals(False) - self.ui.checkBox_rescalePatch.blockSignals(False) - self.ui.checkBox_normWHRatio.blockSignals(False) - self.ui.checkBox_invertPatch.blockSignals(False) - self.ui.checkBox_ordinalIndicator.blockSignals(False) - self.ui.comboBox_binarizationMethod.blockSignals(False) - self.ui.comboBox_formatPrefix.blockSignals(False) - self.ui.checkBox_templatefield.blockSignals(False) - self.ui.lineEdit_templatefield.blockSignals(False) - - def listItemClicked(self, item): - user_role = item.data(Qt.ItemDataRole.UserRole) - if user_role in ["checked", "templatefield"] and item.column() == 0: - # enable the remove box button and disable the make box button - self.ui.pushButton_makeBox.setEnabled(False) - self.ui.pushButton_removeBox.setEnabled(user_role == "checked") - self.ui.groupBox_target_settings.setEnabled(user_role == "checked") - self.populateSettings(item.text()) - else: - # enable the make box button and disable the remove box button - self.ui.pushButton_removeBox.setEnabled(False) - self.ui.pushButton_makeBox.setEnabled(item.column() == 0) - self.ui.groupBox_target_settings.setEnabled(False) - self.populateSettings("") - - if item.column() == 0: - # if this is not a default box - enable the template field checkbox - if item.text() not in [box["name"] for box in default_boxes]: - self.ui.checkBox_templatefield.setEnabled(True) - self.ui.lineEdit_templatefield.setEnabled(True) - else: - self.ui.checkBox_templatefield.setEnabled(False) - self.ui.lineEdit_templatefield.setEnabled(False) - - def openOBSConnectModal(self): - # disable OBS options - self.ui.lineEdit_sceneName.setEnabled(False) - self.ui.checkBox_recreate.setEnabled(False) - self.ui.pushButton_createOBSScene.setEnabled(False) - - # load the ui from "connect_obs.ui" - self.obs_modal_ui = Ui_ConnectObs() - self.obs_connect_modal = QDialog() - self.obs_modal_ui.setupUi(self.obs_connect_modal) - self.obs_connect_modal.setWindowTitle("Connect to OBS") - - # connect the "connect" button to a function - self.obs_modal_ui.pushButton_connect.clicked.connect(self.connectObs) - - # load the saved data from scoresight.json - obs_data = fetch_data("scoresight.json", "obs") - if obs_data: - self.obs_modal_ui.lineEdit_ip.setText(obs_data["ip"]) - self.obs_modal_ui.lineEdit_port.setText(obs_data["port"]) - self.obs_modal_ui.lineEdit_password.setText(obs_data["password"]) - # show the modal - self.obs_connect_modal.show() - # focus the connect button - self.obs_modal_ui.pushButton_connect.setFocus() - - def connectObs(self): - # open a websocket connection to OBS using obs_websocket.py - # enable the save button in the modal if the connection is successful - if self.obs_connect_modal is not None: - self.obs_websocket_client = open_obs_websocket( - { - "ip": self.obs_modal_ui.lineEdit_ip.text(), - "port": self.obs_modal_ui.lineEdit_port.text(), - "password": self.obs_modal_ui.lineEdit_password.text(), - } - ) - else: - self.obs_websocket_client = open_obs_websocket( - fetch_data("scoresight.json", "obs") - ) - if not self.obs_websocket_client: - # show error in label_error - if self.obs_connect_modal: - self.obs_modal_ui.label_error.setText("Cannot connect to OBS") - return - - # connection was successful - if self.obs_connect_modal: - store_data( - "scoresight.json", - "obs", - { - "ip": self.obs_modal_ui.lineEdit_ip.text(), - "port": self.obs_modal_ui.lineEdit_port.text(), - "password": self.obs_modal_ui.lineEdit_password.text(), - }, - ) - self.obs_connect_modal.close() - - self.ui.lineEdit_sceneName.setEnabled(True) - self.ui.checkBox_recreate.setEnabled(True) - self.ui.pushButton_createOBSScene.setEnabled(True) - - logger.info("OBS: Connected") - - @Slot() - def getSources(self): - self.update_sources.emit(get_camera_info()) - - @Slot(list) - def updateSources(self, camera_sources: list[CameraInfo]): - self.reset_playing_source() - # clear all the items after "Screen Capture" - for i in range(4, self.ui.comboBox_camera_source.count()): - self.ui.comboBox_camera_source.removeItem(4) - - # populate the combobox with the sources - for source in camera_sources: - self.ui.comboBox_camera_source.addItem(source.description, source) - - self.ui.comboBox_camera_source.setEnabled(True) - currentIndexChangedSignal = QMetaMethod.fromSignal( - self.ui.comboBox_camera_source.currentIndexChanged - ) - if self.ui.comboBox_camera_source.isSignalConnected(currentIndexChangedSignal): - self.ui.comboBox_camera_source.currentIndexChanged.disconnect() - self.ui.comboBox_camera_source.currentIndexChanged.connect(self.sourceChanged) - - # enable the source view frame - self.ui.frame_source_view.setEnabled(True) - - selected_source_from_storage = fetch_data("scoresight.json", "source_selected") - if type(selected_source_from_storage) == str: - logger.info( - "Source selected from storage: %s", selected_source_from_storage - ) - # check if the source is a file path - if path.exists(selected_source_from_storage): - logger.info("File exists: %s", selected_source_from_storage) - self.ui.comboBox_camera_source.blockSignals(True) - self.ui.comboBox_camera_source.setCurrentIndex(1) - self.ui.comboBox_camera_source.blockSignals(False) - self.source_name = selected_source_from_storage - self.sourceSelectionSucessful() - else: - # select the last selected source - self.ui.comboBox_camera_source.setCurrentText( - selected_source_from_storage - ) - - def reset_playing_source(self): - if self.image_viewer: - # remove the image viewer from the layout frame_for_source_view_label - self.ui.frame_for_source_view_label.layout().removeWidget(self.image_viewer) - self.image_viewer.close() - self.image_viewer = None - # add a label with markdown text - label_select_source = QLabel("### Open a Camera or Load a File") - label_select_source.setTextFormat(Qt.TextFormat.MarkdownText) - label_select_source.setEnabled(False) - label_select_source.setAlignment(Qt.AlignmentFlag.AlignCenter) - clear_layout(self.ui.frame_for_source_view_label.layout()) - self.ui.frame_for_source_view_label.layout().addWidget(label_select_source) - - def sourceChanged(self, index): - # get the source name from the combobox - self.source_name = None - self.ui.groupBox_sb_info.setEnabled(False) - self.ui.tableWidget_boxes.setEnabled(False) - self.ui.pushButton_fourCorner.setEnabled(False) - self.ui.pushButton_binary.setEnabled(False) - if self.ui.comboBox_camera_source.currentIndex() == 0: - self.reset_playing_source() - return - elif self.ui.comboBox_camera_source.currentIndex() == 1: - logger.info("Open File selection dialog") - # open a file dialog to select a video file - file, _ = QFileDialog.getOpenFileName( - self, "Open Video File", "", "Video Files (*.mp4 *.avi *.mov)" - ) - if not file or not path.exists(file): - # no file selected - change source to "Select a source" - logger.error("No file selected") - self.ui.comboBox_camera_source.setCurrentText("Select a source") - return - else: - logger.info("File selected: %s", file) - self.source_name = file - elif self.ui.comboBox_camera_source.currentIndex() == 2: - # open a dialog to enter the url - url_dialog = QDialog() - ui_urlsource = Ui_UrlSource() - ui_urlsource.setupUi(url_dialog) - url_dialog.setWindowTitle("URL Source") - # focus on url input - ui_urlsource.lineEdit_url.setFocus() - url_dialog.exec() # wait for the dialog to close - # check if the dialog was accepted - if url_dialog.result() != QDialog.DialogCode.Accepted: - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - self.source_name = ui_urlsource.lineEdit_url.text() - if self.source_name == "": - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - elif self.ui.comboBox_camera_source.currentIndex() == 3: - # open a dialog to select the screen - screen_dialog = QDialog() - ui_screencapture = Ui_ScreenCapture() - ui_screencapture.setupUi(screen_dialog) - # set width and height of the dialog - screen_dialog.setFixedWidth(400) - - screen_dialog.setWindowTitle( - QCoreApplication.translate( - "MainWindow", "Screen Capture Selection", None - ) - ) - # populate comboBox_window with the available windows - ui_screencapture.comboBox_window.clear() - ui_screencapture.comboBox_window.addItem( - QCoreApplication.translate( - "MainWindow", "Capture the entire screen", None - ), - -1, - ) - for window in ScreenCapture.list_windows(): - ui_screencapture.comboBox_window.addItem(window[0], window[1]) - screen_dialog.exec() - # check if the dialog was accepted - if screen_dialog.result() != QDialog.DialogCode.Accepted: - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - # get the window ID from the comboBox_window - window_id = ui_screencapture.comboBox_window.currentData() - self.source_name = window_id - - # store the source selection in scoresight.json - self.globalSettingsChanged("source_selected", self.source_name) - self.sourceSelectionSucessful() - - def itemSelected(self, item_name): - # select the item in the tableWidget_boxes - items = self.ui.tableWidget_boxes.findItems( - item_name, Qt.MatchFlag.MatchExactly - ) - if len(items) == 0: - return - item = items[0] - item.setSelected(True) - self.ui.tableWidget_boxes.setCurrentItem(item) - self.listItemClicked(item) - - def fourCornersApplied(self, corners): - # check the button - self.ui.pushButton_fourCorner.setChecked(True) - - def sourceSelectionSucessful(self): - if self.ui.comboBox_camera_source.currentIndex() == 0: - return - - self.ui.frame_source_view.setEnabled(False) - - if self.ui.comboBox_camera_source.currentIndex() == 1: - if self.source_name is None or not path.exists(self.source_name): - logger.error("No file selected") - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - logger.info("Loading file selected: %s", self.source_name) - camera_info = CameraInfo( - self.source_name, - self.source_name, - self.source_name, - CameraInfo.CameraType.FILE, - ) - elif self.ui.comboBox_camera_source.currentIndex() == 2: - if self.source_name is None or self.source_name == "": - logger.error("No url entered") - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - logger.info("Loading url: %s", self.source_name) - camera_info = CameraInfo( - self.source_name, - self.source_name, - self.source_name, - CameraInfo.CameraType.URL, - ) - elif self.ui.comboBox_camera_source.currentIndex() == 3: - if self.source_name is None: - logger.error("No screen capture selected") - self.ui.comboBox_camera_source.setCurrentIndex(0) - return - logger.info("Loading screen capture: %s", self.source_name) - camera_info = CameraInfo( - self.source_name, - self.source_name, - self.source_name, - CameraInfo.CameraType.SCREEN_CAPTURE, - ) - else: - if self.ui.comboBox_camera_source.currentData() is None: - return - - logger.info( - "Loading camera: %s", self.ui.comboBox_camera_source.currentText() - ) - camera_info = self.ui.comboBox_camera_source.currentData() - - if self.image_viewer: - # remove the image viewer from the layout frame_for_source_view_label - self.ui.frame_for_source_view_label.layout().removeWidget(self.image_viewer) - self.image_viewer.close() - self.image_viewer = None - - # clear self.ui.frame_for_source_view_label - clear_layout(self.ui.frame_for_source_view_label.layout()) - - # set the pixmap to the image viewer - self.image_viewer = ImageViewer( - camera_info, - self.fourCornersApplied, - self.detectionTargetsStorage, - self.itemSelected, - ) - self.ui.toolButton_videoSettings.setEnabled( - camera_info.type == CameraInfo.CameraType.OPENCV - ) - self.ui.pushButton_fourCorner.setEnabled(True) - self.ui.pushButton_binary.setEnabled(True) - self.ui.pushButton_fourCorner.toggled.connect( - self.image_viewer.toggleFourCorner - ) - self.ui.pushButton_binary.clicked.connect(self.image_viewer.toggleBinary) - if self.image_viewer.timerThread: - self.image_viewer.timerThread.ocr_result_signal.connect(self.ocrResult) - self.image_viewer.timerThread.update_error.connect(self.updateError) - self.image_viewer.first_frame_received_signal.connect( - self.cameraConnectedEnableUI - ) - self.ocrModelChanged(fetch_data("scoresight.json", "ocr_model", 1)) - - # set the image viewer to the layout frame_for_source_view_label - self.ui.frame_for_source_view_label.layout().addWidget(self.image_viewer) - - def cameraConnectedEnableUI(self): - # enable groupBox_sb_info - self.ui.groupBox_sb_info.setEnabled(True) - self.ui.tableWidget_boxes.setEnabled(True) - self.ui.frame_source_view.setEnabled(True) - self.ui.widget_viewTools.setEnabled(True) - - # load the boxes from scoresight.json - self.detectionTargetsStorage.loadBoxesFromStorage() - self.updateError(None) - - def updateError(self, error): - if not error: - return - logger.error(error) - self.ui.frame_source_view.setEnabled(True) - self.ui.widget_viewTools.setEnabled(False) - - def ocrResult(self, results: list[TextDetectionTargetWithResult]): - # update template fields - for targetWithResult in results: - if not targetWithResult.settings["templatefield"]: - continue - targetWithResult.result = evaluate_template_field(results, targetWithResult) - targetWithResult.result_state = ( - TextDetectionTargetWithResult.ResultState.Success - if targetWithResult.result is not None - else TextDetectionTargetWithResult.ResultState.Empty - ) - - # update the table widget value items - for targetWithResult in results: - if ( - targetWithResult.result_state - == TextDetectionTargetWithResult.ResultState.Success - ): - items = self.ui.tableWidget_boxes.findItems( - targetWithResult.name, Qt.MatchFlag.MatchExactly - ) - if len(items) == 0: - continue - item = items[0] - # get the value (1 column) of the item - item = self.ui.tableWidget_boxes.item(item.row(), 1) - item.setText(targetWithResult.result) - - if not self.updateOCRResults: - # don't update the results, the user has disabled updates - return - - update_http_server(results) - - # update vmix - self.vmixUpdater.update_vmix(results) - - if self.out_folder is None: - return - - if not path.exists(path.abspath(self.out_folder)): - self.out_folder = None - remove_data("scoresight.json", "output_folder") - logger.warning("Output folder does not exist") - return - - # check if enough time has passed since last file save according to aggs per second - if ( - datetime.datetime.now() - self.last_aggregate_save - ).total_seconds() < 1.0 / self.ui.horizontalSlider_aggsPerSecond.value(): - return - - self.last_aggregate_save = datetime.datetime.now() - - # update the obs scene sources with the results, use update_text_source - for targetWithResult in results: - if targetWithResult.result is None: - continue - if ( - targetWithResult.settings is not None - and "skip_empty" in targetWithResult.settings - and targetWithResult.settings["skip_empty"] - and len(targetWithResult.result) == 0 - ): - continue - if ( - targetWithResult.result_state - != TextDetectionTargetWithResult.ResultState.Success - ): - continue - - if ( - self.obs_websocket_client is not None - and targetWithResult.settings is not None - ): - # find the source name for the target from the default boxes - update_text_source( - self.obs_websocket_client, - targetWithResult.settings["obs_source_name"], - targetWithResult.result, - ) - - # save the results to text files - file_output.save_text_files( - results, self.out_folder, self.ui.comboBox_appendMethod.currentIndex() - ) - - # save the results to a csv file - if self.ui.checkBox_saveCsv.isChecked(): - file_output.save_csv( - results, - self.out_folder, - self.ui.comboBox_appendMethod.currentIndex(), - self.first_csv_append, - ) - - # save the results to an xml file - if self.ui.checkBox_saveXML.isChecked(): - file_output.save_xml( - results, - self.out_folder, - ) - - def addBox(self): - # add a new box to the tableWidget_boxes - # find the number of custom boxes - custom_boxes = [] - for i in range(self.ui.tableWidget_boxes.rowCount()): - item = self.ui.tableWidget_boxes.item(i, 0) - if item.text() not in [o["name"] for o in default_boxes]: - custom_boxes.append(item.text()) - - i = len(custom_boxes) - new_box_name = f"Custom {i + 1}" - custom_boxes_names = fetch_data("scoresight.json", "custom_boxes_names", []) - # find if the name already exists - while new_box_name in custom_boxes or new_box_name in custom_boxes_names: - i += 1 - new_box_name = f"Custom {i + 1}" - - store_custom_box_name(new_box_name) - item = QTableWidgetItem( - QIcon( - path.abspath(path.join(path.dirname(__file__), "icons/circle-x.svg")) - ), - new_box_name, - ) - item.setData(Qt.ItemDataRole.UserRole, "unchecked") - self.ui.tableWidget_boxes.insertRow(self.ui.tableWidget_boxes.rowCount()) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 0, item - ) - disabledItem = QTableWidgetItem() - disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) - self.ui.tableWidget_boxes.setItem( - self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem - ) - - def removeCustomBox(self): - item = self.ui.tableWidget_boxes.currentItem() - if not item: - logger.info("No item selected") - return - if item.column() != 0: - item = self.ui.tableWidget_boxes.item(item.row(), 0) - self.removeBox() - remove_custom_box_name_in_storage(item.text()) - # only allow removing custom boxes - if item.text() in [o["name"] for o in default_boxes]: - logger.info("Cannot remove default box") - return - # remove the selected item from the tableWidget_boxes - self.ui.tableWidget_boxes.removeRow(item.row()) - - def editBoxName(self, item): - if item.text() in [o["name"] for o in default_boxes]: - # dont allow editing default boxes - return - new_name, ok = QInputDialog.getText( - self, "Edit Box Name", "New Name:", text=item.text() - ) - old_name = item.text() - if ok and new_name != "" and new_name != old_name: - # check if name doesn't exist already - for i in range(self.ui.tableWidget_boxes.rowCount()): - if new_name == self.ui.tableWidget_boxes.item(i, 0).text(): - logger.info("Name '%s' already exists", new_name) - return - # rename the item in the tableWidget_boxes - rename_custom_box_name_in_storage(old_name, new_name) - item.setText(new_name) - # rename the item in the detectionTargetsStorage - if not self.detectionTargetsStorage.rename_item(old_name, new_name): - logger.info("Error renaming item in application storage") - return - else: - # check if the item role isn't "templatefield" - if item.data(Qt.ItemDataRole.UserRole) != "templatefield": - # remove the item from the tableWidget_boxes - self.ui.tableWidget_boxes.removeRow(item.row()) - - def makeBox(self): - item = self.ui.tableWidget_boxes.currentItem() - if not item: - return - # create a new box on self.image_viewer with the name of the selected item from the tableWidget_boxes - # change the list icon to green checkmark - item.setIcon( - QIcon( - path.abspath( - path.join(path.dirname(__file__), "icons/circle-check.svg") - ) - ) - ) - item.setData(Qt.ItemDataRole.UserRole, "checked") - self.listItemClicked(item) - - # get the size of the box from the name - info = default_info_for_box_name(item.text()) - - self.detectionTargetsStorage.add_item( - TextDetectionTarget( - info["x"], - info["y"], - info["width"], - info["height"], - item.text(), - normalize_settings_dict({}, info), - ) - ) - - def removeBox(self): - item = self.ui.tableWidget_boxes.currentItem() - if not item: - return - # change the list icon to red x - item.setIcon( - QIcon(path.abspath(path.join(path.dirname(__file__), "icons/circle-x.svg"))) - ) - item.setData(Qt.ItemDataRole.UserRole, "unchecked") - self.listItemClicked(item) - self.detectionTargetsStorage.remove_item(item.text()) - - def makeTemplateField(self, toggled: bool): - item = self.ui.tableWidget_boxes.currentItem() - if not item: - return - - if not toggled: - self.removeBox() - return - - # create a new box on self.image_viewer with the name of the selected item from the tableWidget_boxes - # change the list icon to green checkmark - item.setIcon( - QIcon( - path.abspath( - path.join(path.dirname(__file__), "icons/template-field.svg") - ) - ) - ) - item.setData(Qt.ItemDataRole.UserRole, "templatefield") - - self.detectionTargetsStorage.add_item( - TextDetectionTarget( - 0, - 0, - 0, - 0, - item.text(), - normalize_settings_dict({"templatefield": True}, None), - ) - ) - - self.listItemClicked(item) - - def createOBSScene(self): - # get the scene name from the lineEdit_sceneName - scene_name = self.ui.lineEdit_sceneName.text() - # clear or create a new scene - create_obs_scene_from_export(self.obs_websocket_client, scene_name) - - # on destroy, close the OBS connection - def closeEvent(self, event): - logger.info("Closing") - if self.image_viewer: - self.image_viewer.close() - self.image_viewer = None - - if self.log_dialog: - self.log_dialog.close() - self.log_dialog = None - - # store the boxes to scoresight.json - self.detectionTargetsStorage.saveBoxesToStorage() - if self.obs_websocket_client: - # destroy the client object - self.obs_websocket_client = None - - super().closeEvent(event) - - -if __name__ == "__main__": - # only attempt splash when not on Mac OSX - os_name = platform.system() - if os_name != "Darwin": - try: - import pyi_splash # type: ignore - - pyi_splash.close() - except ImportError: - pass - app = QApplication(sys.argv) - - # Get system locale - locale = QLocale.system().name() - - # Load the translation file based on the locale - translator = QTranslator() - locale_file = path.abspath( - path.join(path.dirname(__file__), "translations", f"scoresight_{locale}.qm") - ) - # check if the file exists - if not path.exists(locale_file): - # load the default translation file - locale_file = path.abspath( - path.join(path.dirname(__file__), "translations", "scoresight_en_US.qm") - ) - if translator.load(locale_file): - app.installTranslator(translator) - - # show the main window - mainWindow = MainWindow(translator, app) - mainWindow.show() - - app.exec() - logger.info("Exiting...") - - stop_http_server() +from os import path +import platform +import sys +from PySide6.QtWidgets import QApplication +from PySide6.QtCore import ( + QTranslator, + QLocale, +) + +from mainwindow import MainWindow +from sc_logging import logger +from http_server import stop_http_server + +if __name__ == "__main__": + # only attempt splash when not on Mac OSX + os_name = platform.system() + if os_name != "Darwin": + try: + import pyi_splash # type: ignore + + pyi_splash.close() + except ImportError: + pass + app = QApplication(sys.argv) + + # Get system locale + locale = QLocale.system().name() + + # Load the translation file based on the locale + translator = QTranslator() + locale_file = path.abspath( + path.join(path.dirname(__file__), "translations", f"scoresight_{locale}.qm") + ) + # check if the file exists + if not path.exists(locale_file): + # load the default translation file + locale_file = path.abspath( + path.join(path.dirname(__file__), "translations", "scoresight_en_US.qm") + ) + if translator.load(locale_file): + app.installTranslator(translator) + + # show the main window + mainWindow = MainWindow(translator, app) + mainWindow.show() + + app.exec() + logger.info("Exiting...") + + stop_http_server() diff --git a/mainwindow.py b/mainwindow.py new file mode 100644 index 0000000..4563aef --- /dev/null +++ b/mainwindow.py @@ -0,0 +1,1549 @@ +from functools import partial +import os +import platform +import datetime +from PySide6.QtWidgets import ( + QApplication, + QDialog, + QFileDialog, + QInputDialog, + QLabel, + QMainWindow, + QMenu, + QMessageBox, + QTableWidgetItem, +) +from PySide6.QtGui import QIcon, QStandardItemModel, QStandardItem, QDesktopServices +from PySide6.QtCore import ( + Qt, + Signal, + Slot, + QTranslator, + QObject, + QCoreApplication, + QEvent, + QMetaMethod, + QUrl, +) +from dotenv import load_dotenv +from os import path +from platformdirs import user_data_dir + +from camera_info import CameraInfo +from get_camera_info import get_camera_info +from http_server import start_http_server, update_http_server +from screen_capture_source import ScreenCapture +from source_view import ImageViewer +from defaults import ( + default_boxes, + default_info_for_box_name, + normalize_settings_dict, + format_prefixes, +) + +from storage import ( + TextDetectionTargetMemoryStorage, + fetch_data, + remove_data, + store_data, + store_custom_box_name, + rename_custom_box_name_in_storage, + remove_custom_box_name_in_storage, + fetch_custom_box_names, +) +from obs_websocket import ( + create_obs_scene_from_export, + open_obs_websocket, + update_text_source, +) + +from template_fields import evaluate_template_field +from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult +from sc_logging import logger +from update_check import check_for_updates +from log_view import LogViewerDialog +import file_output +from video_settings import VideoSettingsDialog +from vmix_output import VMixAPI +from ui_mainwindow import Ui_MainWindow +from ui_about import Ui_Dialog as Ui_About +from ui_connect_obs import Ui_Dialog as Ui_ConnectObs +from ui_url_source import Ui_Dialog as Ui_UrlSource +from ui_screen_capture import Ui_Dialog as Ui_ScreenCapture + + +def clear_layout(layout): + while layout.count(): + item = layout.takeAt(0) + widget = item.widget() + if widget is not None: + widget.deleteLater() + widget = None + else: + clear_layout(item.layout()) + + +class MainWindow(QMainWindow): + # add a signal to update sources + update_sources = Signal(list) + get_sources = Signal() + + def __init__(self, translator: QTranslator, parent: QObject): + super(MainWindow, self).__init__() + self.parent_object = parent + self.ui = Ui_MainWindow() + logger.info("Starting ScoreSight") + self.ui.setupUi(self) + self.translator = translator + # load env variables + load_dotenv() + self.setWindowTitle(f"ScoreSight - v{os.getenv('LOCAL_RELEASE_TAG')}") + if platform.system() == "Windows": + # set the icon + self.setWindowIcon( + QIcon( + path.abspath( + path.join(path.dirname(__file__), "icons/Windows-icon-open.ico") + ) + ) + ) + + self.menubar = self.menuBar() + file_menu = self.menubar.addMenu("File") + + # check for updates + check_for_updates(False) + file_menu.addAction("Check for Updates", lambda: check_for_updates(True)) + file_menu.addAction("About", self.openAboutDialog) + file_menu.addAction("View Current Log", self.openLogsDialog) + file_menu.addAction("Import Configuration", self.importConfiguration) + file_menu.addAction("Export Configuration", self.exportConfiguration) + file_menu.addAction("Open Configuration Folder", self.openConfigurationFolder) + + # Add "Language" menu + languageMenu = file_menu.addMenu("Language") + + # Add language options + self.addLanguageOption(languageMenu, "English (US)", "en_US") + self.addLanguageOption(languageMenu, "French (France)", "fr_FR") + self.addLanguageOption(languageMenu, "Spanish (Spain)", "es_ES") + self.addLanguageOption(languageMenu, "German", "de_DE") + self.addLanguageOption(languageMenu, "Italian", "it_IT") + self.addLanguageOption(languageMenu, "Japanese", "ja_JP") + self.addLanguageOption(languageMenu, "Korean", "ko_KR") + self.addLanguageOption(languageMenu, "Dutch", "nl_NL") + self.addLanguageOption(languageMenu, "Polish", "pl_PL") + self.addLanguageOption(languageMenu, "Portuguese (Brazil)", "pt_BR") + self.addLanguageOption(languageMenu, "Portuguese (Portugal)", "pt_PT") + self.addLanguageOption(languageMenu, "Russian", "ru_RU") + self.addLanguageOption(languageMenu, "Chinese (Simplified)", "zh_CN") + + # Hide the menu bar by default + self.menubar.setVisible(False) + + # Show the menu bar when the Alt key is pressed + self.installEventFilter(self) + + self.ui.pushButton_connectObs.clicked.connect(self.openOBSConnectModal) + + self.vmixUiSetup() + + start_http_server() + + self.ui.pushButton_stabilize.setEnabled(True) + self.ui.pushButton_stabilize.clicked.connect(self.toggleStabilize) + + self.ui.toolButton_topCrop.clicked.connect(self.cropMode) + # check configuation if crop is enabled + self.ui.toolButton_topCrop.setChecked( + fetch_data("scoresight.json", "crop_mode", False) + ) + self.ui.widget_cropPanel.setVisible(self.ui.toolButton_topCrop.isChecked()) + self.ui.widget_cropPanel.setEnabled(self.ui.toolButton_topCrop.isChecked()) + self.ui.spinBox_leftCrop.valueChanged.connect( + partial(self.globalSettingsChanged, "left_crop") + ) + self.ui.spinBox_leftCrop.setValue(fetch_data("scoresight.json", "left_crop", 0)) + self.ui.spinBox_rightCrop.valueChanged.connect( + partial(self.globalSettingsChanged, "right_crop") + ) + self.ui.spinBox_rightCrop.setValue( + fetch_data("scoresight.json", "right_crop", 0) + ) + self.ui.spinBox_topCrop.valueChanged.connect( + partial(self.globalSettingsChanged, "top_crop") + ) + self.ui.spinBox_topCrop.setValue(fetch_data("scoresight.json", "top_crop", 0)) + self.ui.spinBox_bottomCrop.valueChanged.connect( + partial(self.globalSettingsChanged, "bottom_crop") + ) + self.ui.spinBox_bottomCrop.setValue( + fetch_data("scoresight.json", "bottom_crop", 0) + ) + self.ui.checkBox_enableOutAPI.toggled.connect( + partial(self.globalSettingsChanged, "enable_out_api") + ) + self.ui.checkBox_enableOutAPI.setChecked( + fetch_data("scoresight.json", "enable_out_api", False) + ) + + # connect toolButton_rotate + self.ui.toolButton_rotate.clicked.connect(self.rotateImage) + + self.ui.widget_detectionCadence.setVisible(True) + self.ui.horizontalSlider_detectionCadence.setValue( + fetch_data("scoresight.json", "detection_cadence", 5) + ) + self.ui.horizontalSlider_detectionCadence.valueChanged.connect( + self.detectionCadenceChanged + ) + self.ui.toolButton_addBox.clicked.connect(self.addBox) + self.ui.toolButton_removeBox.clicked.connect(self.removeCustomBox) + + self.video_settings_dialog = None + self.ui.toolButton_videoSettings.clicked.connect(self.openVideoSettings) + + self.obs_websocket_client = None + + ocr_models = [ + "Daktronics", + "General Scoreboard", + "General Fonts (English)", + "General Scoreboard Large", + ] + self.ui.comboBox_ocrModel.addItems(ocr_models) + self.ui.comboBox_ocrModel.setCurrentIndex( + fetch_data("scoresight.json", "ocr_model", 1) + ) # default to General Scoreboard + + self.ui.frame_source_view.setEnabled(False) + self.ui.groupBox_target_settings.setEnabled(False) + self.ui.pushButton_makeBox.clicked.connect(self.makeBox) + self.ui.pushButton_removeBox.clicked.connect(self.removeBox) + self.ui.tableWidget_boxes.itemClicked.connect(self.listItemClicked) + # connect the edit triggers + self.ui.tableWidget_boxes.itemDoubleClicked.connect(self.editBoxName) + self.ui.pushButton_refresh_sources.clicked.connect( + lambda: self.get_sources.emit() + ) + self.detectionTargetsStorage = TextDetectionTargetMemoryStorage() + self.detectionTargetsStorage.data_changed.connect(self.detectionTargetsChanged) + self.ui.pushButton_createOBSScene.clicked.connect(self.createOBSScene) + self.ui.pushButton_selectFolder.clicked.connect(self.selectOutputFolder) + self.ui.toolButton_trashFolder.clicked.connect(self.clearOutputFolder) + self.ui.pushButton_stopUpdates.toggled.connect(self.toggleStopUpdates) + self.ui.comboBox_ocrModel.currentIndexChanged.connect(self.ocrModelChanged) + self.ui.pushButton_restoreDefaults.clicked.connect(self.restoreDefaults) + self.ui.toolButton_zoomReset.clicked.connect(self.resetZoom) + self.ui.toolButton_osd.toggled.connect(self.toggleOSD) + self.ui.toolButton_showOCRrects.toggled.connect(self.toggleOCRRects) + self.ui.checkBox_smoothing.toggled.connect( + partial(self.genericSettingsChanged, "smoothing") + ) + self.ui.checkBox_skip_empty.toggled.connect( + partial(self.genericSettingsChanged, "skip_empty") + ) + self.ui.horizontalSlider_conf_thresh.valueChanged.connect( + self.confThreshChanged + ) + self.ui.lineEdit_format.textChanged.connect( + partial(self.genericSettingsChanged, "format_regex") + ) + self.ui.comboBox_fieldType.currentIndexChanged.connect( + partial(self.genericSettingsChanged, "type") + ) + self.ui.checkBox_skip_similar_image.toggled.connect( + partial(self.genericSettingsChanged, "skip_similar_image") + ) + self.ui.checkBox_autocrop.toggled.connect( + partial(self.genericSettingsChanged, "autocrop") + ) + self.ui.horizontalSlider_cleanup.valueChanged.connect(self.cleanupThreshChanged) + self.ui.horizontalSlider_dilate.valueChanged.connect( + partial(self.genericSettingsChanged, "dilate") + ) + self.ui.horizontalSlider_skew.valueChanged.connect( + partial(self.genericSettingsChanged, "skew") + ) + self.ui.horizontalSlider_vscale.valueChanged.connect( + partial(self.genericSettingsChanged, "vscale") + ) + self.ui.checkBox_removeLeadingZeros.toggled.connect( + partial(self.genericSettingsChanged, "remove_leading_zeros") + ) + self.ui.checkBox_rescalePatch.toggled.connect( + partial(self.genericSettingsChanged, "rescale_patch") + ) + self.ui.checkBox_normWHRatio.toggled.connect( + partial(self.genericSettingsChanged, "normalize_wh_ratio") + ) + self.ui.checkBox_invertPatch.toggled.connect( + partial(self.genericSettingsChanged, "invert_patch") + ) + self.ui.checkBox_dotDetector.toggled.connect( + partial(self.genericSettingsChanged, "dot_detector") + ) + self.ui.checkBox_ordinalIndicator.toggled.connect( + partial(self.genericSettingsChanged, "ordinal_indicator") + ) + self.ui.comboBox_binarizationMethod.currentIndexChanged.connect( + partial(self.genericSettingsChanged, "binarization_method") + ) + self.ui.checkBox_templatefield.toggled.connect(self.makeTemplateField) + self.ui.lineEdit_templatefield.textChanged.connect( + partial(self.genericSettingsChanged, "templatefield_text") + ) + self.ui.comboBox_formatPrefix.currentIndexChanged.connect( + self.formatPrefixChanged + ) + self.ui.checkBox_updateOnchange.toggled.connect(self.toggleUpdateOnChange) + + # populate the tableWidget_boxes with the default and custom boxes + custom_boxes_names = fetch_custom_box_names() + + for box_name in [box["name"] for box in default_boxes] + custom_boxes_names: + item = QTableWidgetItem( + QIcon( + path.abspath( + path.join(path.dirname(__file__), "icons/circle-x.svg") + ) + ), + box_name, + ) + item.setData(Qt.ItemDataRole.UserRole, "unchecked") + self.ui.tableWidget_boxes.insertRow(self.ui.tableWidget_boxes.rowCount()) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 0, item + ) + disabledItem = QTableWidgetItem() + disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem + ) + + self.image_viewer = None + self.obs_connect_modal = None + self.source_name = None + self.updateOCRResults = True + self.log_dialog = None + if fetch_data("scoresight.json", "open_on_startup", False): + logger.info("Opening log dialog on startup") + self.openLogsDialog() + + if fetch_data("scoresight.json", "obs"): + self.connectObs() + + self.out_folder = fetch_data("scoresight.json", "output_folder") + if self.out_folder: + if not path.exists(self.out_folder): + self.out_folder = None + remove_data("scoresight.json", "output_folder") + else: + self.ui.lineEdit_folder.setText(self.out_folder) + + self.first_csv_append = True + self.last_aggregate_save = datetime.datetime.now() + self.ui.checkBox_saveCsv.toggled.connect( + partial(self.globalSettingsChanged, "save_csv") + ) + self.ui.checkBox_saveCsv.setChecked( + fetch_data("scoresight.json", "save_csv", False) + ) + self.ui.checkBox_saveXML.toggled.connect( + partial(self.globalSettingsChanged, "save_xml") + ) + self.ui.checkBox_saveXML.setChecked( + fetch_data("scoresight.json", "save_xml", False) + ) + self.ui.comboBox_appendMethod.currentIndexChanged.connect( + partial(self.globalSettingsChanged, "append_method") + ) + self.ui.horizontalSlider_aggsPerSecond.valueChanged.connect( + partial(self.globalSettingsChanged, "aggs_per_second") + ) + self.ui.comboBox_appendMethod.setCurrentIndex( + fetch_data("scoresight.json", "append_method", 3) + ) + self.ui.horizontalSlider_aggsPerSecond.setValue( + fetch_data("scoresight.json", "aggs_per_second", 5) + ) + self.ui.checkBox_updateOnchange.setChecked( + fetch_data("scoresight.json", "update_on_change", True) + ) + self.ui.checkBox_vmix_send_same.setChecked( + fetch_data("scoresight.json", "vmix_send_same", False) + ) + self.ui.checkBox_vmix_send_same.toggled.connect( + partial(self.globalSettingsChanged, "vmix_send_same") + ) + + self.update_sources.connect(self.updateSources) + self.get_sources.connect(self.getSources) + self.get_sources.emit() + + def openVideoSettings(self): + # only allow opening the video settings for an OpenCV type source + if not self.image_viewer or self.image_viewer is None: + return + + if self.image_viewer.getCameraInfo().type != CameraInfo.CameraType.OPENCV: + return + + if self.video_settings_dialog is None: + # open the logs dialog + self.video_settings_dialog = VideoSettingsDialog() + self.video_settings_dialog.setWindowTitle("Video Settings") + + self.video_settings_dialog.init_ui(self.image_viewer.getCameraCapture()) + + # show the dialog, non modal + self.video_settings_dialog.show() + + def rotateImage(self): + # store the rotation in the scoresight.json + rotation = fetch_data("scoresight.json", "rotation", 0) + rotation += 90 + if rotation >= 360: + rotation = 0 + self.globalSettingsChanged("rotation", rotation) + + def cropMode(self): + # if the toolButton_topCrop is unchecked, go to crop mode + if self.ui.toolButton_topCrop.isChecked(): + self.ui.widget_cropPanel.setVisible(True) + self.ui.widget_cropPanel.setEnabled(True) + self.globalSettingsChanged("crop_mode", True) + else: + self.ui.widget_cropPanel.setVisible(False) + self.ui.widget_cropPanel.setEnabled(False) + self.globalSettingsChanged("crop_mode", False) + + def globalSettingsChanged(self, settingName, value): + store_data("scoresight.json", settingName, value) + + def eventFilter(self, obj, event): + if event.type() == QEvent.KeyPress and event.key() == Qt.Key_Alt: + self.menubar.setVisible(True) + elif event.type() == QEvent.FocusOut and self.menubar.isVisible(): + self.menubar.setVisible(False) + elif event.type() == QEvent.WindowDeactivate and self.menubar.isVisible(): + self.menubar.setVisible(False) + return super().eventFilter(obj, event) + + def focusOutEvent(self, event): + if self.menubar.isVisible(): + self.menubar.setVisible(False) + super().focusOutEvent(event) + + def changeEvent(self, event): + if event.type() == QEvent.WindowDeactivate and self.menubar.isVisible(): + self.menubar.setVisible(False) + super().changeEvent(event) + + def changeLanguage(self, locale): + locale_file = path.abspath( + path.join(path.dirname(__file__), "translations", f"scoresight_{locale}.qm") + ) + logger.info(f"Changing language to {locale_file}") + if not self.translator.load(locale_file): + logger.error(f"Could not load translation for {locale_file}") + return + appInstance = QApplication.instance() + if appInstance: + logger.info(f"installing translator for {locale}") + appInstance.installTranslator(self.translator) + self.ui.retranslateUi(self) + + def addLanguageOption(self, menu: QMenu, language_name: str, locale: str): + menu.addAction(language_name, lambda: self.changeLanguage(locale)) + + def toggleUpdateOnChange(self, value): + self.globalSettingsChanged("update_on_change", value) + if self.image_viewer: + self.image_viewer.setUpdateOnChange(value) + + def formatPrefixChanged(self, index): + if index == 12: + return # do nothing if "Select Preset" is selected + # based on the selected index, set the format prefix + # change lineEdit_format to the selected format prefix + self.ui.lineEdit_format.setText(format_prefixes[index]) + + def importConfiguration(self): + # open a file dialog to select a configuration file + file, _ = QFileDialog.getOpenFileName( + self, "Open Configuration File", "", "Configuration Files (*.json)" + ) + if not file: + return + # load the configuration from the file + if not self.detectionTargetsStorage.loadBoxesFromFile(file): + # show an error qmessagebox + logger.error("Error loading configuration file") + QMessageBox.critical( + self, + "Error", + "Error loading configuration file", + QMessageBox.StandardButton.Ok, + ) + return + + def exportConfiguration(self): + # open a file dialog to select the output file + file, _ = QFileDialog.getSaveFileName( + self, "Save Configuration File", "", "Configuration Files (*.json)" + ) + if not file: + return + # save the configuration to the file + self.detectionTargetsStorage.saveBoxesToFile(file) + + def openConfigurationFolder(self): + # open the configuration folder in the file explorer + QDesktopServices.openUrl( + QUrl( + "file:///" + user_data_dir("scoresight"), QUrl.ParsingMode.TolerantMode + ) + ) + + def toggleOSD(self, value): + if self.image_viewer: + self.image_viewer.toggleOSD(value) + + def toggleOCRRects(self, value): + if self.image_viewer: + self.image_viewer.toggleOCRRects(value) + + def resetZoom(self): + if self.image_viewer: + self.image_viewer.resetZoom() + + def detectionCadenceChanged(self, detections_per_second): + self.globalSettingsChanged("detection_cadence", detections_per_second) + if self.image_viewer and self.image_viewer.timerThread: + # convert the detections_per_second to milliseconds + self.image_viewer.timerThread.update_frame_interval = ( + 1000 / detections_per_second + ) + + def ocrModelChanged(self, index): + self.globalSettingsChanged("ocr_model", index) + # update the ocr model in the text detector + if ( + self.image_viewer + and self.image_viewer.timerThread + and self.image_viewer.timerThread.textDetector + ): + self.image_viewer.timerThread.textDetector.setOcrModel(index) + + def openLogsDialog(self): + if self.log_dialog is None: + # open the logs dialog + self.log_dialog = LogViewerDialog() + self.log_dialog.setWindowTitle("Logs") + + # show the dialog, non modal + self.log_dialog.show() + + def openAboutDialog(self): + # open the about dialog + about_dialog = QDialog() + about_dialog_ui = Ui_About() + about_dialog_ui.setupUi(about_dialog) + about_dialog.setWindowTitle("About ScoreSight") + about_dialog.exec() + + def toggleStabilize(self): + if not self.image_viewer: + return + # start or stop the stabilization + self.image_viewer.toggleStabilization(self.ui.pushButton_stabilize.isChecked()) + + def toggleStopUpdates(self, value): + self.updateOCRResults = not value + # change the text on the button + self.ui.pushButton_stopUpdates.setText( + self.translator.translate("main", "Resume updates") + if value + else self.translator.translate("main", "Stop updates") + ) + + def selectOutputFolder(self): + # open a Qt dialog to select the output folder + folder = QFileDialog.getExistingDirectory( + self, + "Select Output Folder", + fetch_data("scoresight.json", "output_folder", ""), + options=QFileDialog.Option.ShowDirsOnly, + ) + if folder and len(folder) > 0: + self.ui.lineEdit_folder.setText(folder) + self.out_folder = folder + self.globalSettingsChanged("output_folder", folder) + + def clearOutputFolder(self): + # clear the output folder + self.ui.lineEdit_folder.setText("") + self.out_folder = None + remove_data("scoresight.json", "output_folder") + + def editSettings(self, settingsMutatorCallback): + # update the selected item's settings in the detectionTargetsStorage + item = self.ui.tableWidget_boxes.currentItem() + if item is None: + logger.info("no item selected") + return + item_name = item.text() + item_obj = self.detectionTargetsStorage.find_item_by_name(item_name) + if item_obj is None: + logger.info("item not found: %s", item_name) + return + item_obj = settingsMutatorCallback(item_obj) + self.detectionTargetsStorage.edit_item(item_name, item_obj) + + def restoreDefaults(self): + # restore the default settings for the selected item + def restoreDefaultsSettings(item_obj): + info = default_info_for_box_name(item_obj.name) + item_obj.settings = normalize_settings_dict({}, info) + return item_obj + + self.editSettings(restoreDefaultsSettings) + self.populateSettings(self.ui.tableWidget_boxes.currentItem().text()) + + def confThreshChanged(self): + def editConfThreshSettings(item_obj): + item_obj.settings["conf_thresh"] = ( + float(self.ui.horizontalSlider_conf_thresh.value()) / 100.0 + ) + return item_obj + + self.editSettings(editConfThreshSettings) + + def cleanupThreshChanged(self): + def editCleanupThreshSettings(item_obj): + item_obj.settings["cleanup_thresh"] = ( + float(self.ui.horizontalSlider_cleanup.value()) / 100.0 + ) + return item_obj + + self.editSettings(editCleanupThreshSettings) + + def genericSettingsChanged(self, settingName, value): + def editGenericSettings(item_obj): + item_obj.settings[settingName] = value + return item_obj + + self.editSettings(editGenericSettings) + + def vmixConnectionChanged(self): + self.vmixUpdater = VMixAPI( + self.ui.lineEdit_vmixHost.text(), + self.ui.lineEdit_vmixPort.text(), + self.ui.inputLineEdit_vmix.text(), + {}, + ) + self.globalSettingsChanged("vmix_host", self.ui.lineEdit_vmixHost.text()) + self.globalSettingsChanged("vmix_port", self.ui.lineEdit_vmixPort.text()) + self.globalSettingsChanged("vmix_input", self.ui.inputLineEdit_vmix.text()) + + def vmixMappingChanged(self, _): + # store entire mapping data in scoresight.json + mapping = {} + model = self.ui.tableView_vmixMapping.model() + if isinstance(model, QStandardItemModel): + for i in range(model.rowCount()): + item = model.item(i, 0) + value = model.item(i, 1) + if item and value: + mapping[item.text()] = value.text() + self.globalSettingsChanged("vmix_mapping", mapping) + self.vmixUpdater.set_field_mapping(mapping) + else: + logger.error("vmixMappingChanged: model is not a QStandardItemModel") + + def vmixUiSetup(self): + # populate the vmix connection from storage + self.ui.lineEdit_vmixHost.setText( + fetch_data("scoresight.json", "vmix_host", "localhost") + ) + self.ui.lineEdit_vmixPort.setText( + fetch_data("scoresight.json", "vmix_port", "8099") + ) + self.ui.inputLineEdit_vmix.setText( + fetch_data("scoresight.json", "vmix_input", "1") + ) + # connect the lineEdits to vmixConnectionChanged + self.ui.lineEdit_vmixHost.textChanged.connect(self.vmixConnectionChanged) + self.ui.lineEdit_vmixPort.textChanged.connect(self.vmixConnectionChanged) + self.ui.inputLineEdit_vmix.textChanged.connect(self.vmixConnectionChanged) + + # create the vmixUpdater + self.vmixUpdater = VMixAPI( + self.ui.lineEdit_vmixHost.text(), + self.ui.lineEdit_vmixPort.text(), + self.ui.inputLineEdit_vmix.text(), + {}, + ) + # add standard item model to the tableView_vmixMapping + self.ui.tableView_vmixMapping.setModel(QStandardItemModel()) + mapping = fetch_data("scoresight.json", "vmix_mapping", {}) + if mapping: + self.vmixUpdater.set_field_mapping(mapping) + + self.ui.tableView_vmixMapping.model().dataChanged.connect( + self.vmixMappingChanged + ) + + self.ui.pushButton_startvmix.toggled.connect(self.togglevMix) + + def togglevMix(self, value): + if not self.vmixUpdater: + return + if value: + self.ui.pushButton_startvmix.setText("🛑 Stop vMix") + self.vmixUpdater.running = True + else: + self.ui.pushButton_startvmix.setText("▶️ Start vMix") + self.vmixUpdater.running = False + + def updatevMixTable(self, detectionTargets): + mapping_storage = fetch_data("scoresight.json", "vmix_mapping") + model = QStandardItemModel() + model.blockSignals(True) + + for box in detectionTargets: + # add the detection to the vmix output mapping: tableView_vmixMapping + # check if the table already has the detectionTarget + items = model.findItems(box.name, Qt.MatchFlag.MatchExactly) + if len(items) == 0: + # add the item to the list + row = model.rowCount() + model.insertRow(row) + model.setItem(row, 0, QStandardItem(box.name)) + # the first item shouldn't be editable + model.item(row, 0).setFlags(Qt.ItemFlag.NoItemFlags) + else: + # update the item in the list + item = items[0] + row = item.row() + + # get value from storage or use the box name + if mapping_storage and box.name in mapping_storage: + model.setItem(row, 1, QStandardItem(mapping_storage[box.name])) + else: + model.setItem(row, 1, QStandardItem(box.name)) + # remove the items that are not in the detectionTargets + for i in range(model.rowCount()): + item = model.item(i, 0) + if not any([box.name == item.text() for box in detectionTargets]): + model.removeRow(i) + + model.blockSignals(False) + self.ui.tableView_vmixMapping.setModel(model) + self.ui.tableView_vmixMapping.model().dataChanged.connect( + self.vmixMappingChanged + ) + + def detectionTargetsChanged(self, detectionTargets): + for box in detectionTargets: + logger.debug(f"Change: Detection target: {box.name}") + # change the list icon to green checkmark + items = self.ui.tableWidget_boxes.findItems( + box.name, Qt.MatchFlag.MatchExactly + ) + if len(items) == 0: + logger.warning(f"Item not found: {box.name}. Adding it to the list.") + # add the item to the list + self.ui.tableWidget_boxes.insertRow( + self.ui.tableWidget_boxes.rowCount() + ) + item = QTableWidgetItem(box.name) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 0, item + ) + disabledItem = QTableWidgetItem() + disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem + ) + else: + item = items[0] + + if not box.settings["templatefield"]: + # this is a detection target + item.setIcon( + QIcon( + path.abspath( + path.join(path.dirname(__file__), "icons/circle-check.svg") + ) + ) + ) + item.setData(Qt.ItemDataRole.UserRole, "checked") + else: + # this is a template field + item.setIcon( + QIcon( + path.abspath( + path.join( + path.dirname(__file__), "icons/template-field.svg" + ) + ) + ) + ) + item.setData(Qt.ItemDataRole.UserRole, "templatefield") + + self.updatevMixTable(detectionTargets) + + # if save_csv is enabled, truncate the aggregate file + if self.ui.checkBox_saveCsv.isChecked() and self.out_folder: + csv_output_file_path = path.abspath( + path.join(self.out_folder, "results.csv") + ) + try: + with open(csv_output_file_path, "w") as f: + f.write("") + self.first_csv_append = True + self.last_aggregate_save = datetime.datetime.now() + except Exception as e: + logger.error(f"Error truncating aggregate file: {e}") + + def populateSettings(self, name): + self.ui.lineEdit_format.blockSignals(True) + self.ui.comboBox_fieldType.blockSignals(True) + self.ui.checkBox_smoothing.blockSignals(True) + self.ui.checkBox_skip_empty.blockSignals(True) + self.ui.horizontalSlider_conf_thresh.blockSignals(True) + self.ui.checkBox_autocrop.blockSignals(True) + self.ui.checkBox_skip_similar_image.blockSignals(True) + self.ui.horizontalSlider_cleanup.blockSignals(True) + self.ui.horizontalSlider_dilate.blockSignals(True) + self.ui.horizontalSlider_skew.blockSignals(True) + self.ui.horizontalSlider_vscale.blockSignals(True) + self.ui.checkBox_removeLeadingZeros.blockSignals(True) + self.ui.checkBox_rescalePatch.blockSignals(True) + self.ui.checkBox_normWHRatio.blockSignals(True) + self.ui.checkBox_invertPatch.blockSignals(True) + self.ui.checkBox_ordinalIndicator.blockSignals(True) + self.ui.comboBox_binarizationMethod.blockSignals(True) + self.ui.comboBox_formatPrefix.blockSignals(True) + self.ui.checkBox_templatefield.blockSignals(True) + self.ui.lineEdit_templatefield.blockSignals(True) + + # populate the settings from the detectionTargetsStorage + item_obj = self.detectionTargetsStorage.find_item_by_name(name) + if item_obj is None: + self.ui.lineEdit_format.setText("") + self.ui.comboBox_fieldType.setCurrentIndex(0) + self.ui.checkBox_smoothing.setChecked(True) + self.ui.checkBox_skip_empty.setChecked(True) + self.ui.horizontalSlider_conf_thresh.setValue(50) + self.ui.checkBox_autocrop.setChecked(False) + self.ui.checkBox_skip_similar_image.setChecked(False) + self.ui.horizontalSlider_cleanup.setValue(0) + self.ui.horizontalSlider_dilate.setValue(1) + self.ui.horizontalSlider_skew.setValue(0) + self.ui.horizontalSlider_vscale.setValue(10) + self.ui.label_selectedInfo.setText("") + self.ui.checkBox_removeLeadingZeros.setChecked(False) + self.ui.checkBox_rescalePatch.setChecked(False) + self.ui.checkBox_normWHRatio.setChecked(False) + self.ui.checkBox_invertPatch.setChecked(False) + self.ui.checkBox_ordinalIndicator.setChecked(False) + self.ui.comboBox_binarizationMethod.setCurrentIndex(0) + self.ui.checkBox_templatefield.setChecked(False) + self.ui.lineEdit_templatefield.setText("") + else: + item_obj.settings = normalize_settings_dict( + item_obj.settings, default_info_for_box_name(item_obj.name) + ) + self.ui.label_selectedInfo.setText(f"{item_obj.name}") + self.ui.lineEdit_format.setText(item_obj.settings["format_regex"]) + self.ui.comboBox_fieldType.setCurrentIndex(item_obj.settings["type"]) + self.ui.checkBox_smoothing.setChecked(item_obj.settings["smoothing"]) + self.ui.checkBox_skip_empty.setChecked(item_obj.settings["skip_empty"]) + self.ui.horizontalSlider_conf_thresh.setValue( + int(item_obj.settings["conf_thresh"] * 100) + ) + self.ui.checkBox_autocrop.setChecked(item_obj.settings["autocrop"]) + self.ui.checkBox_skip_similar_image.setChecked( + item_obj.settings["skip_similar_image"] + ) + self.ui.horizontalSlider_cleanup.setValue( + int(item_obj.settings["cleanup_thresh"] * 100) + ) + self.ui.horizontalSlider_dilate.setValue(item_obj.settings["dilate"]) + self.ui.horizontalSlider_skew.setValue(item_obj.settings["skew"]) + self.ui.horizontalSlider_vscale.setValue(item_obj.settings["vscale"]) + self.ui.checkBox_removeLeadingZeros.setChecked( + item_obj.settings["remove_leading_zeros"] + ) + self.ui.checkBox_rescalePatch.setChecked(item_obj.settings["rescale_patch"]) + self.ui.checkBox_normWHRatio.setChecked( + item_obj.settings["normalize_wh_ratio"] + ) + self.ui.checkBox_invertPatch.setChecked(item_obj.settings["invert_patch"]) + self.ui.checkBox_dotDetector.setChecked(item_obj.settings["dot_detector"]) + self.ui.checkBox_ordinalIndicator.setChecked( + item_obj.settings["ordinal_indicator"] + ) + self.ui.comboBox_binarizationMethod.setCurrentIndex( + item_obj.settings["binarization_method"] + ) + self.ui.checkBox_templatefield.setChecked( + item_obj.settings["templatefield"] + ) + self.ui.lineEdit_templatefield.setText( + item_obj.settings["templatefield_text"] + ) + + self.ui.comboBox_formatPrefix.setCurrentIndex(12) + + self.ui.lineEdit_format.blockSignals(False) + self.ui.comboBox_fieldType.blockSignals(False) + self.ui.checkBox_smoothing.blockSignals(False) + self.ui.checkBox_skip_empty.blockSignals(False) + self.ui.horizontalSlider_conf_thresh.blockSignals(False) + self.ui.checkBox_autocrop.blockSignals(False) + self.ui.checkBox_skip_similar_image.blockSignals(False) + self.ui.horizontalSlider_cleanup.blockSignals(False) + self.ui.horizontalSlider_dilate.blockSignals(False) + self.ui.horizontalSlider_skew.blockSignals(False) + self.ui.horizontalSlider_vscale.blockSignals(False) + self.ui.checkBox_removeLeadingZeros.blockSignals(False) + self.ui.checkBox_rescalePatch.blockSignals(False) + self.ui.checkBox_normWHRatio.blockSignals(False) + self.ui.checkBox_invertPatch.blockSignals(False) + self.ui.checkBox_ordinalIndicator.blockSignals(False) + self.ui.comboBox_binarizationMethod.blockSignals(False) + self.ui.comboBox_formatPrefix.blockSignals(False) + self.ui.checkBox_templatefield.blockSignals(False) + self.ui.lineEdit_templatefield.blockSignals(False) + + def listItemClicked(self, item): + user_role = item.data(Qt.ItemDataRole.UserRole) + if user_role in ["checked", "templatefield"] and item.column() == 0: + # enable the remove box button and disable the make box button + self.ui.pushButton_makeBox.setEnabled(False) + self.ui.pushButton_removeBox.setEnabled(user_role == "checked") + self.ui.groupBox_target_settings.setEnabled(user_role == "checked") + self.populateSettings(item.text()) + else: + # enable the make box button and disable the remove box button + self.ui.pushButton_removeBox.setEnabled(False) + self.ui.pushButton_makeBox.setEnabled(item.column() == 0) + self.ui.groupBox_target_settings.setEnabled(False) + self.populateSettings("") + + if item.column() == 0: + # if this is not a default box - enable the template field checkbox + if item.text() not in [box["name"] for box in default_boxes]: + self.ui.checkBox_templatefield.setEnabled(True) + self.ui.lineEdit_templatefield.setEnabled(True) + else: + self.ui.checkBox_templatefield.setEnabled(False) + self.ui.lineEdit_templatefield.setEnabled(False) + + def openOBSConnectModal(self): + # disable OBS options + self.ui.lineEdit_sceneName.setEnabled(False) + self.ui.checkBox_recreate.setEnabled(False) + self.ui.pushButton_createOBSScene.setEnabled(False) + + # load the ui from "connect_obs.ui" + self.obs_modal_ui = Ui_ConnectObs() + self.obs_connect_modal = QDialog() + self.obs_modal_ui.setupUi(self.obs_connect_modal) + self.obs_connect_modal.setWindowTitle("Connect to OBS") + + # connect the "connect" button to a function + self.obs_modal_ui.pushButton_connect.clicked.connect(self.connectObs) + + # load the saved data from scoresight.json + obs_data = fetch_data("scoresight.json", "obs") + if obs_data: + self.obs_modal_ui.lineEdit_ip.setText(obs_data["ip"]) + self.obs_modal_ui.lineEdit_port.setText(obs_data["port"]) + self.obs_modal_ui.lineEdit_password.setText(obs_data["password"]) + # show the modal + self.obs_connect_modal.show() + # focus the connect button + self.obs_modal_ui.pushButton_connect.setFocus() + + def connectObs(self): + # open a websocket connection to OBS using obs_websocket.py + # enable the save button in the modal if the connection is successful + if self.obs_connect_modal is not None: + self.obs_websocket_client = open_obs_websocket( + { + "ip": self.obs_modal_ui.lineEdit_ip.text(), + "port": self.obs_modal_ui.lineEdit_port.text(), + "password": self.obs_modal_ui.lineEdit_password.text(), + } + ) + else: + self.obs_websocket_client = open_obs_websocket( + fetch_data("scoresight.json", "obs") + ) + if not self.obs_websocket_client: + # show error in label_error + if self.obs_connect_modal: + self.obs_modal_ui.label_error.setText("Cannot connect to OBS") + return + + # connection was successful + if self.obs_connect_modal: + store_data( + "scoresight.json", + "obs", + { + "ip": self.obs_modal_ui.lineEdit_ip.text(), + "port": self.obs_modal_ui.lineEdit_port.text(), + "password": self.obs_modal_ui.lineEdit_password.text(), + }, + ) + self.obs_connect_modal.close() + + self.ui.lineEdit_sceneName.setEnabled(True) + self.ui.checkBox_recreate.setEnabled(True) + self.ui.pushButton_createOBSScene.setEnabled(True) + + logger.info("OBS: Connected") + + @Slot() + def getSources(self): + self.update_sources.emit(get_camera_info()) + + @Slot(list) + def updateSources(self, camera_sources: list[CameraInfo]): + self.reset_playing_source() + # clear all the items after "Screen Capture" + for i in range(4, self.ui.comboBox_camera_source.count()): + self.ui.comboBox_camera_source.removeItem(4) + + # populate the combobox with the sources + for source in camera_sources: + self.ui.comboBox_camera_source.addItem(source.description, source) + + self.ui.comboBox_camera_source.setEnabled(True) + currentIndexChangedSignal = QMetaMethod.fromSignal( + self.ui.comboBox_camera_source.currentIndexChanged + ) + if self.ui.comboBox_camera_source.isSignalConnected(currentIndexChangedSignal): + self.ui.comboBox_camera_source.currentIndexChanged.disconnect() + self.ui.comboBox_camera_source.currentIndexChanged.connect(self.sourceChanged) + + # enable the source view frame + self.ui.frame_source_view.setEnabled(True) + + selected_source_from_storage = fetch_data("scoresight.json", "source_selected") + if type(selected_source_from_storage) == str: + logger.info( + "Source selected from storage: %s", selected_source_from_storage + ) + # check if the source is a file path + if path.exists(selected_source_from_storage): + logger.info("File exists: %s", selected_source_from_storage) + self.ui.comboBox_camera_source.blockSignals(True) + self.ui.comboBox_camera_source.setCurrentIndex(1) + self.ui.comboBox_camera_source.blockSignals(False) + self.source_name = selected_source_from_storage + self.sourceSelectionSucessful() + else: + # select the last selected source + self.ui.comboBox_camera_source.setCurrentText( + selected_source_from_storage + ) + + def reset_playing_source(self): + if self.image_viewer: + # remove the image viewer from the layout frame_for_source_view_label + self.ui.frame_for_source_view_label.layout().removeWidget(self.image_viewer) + self.image_viewer.close() + self.image_viewer = None + # add a label with markdown text + label_select_source = QLabel("### Open a Camera or Load a File") + label_select_source.setTextFormat(Qt.TextFormat.MarkdownText) + label_select_source.setEnabled(False) + label_select_source.setAlignment(Qt.AlignmentFlag.AlignCenter) + clear_layout(self.ui.frame_for_source_view_label.layout()) + self.ui.frame_for_source_view_label.layout().addWidget(label_select_source) + + def sourceChanged(self, index): + # get the source name from the combobox + self.source_name = None + self.ui.groupBox_sb_info.setEnabled(False) + self.ui.tableWidget_boxes.setEnabled(False) + self.ui.pushButton_fourCorner.setEnabled(False) + self.ui.pushButton_binary.setEnabled(False) + if self.ui.comboBox_camera_source.currentIndex() == 0: + self.reset_playing_source() + return + elif self.ui.comboBox_camera_source.currentIndex() == 1: + logger.info("Open File selection dialog") + # open a file dialog to select a video file + file, _ = QFileDialog.getOpenFileName( + self, "Open Video File", "", "Video Files (*.mp4 *.avi *.mov)" + ) + if not file or not path.exists(file): + # no file selected - change source to "Select a source" + logger.error("No file selected") + self.ui.comboBox_camera_source.setCurrentText("Select a source") + return + else: + logger.info("File selected: %s", file) + self.source_name = file + elif self.ui.comboBox_camera_source.currentIndex() == 2: + # open a dialog to enter the url + url_dialog = QDialog() + ui_urlsource = Ui_UrlSource() + ui_urlsource.setupUi(url_dialog) + url_dialog.setWindowTitle("URL Source") + # focus on url input + ui_urlsource.lineEdit_url.setFocus() + url_dialog.exec() # wait for the dialog to close + # check if the dialog was accepted + if url_dialog.result() != QDialog.DialogCode.Accepted: + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + self.source_name = ui_urlsource.lineEdit_url.text() + if self.source_name == "": + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + elif self.ui.comboBox_camera_source.currentIndex() == 3: + # open a dialog to select the screen + screen_dialog = QDialog() + ui_screencapture = Ui_ScreenCapture() + ui_screencapture.setupUi(screen_dialog) + # set width and height of the dialog + screen_dialog.setFixedWidth(400) + + screen_dialog.setWindowTitle( + QCoreApplication.translate( + "MainWindow", "Screen Capture Selection", None + ) + ) + # populate comboBox_window with the available windows + ui_screencapture.comboBox_window.clear() + ui_screencapture.comboBox_window.addItem( + QCoreApplication.translate( + "MainWindow", "Capture the entire screen", None + ), + -1, + ) + for window in ScreenCapture.list_windows(): + ui_screencapture.comboBox_window.addItem(window[0], window[1]) + screen_dialog.exec() + # check if the dialog was accepted + if screen_dialog.result() != QDialog.DialogCode.Accepted: + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + # get the window ID from the comboBox_window + window_id = ui_screencapture.comboBox_window.currentData() + self.source_name = window_id + + # store the source selection in scoresight.json + self.globalSettingsChanged("source_selected", self.source_name) + self.sourceSelectionSucessful() + + def itemSelected(self, item_name): + # select the item in the tableWidget_boxes + items = self.ui.tableWidget_boxes.findItems( + item_name, Qt.MatchFlag.MatchExactly + ) + if len(items) == 0: + return + item = items[0] + item.setSelected(True) + self.ui.tableWidget_boxes.setCurrentItem(item) + self.listItemClicked(item) + + def fourCornersApplied(self, corners): + # check the button + self.ui.pushButton_fourCorner.setChecked(True) + + def sourceSelectionSucessful(self): + if self.ui.comboBox_camera_source.currentIndex() == 0: + return + + self.ui.frame_source_view.setEnabled(False) + + if self.ui.comboBox_camera_source.currentIndex() == 1: + if self.source_name is None or not path.exists(self.source_name): + logger.error("No file selected") + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + logger.info("Loading file selected: %s", self.source_name) + camera_info = CameraInfo( + self.source_name, + self.source_name, + self.source_name, + CameraInfo.CameraType.FILE, + ) + elif self.ui.comboBox_camera_source.currentIndex() == 2: + if self.source_name is None or self.source_name == "": + logger.error("No url entered") + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + logger.info("Loading url: %s", self.source_name) + camera_info = CameraInfo( + self.source_name, + self.source_name, + self.source_name, + CameraInfo.CameraType.URL, + ) + elif self.ui.comboBox_camera_source.currentIndex() == 3: + if self.source_name is None: + logger.error("No screen capture selected") + self.ui.comboBox_camera_source.setCurrentIndex(0) + return + logger.info("Loading screen capture: %s", self.source_name) + camera_info = CameraInfo( + self.source_name, + self.source_name, + self.source_name, + CameraInfo.CameraType.SCREEN_CAPTURE, + ) + else: + if self.ui.comboBox_camera_source.currentData() is None: + return + + logger.info( + "Loading camera: %s", self.ui.comboBox_camera_source.currentText() + ) + camera_info = self.ui.comboBox_camera_source.currentData() + + if self.image_viewer: + # remove the image viewer from the layout frame_for_source_view_label + self.ui.frame_for_source_view_label.layout().removeWidget(self.image_viewer) + self.image_viewer.close() + self.image_viewer = None + + # clear self.ui.frame_for_source_view_label + clear_layout(self.ui.frame_for_source_view_label.layout()) + + # set the pixmap to the image viewer + self.image_viewer = ImageViewer( + camera_info, + self.fourCornersApplied, + self.detectionTargetsStorage, + self.itemSelected, + ) + self.ui.toolButton_videoSettings.setEnabled( + camera_info.type == CameraInfo.CameraType.OPENCV + ) + self.ui.pushButton_fourCorner.setEnabled(True) + self.ui.pushButton_binary.setEnabled(True) + self.ui.pushButton_fourCorner.toggled.connect( + self.image_viewer.toggleFourCorner + ) + self.ui.pushButton_binary.clicked.connect(self.image_viewer.toggleBinary) + if self.image_viewer.timerThread: + self.image_viewer.timerThread.ocr_result_signal.connect(self.ocrResult) + self.image_viewer.timerThread.update_error.connect(self.updateError) + self.image_viewer.first_frame_received_signal.connect( + self.cameraConnectedEnableUI + ) + self.ocrModelChanged(fetch_data("scoresight.json", "ocr_model", 1)) + + # set the image viewer to the layout frame_for_source_view_label + self.ui.frame_for_source_view_label.layout().addWidget(self.image_viewer) + + def cameraConnectedEnableUI(self): + # enable groupBox_sb_info + self.ui.groupBox_sb_info.setEnabled(True) + self.ui.tableWidget_boxes.setEnabled(True) + self.ui.frame_source_view.setEnabled(True) + self.ui.widget_viewTools.setEnabled(True) + + # load the boxes from scoresight.json + self.detectionTargetsStorage.loadBoxesFromStorage() + self.updateError(None) + + def updateError(self, error): + if not error: + return + logger.error(error) + self.ui.frame_source_view.setEnabled(True) + self.ui.widget_viewTools.setEnabled(False) + + def ocrResult(self, results: list[TextDetectionTargetWithResult]): + # update template fields + for targetWithResult in results: + if not targetWithResult.settings["templatefield"]: + continue + targetWithResult.result = evaluate_template_field(results, targetWithResult) + targetWithResult.result_state = ( + TextDetectionTargetWithResult.ResultState.Success + if targetWithResult.result is not None + else TextDetectionTargetWithResult.ResultState.Empty + ) + + # update the table widget value items + for targetWithResult in results: + if ( + targetWithResult.result_state + == TextDetectionTargetWithResult.ResultState.Success + ): + items = self.ui.tableWidget_boxes.findItems( + targetWithResult.name, Qt.MatchFlag.MatchExactly + ) + if len(items) == 0: + continue + item = items[0] + # get the value (1 column) of the item + item = self.ui.tableWidget_boxes.item(item.row(), 1) + item.setText(targetWithResult.result) + + if not self.updateOCRResults: + # don't update the results, the user has disabled updates + return + + update_http_server(results) + + # update vmix + self.vmixUpdater.update_vmix(results) + + if self.out_folder is None: + return + + if not path.exists(path.abspath(self.out_folder)): + self.out_folder = None + remove_data("scoresight.json", "output_folder") + logger.warning("Output folder does not exist") + return + + # check if enough time has passed since last file save according to aggs per second + if ( + datetime.datetime.now() - self.last_aggregate_save + ).total_seconds() < 1.0 / self.ui.horizontalSlider_aggsPerSecond.value(): + return + + self.last_aggregate_save = datetime.datetime.now() + + # update the obs scene sources with the results, use update_text_source + for targetWithResult in results: + if targetWithResult.result is None: + continue + if ( + targetWithResult.settings is not None + and "skip_empty" in targetWithResult.settings + and targetWithResult.settings["skip_empty"] + and len(targetWithResult.result) == 0 + ): + continue + if ( + targetWithResult.result_state + != TextDetectionTargetWithResult.ResultState.Success + ): + continue + + if ( + self.obs_websocket_client is not None + and targetWithResult.settings is not None + ): + # find the source name for the target from the default boxes + update_text_source( + self.obs_websocket_client, + targetWithResult.settings["obs_source_name"], + targetWithResult.result, + ) + + # save the results to text files + file_output.save_text_files( + results, self.out_folder, self.ui.comboBox_appendMethod.currentIndex() + ) + + # save the results to a csv file + if self.ui.checkBox_saveCsv.isChecked(): + file_output.save_csv( + results, + self.out_folder, + self.ui.comboBox_appendMethod.currentIndex(), + self.first_csv_append, + ) + + # save the results to an xml file + if self.ui.checkBox_saveXML.isChecked(): + file_output.save_xml( + results, + self.out_folder, + ) + + def addBox(self): + # add a new box to the tableWidget_boxes + # find the number of custom boxes + custom_boxes = [] + for i in range(self.ui.tableWidget_boxes.rowCount()): + item = self.ui.tableWidget_boxes.item(i, 0) + if item.text() not in [o["name"] for o in default_boxes]: + custom_boxes.append(item.text()) + + i = len(custom_boxes) + new_box_name = f"Custom {i + 1}" + custom_boxes_names = fetch_data("scoresight.json", "custom_boxes_names", []) + # find if the name already exists + while new_box_name in custom_boxes or new_box_name in custom_boxes_names: + i += 1 + new_box_name = f"Custom {i + 1}" + + store_custom_box_name(new_box_name) + item = QTableWidgetItem( + QIcon( + path.abspath(path.join(path.dirname(__file__), "icons/circle-x.svg")) + ), + new_box_name, + ) + item.setData(Qt.ItemDataRole.UserRole, "unchecked") + self.ui.tableWidget_boxes.insertRow(self.ui.tableWidget_boxes.rowCount()) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 0, item + ) + disabledItem = QTableWidgetItem() + disabledItem.setFlags(Qt.ItemFlag.NoItemFlags) + self.ui.tableWidget_boxes.setItem( + self.ui.tableWidget_boxes.rowCount() - 1, 1, disabledItem + ) + + def removeCustomBox(self): + item = self.ui.tableWidget_boxes.currentItem() + if not item: + logger.info("No item selected") + return + if item.column() != 0: + item = self.ui.tableWidget_boxes.item(item.row(), 0) + self.removeBox() + remove_custom_box_name_in_storage(item.text()) + # only allow removing custom boxes + if item.text() in [o["name"] for o in default_boxes]: + logger.info("Cannot remove default box") + return + # remove the selected item from the tableWidget_boxes + self.ui.tableWidget_boxes.removeRow(item.row()) + + def editBoxName(self, item): + if item.text() in [o["name"] for o in default_boxes]: + # dont allow editing default boxes + return + new_name, ok = QInputDialog.getText( + self, "Edit Box Name", "New Name:", text=item.text() + ) + old_name = item.text() + if ok and new_name != "" and new_name != old_name: + # check if name doesn't exist already + for i in range(self.ui.tableWidget_boxes.rowCount()): + if new_name == self.ui.tableWidget_boxes.item(i, 0).text(): + logger.info("Name '%s' already exists", new_name) + return + # rename the item in the tableWidget_boxes + rename_custom_box_name_in_storage(old_name, new_name) + item.setText(new_name) + # rename the item in the detectionTargetsStorage + if not self.detectionTargetsStorage.rename_item(old_name, new_name): + logger.info("Error renaming item in application storage") + return + else: + # check if the item role isn't "templatefield" + if item.data(Qt.ItemDataRole.UserRole) != "templatefield": + # remove the item from the tableWidget_boxes + self.ui.tableWidget_boxes.removeRow(item.row()) + + def makeBox(self): + item = self.ui.tableWidget_boxes.currentItem() + if not item: + return + # create a new box on self.image_viewer with the name of the selected item from the tableWidget_boxes + # change the list icon to green checkmark + item.setIcon( + QIcon( + path.abspath( + path.join(path.dirname(__file__), "icons/circle-check.svg") + ) + ) + ) + item.setData(Qt.ItemDataRole.UserRole, "checked") + self.listItemClicked(item) + + # get the size of the box from the name + info = default_info_for_box_name(item.text()) + + self.detectionTargetsStorage.add_item( + TextDetectionTarget( + info["x"], + info["y"], + info["width"], + info["height"], + item.text(), + normalize_settings_dict({}, info), + ) + ) + + def removeBox(self): + item = self.ui.tableWidget_boxes.currentItem() + if not item: + return + # change the list icon to red x + item.setIcon( + QIcon(path.abspath(path.join(path.dirname(__file__), "icons/circle-x.svg"))) + ) + item.setData(Qt.ItemDataRole.UserRole, "unchecked") + self.listItemClicked(item) + self.detectionTargetsStorage.remove_item(item.text()) + + def makeTemplateField(self, toggled: bool): + item = self.ui.tableWidget_boxes.currentItem() + if not item: + return + + if not toggled: + self.removeBox() + return + + # create a new box on self.image_viewer with the name of the selected item from the tableWidget_boxes + # change the list icon to green checkmark + item.setIcon( + QIcon( + path.abspath( + path.join(path.dirname(__file__), "icons/template-field.svg") + ) + ) + ) + item.setData(Qt.ItemDataRole.UserRole, "templatefield") + + self.detectionTargetsStorage.add_item( + TextDetectionTarget( + 0, + 0, + 0, + 0, + item.text(), + normalize_settings_dict({"templatefield": True}, None), + ) + ) + + self.listItemClicked(item) + + def createOBSScene(self): + # get the scene name from the lineEdit_sceneName + scene_name = self.ui.lineEdit_sceneName.text() + # clear or create a new scene + create_obs_scene_from_export(self.obs_websocket_client, scene_name) + + # on destroy, close the OBS connection + def closeEvent(self, event): + logger.info("Closing") + if self.image_viewer: + self.image_viewer.close() + self.image_viewer = None + + if self.log_dialog: + self.log_dialog.close() + self.log_dialog = None + + # store the boxes to scoresight.json + self.detectionTargetsStorage.saveBoxesToStorage() + if self.obs_websocket_client: + # destroy the client object + self.obs_websocket_client = None + + super().closeEvent(event) diff --git a/requirements.txt b/requirements.txt index f5918b5..acc563a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,10 +2,10 @@ cyndilib fastapi numpy==1.26.4 obsws-python -opencv-python==4.9.0.80 +opencv-python==4.10.0.84 pillow platformdirs -pyinstaller==6.8.0 +pyinstaller==6.10.0 pyside6 python-dotenv requests diff --git a/scoresight.spec b/scoresight.spec index 0ff8ff9..f8ec012 100644 --- a/scoresight.spec +++ b/scoresight.spec @@ -50,6 +50,7 @@ datas = [ sources = [ 'camera_info.py', 'camera_view.py', + 'camera_thread.py', 'defaults.py', 'file_output.py', 'frame_stabilizer.py', @@ -57,6 +58,7 @@ sources = [ 'http_server.py', 'log_view.py', 'main.py', + 'mainwindow.py', 'ndi.py', 'obs_websocket.py', 'resizable_rect.py', diff --git a/screen_capture_source.py b/screen_capture_source.py index f6c7142..8f801d9 100644 --- a/screen_capture_source.py +++ b/screen_capture_source.py @@ -1,4 +1,5 @@ import platform +from typing import Union, Type class ScreenCaptureNotImplemented: @@ -19,6 +20,25 @@ def read(self): return False, None +# Define a base class or protocol for screen capture +class ScreenCaptureBase: + @staticmethod + def list_windows(): + raise NotImplementedError + + def __init__(self, window_name): + raise NotImplementedError + + def isOpened(self): + raise NotImplementedError + + def release(self): + raise NotImplementedError + + def read(self): + raise NotImplementedError + + # This is a simple example of how to use the screen capture source in the # platform-independent part of the code. if platform.system() == "Darwin": @@ -27,3 +47,5 @@ def read(self): from screen_capture_source_windows import ScreenCaptureWindows as ScreenCapture else: ScreenCapture = ScreenCaptureNotImplemented + +ScreenCaptureType = Union[Type[ScreenCapture], Type[ScreenCaptureBase]] diff --git a/storage.py b/storage.py index 0c1e042..e6b34ec 100644 --- a/storage.py +++ b/storage.py @@ -11,7 +11,7 @@ data_subscribers = {} -def subscribe_to_data(file_path, document_name, callback): +def subscribe_to_data(file_path: str, document_name: str, callback: callable): # Subscribe to data changes in a JSON file # prepend the user data directory file_path = os.path.join(user_data_dir("scoresight"), file_path) diff --git a/ui_about.py b/ui_about.py index 5c8eac5..c37899a 100644 --- a/ui_about.py +++ b/ui_about.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'about.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_connect_obs.py b/ui_connect_obs.py index ab04702..d98e064 100644 --- a/ui_connect_obs.py +++ b/ui_connect_obs.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'connect_obs.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_log_view.py b/ui_log_view.py index 028d293..6427dfa 100644 --- a/ui_log_view.py +++ b/ui_log_view.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'log_view.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_mainwindow.py b/ui_mainwindow.py index 13231fc..77db006 100644 --- a/ui_mainwindow.py +++ b/ui_mainwindow.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'mainwindow.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_screen_capture.py b/ui_screen_capture.py index cb3513a..70bd188 100644 --- a/ui_screen_capture.py +++ b/ui_screen_capture.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'screen_capture.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_update_available.py b/ui_update_available.py index caf9763..368a569 100644 --- a/ui_update_available.py +++ b/ui_update_available.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'update_available.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_url_source.py b/ui_url_source.py index e3482b2..f6fa3ac 100644 --- a/ui_url_source.py +++ b/ui_url_source.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'url_source.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ diff --git a/ui_video_settings.py b/ui_video_settings.py index ad28b3b..8f4aa3a 100644 --- a/ui_video_settings.py +++ b/ui_video_settings.py @@ -3,7 +3,7 @@ ################################################################################ ## Form generated from reading UI file 'video_settings.ui' ## -## Created by: Qt User Interface Compiler version 6.7.1 +## Created by: Qt User Interface Compiler version 6.7.2 ## ## WARNING! All changes made in this file will be lost when recompiling UI file! ################################################################################ @@ -17,22 +17,15 @@ QPalette, QPixmap, QRadialGradient, QTransform) from PySide6.QtWidgets import (QAbstractButton, QApplication, QComboBox, QDialog, QDialogButtonBox, QFormLayout, QGridLayout, QLabel, - QSizePolicy, QSpinBox, QWidget) + QPlainTextEdit, QSizePolicy, QSpinBox, QWidget) class Ui_Dialog(object): def setupUi(self, Dialog): if not Dialog.objectName(): Dialog.setObjectName(u"Dialog") - Dialog.resize(400, 153) + Dialog.resize(400, 259) self.gridLayout = QGridLayout(Dialog) self.gridLayout.setObjectName(u"gridLayout") - self.buttonBox = QDialogButtonBox(Dialog) - self.buttonBox.setObjectName(u"buttonBox") - self.buttonBox.setOrientation(Qt.Horizontal) - self.buttonBox.setStandardButtons(QDialogButtonBox.Cancel|QDialogButtonBox.Ok) - - self.gridLayout.addWidget(self.buttonBox, 2, 0, 1, 1) - self.widget = QWidget(Dialog) self.widget.setObjectName(u"widget") self.formLayout = QFormLayout(self.widget) @@ -40,18 +33,18 @@ def setupUi(self, Dialog): self.label = QLabel(self.widget) self.label.setObjectName(u"label") - self.formLayout.setWidget(1, QFormLayout.LabelRole, self.label) + self.formLayout.setWidget(2, QFormLayout.LabelRole, self.label) self.comboBox_fourcc = QComboBox(self.widget) self.comboBox_fourcc.addItem("") self.comboBox_fourcc.setObjectName(u"comboBox_fourcc") - self.formLayout.setWidget(1, QFormLayout.FieldRole, self.comboBox_fourcc) + self.formLayout.setWidget(2, QFormLayout.FieldRole, self.comboBox_fourcc) self.label_2 = QLabel(self.widget) self.label_2.setObjectName(u"label_2") - self.formLayout.setWidget(2, QFormLayout.LabelRole, self.label_2) + self.formLayout.setWidget(3, QFormLayout.LabelRole, self.label_2) self.spinBox_fps = QSpinBox(self.widget) self.spinBox_fps.setObjectName(u"spinBox_fps") @@ -59,12 +52,12 @@ def setupUi(self, Dialog): self.spinBox_fps.setMaximum(60) self.spinBox_fps.setValue(30) - self.formLayout.setWidget(2, QFormLayout.FieldRole, self.spinBox_fps) + self.formLayout.setWidget(3, QFormLayout.FieldRole, self.spinBox_fps) self.label_3 = QLabel(self.widget) self.label_3.setObjectName(u"label_3") - self.formLayout.setWidget(3, QFormLayout.LabelRole, self.label_3) + self.formLayout.setWidget(4, QFormLayout.LabelRole, self.label_3) self.comboBox_resolution = QComboBox(self.widget) self.comboBox_resolution.addItem("") @@ -81,11 +74,34 @@ def setupUi(self, Dialog): self.comboBox_resolution.addItem("") self.comboBox_resolution.setObjectName(u"comboBox_resolution") - self.formLayout.setWidget(3, QFormLayout.FieldRole, self.comboBox_resolution) + self.formLayout.setWidget(4, QFormLayout.FieldRole, self.comboBox_resolution) + + self.plainTextEdit_videoProps = QPlainTextEdit(self.widget) + self.plainTextEdit_videoProps.setObjectName(u"plainTextEdit_videoProps") + + self.formLayout.setWidget(5, QFormLayout.FieldRole, self.plainTextEdit_videoProps) + + self.label_4 = QLabel(self.widget) + self.label_4.setObjectName(u"label_4") + + self.formLayout.setWidget(1, QFormLayout.LabelRole, self.label_4) + + self.comboBox_captureBackend = QComboBox(self.widget) + self.comboBox_captureBackend.addItem("") + self.comboBox_captureBackend.setObjectName(u"comboBox_captureBackend") + + self.formLayout.setWidget(1, QFormLayout.FieldRole, self.comboBox_captureBackend) self.gridLayout.addWidget(self.widget, 1, 0, 1, 1) + self.buttonBox = QDialogButtonBox(Dialog) + self.buttonBox.setObjectName(u"buttonBox") + self.buttonBox.setOrientation(Qt.Horizontal) + self.buttonBox.setStandardButtons(QDialogButtonBox.Cancel|QDialogButtonBox.Ok) + + self.gridLayout.addWidget(self.buttonBox, 3, 0, 1, 1) + self.retranslateUi(Dialog) self.buttonBox.accepted.connect(Dialog.accept) @@ -114,5 +130,8 @@ def retranslateUi(self, Dialog): self.comboBox_resolution.setItemText(10, QCoreApplication.translate("Dialog", u"2048x1080", None)) self.comboBox_resolution.setItemText(11, QCoreApplication.translate("Dialog", u"4096x2160", None)) + self.label_4.setText(QCoreApplication.translate("Dialog", u"Capture Backend", None)) + self.comboBox_captureBackend.setItemText(0, QCoreApplication.translate("Dialog", u"Default / Any", None)) + # retranslateUi diff --git a/video_settings.py b/video_settings.py index d0b82ba..b8aecf8 100644 --- a/video_settings.py +++ b/video_settings.py @@ -1,6 +1,8 @@ import cv2 +import platform from PySide6.QtWidgets import QDialog +from camera_thread import OpenCVVideoCaptureWithSettings from sc_logging import logger from ui_video_settings import Ui_Dialog as Ui_VideoSettingsDialog from storage import store_data @@ -55,16 +57,125 @@ def __init__(self): super().__init__() self.ui = Ui_VideoSettingsDialog() self.ui.setupUi(self) - - def init_ui(self, cap): - self.cap = cap self.ui.buttonBox.accepted.connect(self.save_settings) self.ui.buttonBox.rejected.connect(self.close) - self.ui.spinBox_fps.setValue(int(cap.get(cv2.CAP_PROP_FPS))) + + def init_ui(self, cap: OpenCVVideoCaptureWithSettings): + self.cap = cap + self.ui.spinBox_fps.setValue(int(self.cap.get(cv2.CAP_PROP_FPS))) self.ui.comboBox_resolution.setCurrentText( - f"{cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}" + f"{self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}" ) self.populate_supported_fourcc() + self.populate_video_props() + self.populate_backend_combo() + + def populate_backend_combo(self): + self.ui.comboBox_captureBackend.clear() + self.ui.comboBox_captureBackend.addItem("Auto", cv2.CAP_ANY) + if platform.system() == "Windows": + self.ui.comboBox_captureBackend.addItem("DShow", cv2.CAP_DSHOW) + self.ui.comboBox_captureBackend.addItem( + "Microsoft Media Foundation", cv2.CAP_MSMF + ) + self.ui.comboBox_captureBackend.addItem( + "Microsoft Windows Runtime", cv2.CAP_WINRT + ) + elif platform.system() == "Linux": + self.ui.comboBox_captureBackend.addItem("V4L", cv2.CAP_V4L) + self.ui.comboBox_captureBackend.addItem("GStreamer", cv2.CAP_GSTREAMER) + elif platform.system() == "Darwin": + self.ui.comboBox_captureBackend.addItem( + "AVFoundation", cv2.CAP_AVFOUNDATION + ) + self.ui.comboBox_captureBackend.setCurrentIndex( + self.ui.comboBox_captureBackend.findData(self.cap.get(cv2.CAP_PROP_BACKEND)) + ) + + def populate_video_props(self): + # scan all opencv properties and write to plainTextEdit_videoProps + video_props = [ + (cv2.CAP_PROP_POS_MSEC, "CAP_PROP_POS_MSEC"), + (cv2.CAP_PROP_POS_FRAMES, "CAP_PROP_POS_FRAMES"), + (cv2.CAP_PROP_POS_AVI_RATIO, "CAP_PROP_POS_AVI_RATIO"), + (cv2.CAP_PROP_FRAME_WIDTH, "CAP_PROP_FRAME_WIDTH"), + (cv2.CAP_PROP_FRAME_HEIGHT, "CAP_PROP_FRAME_HEIGHT"), + (cv2.CAP_PROP_FPS, "CAP_PROP_FPS"), + (cv2.CAP_PROP_FOURCC, "CAP_PROP_FOURCC"), + (cv2.CAP_PROP_FRAME_COUNT, "CAP_PROP_FRAME_COUNT"), + (cv2.CAP_PROP_FORMAT, "CAP_PROP_FORMAT"), + (cv2.CAP_PROP_MODE, "CAP_PROP_MODE"), + (cv2.CAP_PROP_BRIGHTNESS, "CAP_PROP_BRIGHTNESS"), + (cv2.CAP_PROP_CONTRAST, "CAP_PROP_CONTRAST"), + (cv2.CAP_PROP_SATURATION, "CAP_PROP_SATURATION"), + (cv2.CAP_PROP_HUE, "CAP_PROP_HUE"), + (cv2.CAP_PROP_GAIN, "CAP_PROP_GAIN"), + (cv2.CAP_PROP_EXPOSURE, "CAP_PROP_EXPOSURE"), + (cv2.CAP_PROP_CONVERT_RGB, "CAP_PROP_CONVERT_RGB"), + (cv2.CAP_PROP_WHITE_BALANCE_BLUE_U, "CAP_PROP_WHITE_BALANCE_BLUE_U"), + (cv2.CAP_PROP_RECTIFICATION, "CAP_PROP_RECTIFICATION"), + (cv2.CAP_PROP_MONOCHROME, "CAP_PROP_MONOCHROME"), + (cv2.CAP_PROP_SHARPNESS, "CAP_PROP_SHARPNESS"), + (cv2.CAP_PROP_AUTO_EXPOSURE, "CAP_PROP_AUTO_EXPOSURE"), + (cv2.CAP_PROP_GAMMA, "CAP_PROP_GAMMA"), + (cv2.CAP_PROP_TEMPERATURE, "CAP_PROP_TEMPERATURE"), + (cv2.CAP_PROP_TRIGGER, "CAP_PROP_TRIGGER"), + (cv2.CAP_PROP_TRIGGER_DELAY, "CAP_PROP_TRIGGER_DELAY"), + (cv2.CAP_PROP_WHITE_BALANCE_RED_V, "CAP_PROP_WHITE_BALANCE_RED_V"), + (cv2.CAP_PROP_ZOOM, "CAP_PROP_ZOOM"), + (cv2.CAP_PROP_FOCUS, "CAP_PROP_FOCUS"), + (cv2.CAP_PROP_GUID, "CAP_PROP_GUID"), + (cv2.CAP_PROP_ISO_SPEED, "CAP_PROP_ISO_SPEED"), + (cv2.CAP_PROP_BACKLIGHT, "CAP_PROP_BACKLIGHT"), + (cv2.CAP_PROP_PAN, "CAP_PROP_PAN"), + (cv2.CAP_PROP_TILT, "CAP_PROP_TILT"), + (cv2.CAP_PROP_ROLL, "CAP_PROP_ROLL"), + (cv2.CAP_PROP_IRIS, "CAP_PROP_IRIS"), + (cv2.CAP_PROP_SETTINGS, "CAP_PROP_SETTINGS"), + (cv2.CAP_PROP_BUFFERSIZE, "CAP_PROP_BUFFERSIZE"), + (cv2.CAP_PROP_AUTOFOCUS, "CAP_PROP_AUTOFOCUS"), + (cv2.CAP_PROP_SAR_NUM, "CAP_PROP_SAR_NUM"), + (cv2.CAP_PROP_SAR_DEN, "CAP_PROP_SAR_DEN"), + (cv2.CAP_PROP_BACKEND, "CAP_PROP_BACKEND"), + (cv2.CAP_PROP_CHANNEL, "CAP_PROP_CHANNEL"), + (cv2.CAP_PROP_AUTO_WB, "CAP_PROP_AUTO_WB"), + (cv2.CAP_PROP_WB_TEMPERATURE, "CAP_PROP_WB_TEMPERATURE"), + (cv2.CAP_PROP_CODEC_PIXEL_FORMAT, "CAP_PROP_CODEC_PIXEL_FORMAT"), + (cv2.CAP_PROP_BITRATE, "CAP_PROP_BITRATE"), + (cv2.CAP_PROP_ORIENTATION_META, "CAP_PROP_ORIENTATION_META"), + (cv2.CAP_PROP_ORIENTATION_AUTO, "CAP_PROP_ORIENTATION_AUTO"), + (cv2.CAP_PROP_HW_ACCELERATION, "CAP_PROP_HW_ACCELERATION"), + (cv2.CAP_PROP_HW_DEVICE, "CAP_PROP_HW_DEVICE"), + ( + cv2.CAP_PROP_HW_ACCELERATION_USE_OPENCL, + "CAP_PROP_HW_ACCELERATION_USE_OPENCL", + ), + (cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, "CAP_PROP_OPEN_TIMEOUT_MSEC"), + (cv2.CAP_PROP_READ_TIMEOUT_MSEC, "CAP_PROP_READ_TIMEOUT_MSEC"), + (cv2.CAP_PROP_STREAM_OPEN_TIME_USEC, "CAP_PROP_STREAM_OPEN_TIME_USEC"), + (cv2.CAP_PROP_VIDEO_TOTAL_CHANNELS, "CAP_PROP_VIDEO_TOTAL_CHANNELS"), + (cv2.CAP_PROP_VIDEO_STREAM, "CAP_PROP_VIDEO_STREAM"), + (cv2.CAP_PROP_AUDIO_STREAM, "CAP_PROP_AUDIO_STREAM"), + (cv2.CAP_PROP_AUDIO_POS, "CAP_PROP_AUDIO_POS"), + (cv2.CAP_PROP_AUDIO_SHIFT_NSEC, "CAP_PROP_AUDIO_SHIFT_NSEC"), + (cv2.CAP_PROP_AUDIO_DATA_DEPTH, "CAP_PROP_AUDIO_DATA_DEPTH"), + ( + cv2.CAP_PROP_AUDIO_SAMPLES_PER_SECOND, + "CAP_PROP_AUDIO_SAMPLES_PER_SECOND", + ), + (cv2.CAP_PROP_AUDIO_BASE_INDEX, "CAP_PROP_AUDIO_BASE_INDEX"), + (cv2.CAP_PROP_AUDIO_TOTAL_CHANNELS, "CAP_PROP_AUDIO_TOTAL_CHANNELS"), + (cv2.CAP_PROP_AUDIO_TOTAL_STREAMS, "CAP_PROP_AUDIO_TOTAL_STREAMS"), + (cv2.CAP_PROP_AUDIO_SYNCHRONIZE, "CAP_PROP_AUDIO_SYNCHRONIZE"), + (cv2.CAP_PROP_LRF_HAS_KEY_FRAME, "CAP_PROP_LRF_HAS_KEY_FRAME"), + (cv2.CAP_PROP_CODEC_EXTRADATA_INDEX, "CAP_PROP_CODEC_EXTRADATA_INDEX"), + (cv2.CAP_PROP_FRAME_TYPE, "CAP_PROP_FRAME_TYPE"), + (cv2.CAP_PROP_N_THREADS, "CAP_PROP_N_THREADS"), + ] + video_props_str = [ + f"{prop[1]}: {self.cap.get(prop[0])}" for prop in video_props + ] + self.ui.plainTextEdit_videoProps.setPlainText("\n".join(video_props_str)) def populate_supported_fourcc(self): supported_fourccs = get_supported_fourcc(self.cap) @@ -82,27 +193,12 @@ def save_settings(self): width = -1 height = -1 fourcc = self.ui.comboBox_fourcc.currentText() + store_data("scoresight.json", "fps", fps) + store_data("scoresight.json", "width", width) + store_data("scoresight.json", "height", height) + store_data("scoresight.json", "fourcc", fourcc) store_data( - "scoresight.json", - "video_settings", - { - "fps": fps, - "width": width, - "height": height, - "fourcc": fourcc, - }, - ) - self.cap.set(cv2.CAP_PROP_FPS, fps) - if width > 0: - self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) - if height > 0: - self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) - fourcc_int = int.from_bytes(fourcc.encode(), "little") - self.cap.set(cv2.CAP_PROP_FOURCC, fourcc_int) - print(f"FPS: {fps}, Resolution: {width}x{height}, FourCC: {fourcc}") - logger.info(f"FPS: {fps}, Resolution: {width}x{height}, FourCC: {fourcc}") - print( - f"FPS: {self.cap.get(cv2.CAP_PROP_FPS)}, Resolution: {self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)}x{self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)}, FourCC: {fourcc_to_str(self.cap.get(cv2.CAP_PROP_FOURCC))}" + "scoresight.json", "backend", self.ui.comboBox_captureBackend.currentData() ) self.close() diff --git a/video_settings.ui b/video_settings.ui index e5a3213..cc97d9e 100644 --- a/video_settings.ui +++ b/video_settings.ui @@ -7,34 +7,24 @@ 0 0 400 - 153 + 259 Dialog - - - - Qt::Horizontal - - - QDialogButtonBox::Cancel|QDialogButtonBox::Ok - - - - + Video Format - + @@ -43,14 +33,14 @@ - + FPS - + 1 @@ -63,14 +53,14 @@ - + Resolution - + @@ -134,9 +124,38 @@ + + + + + + + Capture Backend + + + + + + + + Default / Any + + + + + + + + Qt::Horizontal + + + QDialogButtonBox::Cancel|QDialogButtonBox::Ok + + + From f4de365ab82033c0a2c9c29b23fab8229b7c4224 Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Sat, 21 Sep 2024 22:13:54 -0400 Subject: [PATCH 3/6] Refactor: Remove print statement in camera_thread.py and unused code in video_settings.py --- camera_thread.py | 1 - video_settings.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/camera_thread.py b/camera_thread.py index 39a5b22..730d187 100644 --- a/camera_thread.py +++ b/camera_thread.py @@ -143,7 +143,6 @@ def setBackend(self, backend): self.backend = backend self.release() with self.video_capture_mutex: - print("Setting backend to", backend) self.video_capture = cv2.VideoCapture(self.capture_id, self.backend) def isOpened(self): diff --git a/video_settings.py b/video_settings.py index b8aecf8..43b3d79 100644 --- a/video_settings.py +++ b/video_settings.py @@ -13,8 +13,6 @@ def fourcc_to_str(fourcc): return "NULL" if type(fourcc) == str: return fourcc - if type(fourcc) == bytes: - return fourcc.decode() if type(fourcc) != int: fourcc = int(fourcc) return "".join([chr((fourcc >> 8 * i) & 0xFF) for i in range(4)]) From 2075e3d00d5d94446fac6f7b61571c67bea05325 Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Sun, 22 Sep 2024 21:00:47 -0400 Subject: [PATCH 4/6] Refactor: Add API output functionality and UI --- api_output.py | 120 +++++++++++++++++++++++++++++++ camera_thread.py | 3 +- mainwindow.py | 30 +++++++- mainwindow.ui | 13 ++-- scoresight.spec | 1 + text_detection_target.py | 6 +- translations/scoresight_en_US.qm | Bin 27315 -> 27396 bytes translations/scoresight_en_US.ts | 7 +- ui_mainwindow.py | 8 ++- 9 files changed, 171 insertions(+), 17 deletions(-) create mode 100644 api_output.py diff --git a/api_output.py b/api_output.py new file mode 100644 index 0000000..abc2fe4 --- /dev/null +++ b/api_output.py @@ -0,0 +1,120 @@ +import requests +import json +import csv +import io +from functools import partial +import xml.etree.ElementTree as ET +from urllib.parse import urlparse +import threading + +from sc_logging import logger +from text_detection_target import TextDetectionTargetWithResult +from storage import fetch_data, subscribe_to_data + +out_api_url = fetch_data("scoresight.json", "out_api_url", None) +out_api_encoding = fetch_data("scoresight.json", "out_api_encoding", "JSON") + + +def is_valid_url_urllib(url): + try: + result = urlparse(url) + return all([result.scheme, result.netloc]) + except ValueError: + return False + + +def setup_out_api_url(url): + global out_api_url + out_api_url = url + + +def setup_out_api_encoding(encoding): + global out_api_encoding + out_api_encoding = encoding + + +subscribe_to_data("scoresight.json", "out_api_url", setup_out_api_url) +subscribe_to_data("scoresight.json", "out_api_encoding", setup_out_api_encoding) + + +def update_out_api(data: list[TextDetectionTargetWithResult]): + if out_api_url is None or out_api_encoding is None: + logger.error(f"Output API not set up: {out_api_url}, {out_api_encoding}") + return + + # validate the URL + if not is_valid_url_urllib(out_api_url): + logger.error(f"Invalid URL: {out_api_url}") + return + + logger.debug(f"Sending data to output API: {out_api_url}") + + def send_data(): + try: + if out_api_encoding == "JSON": + response = send_json(data) + elif out_api_encoding == "XML": + response = send_xml(data) + elif out_api_encoding == "CSV": + response = send_csv(data) + else: + logger.error("Invalid encoding: %s", out_api_encoding) + return + + if response.status_code != 200: + logger.error( + f"Error sending data to output API: {out_api_url}, {response.status_code}" + ) + except Exception as e: + logger.error(f"Error sending data to output API: {out_api_url}, {e}") + + thread = threading.Thread(target=send_data) + thread.start() + + +def send_json(data: list[TextDetectionTargetWithResult]): + headers = {"Content-Type": "application/json"} + response = requests.post( + out_api_url, + headers=headers, + data=json.dumps([result.to_dict() for result in data]), + ) + return response + + +def send_xml(data: list[TextDetectionTargetWithResult]): + headers = {"Content-Type": "application/xml"} + root = ET.Element("root") + for targetWithResult in data: + resultEl = ET.SubElement(root, "result") + resultEl.set("name", targetWithResult.name) + resultEl.set("result", targetWithResult.result) + resultEl.set("result_state", targetWithResult.result_state.name) + resultEl.set("x", str(targetWithResult.x())) + resultEl.set("y", str(targetWithResult.y())) + resultEl.set("width", str(targetWithResult.width())) + resultEl.set("height", str(targetWithResult.height())) + xml_data = ET.tostring(root, encoding="utf-8") + response = requests.post(out_api_url, headers=headers, data=xml_data) + return response + + +def send_csv(data: list[TextDetectionTargetWithResult]): + headers = {"Content-Type": "text/csv"} + output = io.StringIO() + writer = csv.writer(output) + writer.writerow(["Name", "Text", "State", "X", "Y", "Width", "Height"]) + for result in data: + writer.writerow( + [ + result.name, + result.result, + result.result_state.name, + result.x(), + result.y(), + result.width(), + result.height(), + ] + ) + response = requests.post(out_api_url, headers=headers, data=output.getvalue()) + return response diff --git a/camera_thread.py b/camera_thread.py index 730d187..6356f52 100644 --- a/camera_thread.py +++ b/camera_thread.py @@ -306,8 +306,7 @@ def run(self): self.retry_count += 1 if self.camera_info.type == CameraInfo.CameraType.FILE: logger.debug("Restarting video file") - with self.video_capture_mutex: - self.video_capture.set(cv2.CAP_PROP_POS_FRAMES, 0) + self.video_capture.set(cv2.CAP_PROP_POS_FRAMES, 0) self.sleep_fps_target() continue logger.warn( diff --git a/mainwindow.py b/mainwindow.py index 4563aef..030968c 100644 --- a/mainwindow.py +++ b/mainwindow.py @@ -29,6 +29,7 @@ from os import path from platformdirs import user_data_dir +from api_output import update_out_api from camera_info import CameraInfo from get_camera_info import get_camera_info from http_server import start_http_server, update_http_server @@ -203,6 +204,28 @@ def __init__(self, translator: QTranslator, parent: QObject): self.video_settings_dialog = None self.ui.toolButton_videoSettings.clicked.connect(self.openVideoSettings) + self.ui.lineEdit_api_url.textChanged.connect( + partial(self.globalSettingsChanged, "out_api_url") + ) + self.ui.lineEdit_api_url.setText( + fetch_data("scoresight.json", "out_api_url", "") + ) + + self.ui.checkBox_enableOutAPI.toggled.connect( + partial(self.globalSettingsChanged, "enable_out_api") + ) + self.ui.checkBox_enableOutAPI.setChecked( + fetch_data("scoresight.json", "enable_out_api", False) + ) + self.ui.comboBox_api_encode.currentTextChanged.connect( + partial(self.globalSettingsChanged, "out_api_encoding") + ) + self.ui.comboBox_api_encode.setCurrentIndex( + self.ui.comboBox_api_encode.findText( + fetch_data("scoresight.json", "out_api_encoding", "JSON") + ) + ) + self.obs_websocket_client = None ocr_models = [ @@ -563,9 +586,9 @@ def toggleStopUpdates(self, value): self.updateOCRResults = not value # change the text on the button self.ui.pushButton_stopUpdates.setText( - self.translator.translate("main", "Resume updates") + self.translator.translate("MainWindow", "Resume Updates") if value - else self.translator.translate("main", "Stop updates") + else self.translator.translate("MainWindow", "Stop Updates") ) def selectOutputFolder(self): @@ -1301,6 +1324,9 @@ def ocrResult(self, results: list[TextDetectionTargetWithResult]): update_http_server(results) + if self.ui.checkBox_enableOutAPI.isChecked(): + update_out_api(results) + # update vmix self.vmixUpdater.update_vmix(results) diff --git a/mainwindow.ui b/mainwindow.ui index 293c227..d8780b3 100644 --- a/mainwindow.ui +++ b/mainwindow.ui @@ -1467,12 +1467,7 @@ - Key-Value - - - - - Plain Text + CSV @@ -1517,6 +1512,12 @@ + + false + + + Not implemented yet. + Websocket diff --git a/scoresight.spec b/scoresight.spec index f8ec012..36969c9 100644 --- a/scoresight.spec +++ b/scoresight.spec @@ -48,6 +48,7 @@ datas = [ ] sources = [ + 'api_output.py', 'camera_info.py', 'camera_view.py', 'camera_thread.py', diff --git a/text_detection_target.py b/text_detection_target.py index da9a4d4..c626bb7 100644 --- a/text_detection_target.py +++ b/text_detection_target.py @@ -39,7 +39,7 @@ def clear(self): class TextDetectionTarget(QRectF): - def __init__(self, x, y, width, height, name, settings: dict | None = None): + def __init__(self, x, y, width, height, name: str, settings: dict | None = None): super().__init__(x, y, width, height) self.name = name self.settings = settings @@ -58,8 +58,8 @@ class ResultState(enum.Enum): def __init__( self, detection_target: TextDetectionTarget, - result, - result_state, + result: str, + result_state: ResultState, effectiveRect=None, extras=None, ): diff --git a/translations/scoresight_en_US.qm b/translations/scoresight_en_US.qm index ecbbdfda81a90a4cd72c25269626876b49c48414..943d604b861d8cb08c9c9b36cb4722bea7ffbbc5 100644 GIT binary patch delta 507 zcmdmdm9ga-;{-Xz6BFgjC7&{|uKB{ikbInhbqfarLs}gJ>w#1ThSZ#i>)q;;l^9~0 zw*p01GOP=k%)k&`!EkJcFatx362q%!p+LSYqg6sM14EJkqr?s}XbO z!h8mXI4=U||0$ull{orI1pe^q?^vqd-{3MQ5kDmZtwSZ$y z`Vn7N+Y0r1gGVUQlRg>IGdIN!^o!rLZ@Ayyq!s2 zFo5&VX?F&OfL<;(UMB{ISd+=mndB8GbN!nS456T}+#(MbFfjOE=9bjx0@|I)t;7Iy zOU#zZ>dgLg5jc@xe``#0RYQmm|y?^ delta 421 zcmZp<#<=+^;{-Xz0~6)TCFd}(uKB{ikX*;Wx`l&*AjsMfg$B3bLYZ*28P&!lk*wlnBsp7?S3* zADnCgG>9aO9qCpevacNj~N(Zws2frX#|vy z;Iy1k3iKW$XVX$(Aow&u=(N7c`E`{F2n>)Q uaXy<1K>JcR@s;NS`4Mq^U;kcZV2GEVyq{T{Dec7OH_RF_n`0A}a{&NFWr9lp diff --git a/translations/scoresight_en_US.ts b/translations/scoresight_en_US.ts index 70190bb..6cab4bc 100644 --- a/translations/scoresight_en_US.ts +++ b/translations/scoresight_en_US.ts @@ -519,10 +519,15 @@ 1 - + Stop Updates Stop Updates + + + Resume Updates + Resume Updates + Detections / s diff --git a/ui_mainwindow.py b/ui_mainwindow.py index 77db006..096bd83 100644 --- a/ui_mainwindow.py +++ b/ui_mainwindow.py @@ -774,7 +774,6 @@ def setupUi(self, MainWindow): self.comboBox_api_encode.addItem("") self.comboBox_api_encode.addItem("") self.comboBox_api_encode.addItem("") - self.comboBox_api_encode.addItem("") self.comboBox_api_encode.setObjectName(u"comboBox_api_encode") self.formLayout_3.setWidget(4, QFormLayout.FieldRole, self.comboBox_api_encode) @@ -796,6 +795,7 @@ def setupUi(self, MainWindow): self.checkBox_is_websocket = QCheckBox(self.widget_24) self.checkBox_is_websocket.setObjectName(u"checkBox_is_websocket") + self.checkBox_is_websocket.setEnabled(False) self.horizontalLayout_27.addWidget(self.checkBox_is_websocket) @@ -1209,10 +1209,12 @@ def retranslateUi(self, MainWindow): self.label_21.setText(QCoreApplication.translate("MainWindow", u"Encode", None)) self.comboBox_api_encode.setItemText(0, QCoreApplication.translate("MainWindow", u"JSON", None)) self.comboBox_api_encode.setItemText(1, QCoreApplication.translate("MainWindow", u"XML", None)) - self.comboBox_api_encode.setItemText(2, QCoreApplication.translate("MainWindow", u"Key-Value", None)) - self.comboBox_api_encode.setItemText(3, QCoreApplication.translate("MainWindow", u"Plain Text", None)) + self.comboBox_api_encode.setItemText(2, QCoreApplication.translate("MainWindow", u"CSV", None)) self.lineEdit_api_url.setPlaceholderText(QCoreApplication.translate("MainWindow", u"http://", None)) +#if QT_CONFIG(tooltip) + self.checkBox_is_websocket.setToolTip(QCoreApplication.translate("MainWindow", u"Not implemented yet.", None)) +#endif // QT_CONFIG(tooltip) self.checkBox_is_websocket.setText(QCoreApplication.translate("MainWindow", u"Websocket", None)) self.label_20.setText(QCoreApplication.translate("MainWindow", u"URL", None)) self.tabWidget_outputs.setTabText(self.tabWidget_outputs.indexOf(self.tab_api), QCoreApplication.translate("MainWindow", u"API", None)) From c434745c5c2cb3e950f879a05d3ae752314655e4 Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Sun, 22 Sep 2024 23:08:41 -0400 Subject: [PATCH 5/6] Refactor: Add OCR training data functionality and UI --- mainwindow.py | 23 ++++- mainwindow.ui | 14 ++- ocr_training_data.py | 184 +++++++++++++++++++++++++++++++++ ocr_training_data_dialog.ui | 175 +++++++++++++++++++++++++++++++ scoresight.spec | 2 + tesseract.py | 44 +++++--- text_detection_target.py | 14 +++ ui_mainwindow.py | 11 +- ui_ocr_training_data_dialog.py | 121 ++++++++++++++++++++++ 9 files changed, 566 insertions(+), 22 deletions(-) create mode 100644 ocr_training_data.py create mode 100644 ocr_training_data_dialog.ui create mode 100644 ui_ocr_training_data_dialog.py diff --git a/mainwindow.py b/mainwindow.py index 030968c..a4497be 100644 --- a/mainwindow.py +++ b/mainwindow.py @@ -33,6 +33,7 @@ from camera_info import CameraInfo from get_camera_info import get_camera_info from http_server import start_http_server, update_http_server +from ocr_training_data import OCRTrainingDataDialog from screen_capture_source import ScreenCapture from source_view import ImageViewer from defaults import ( @@ -120,6 +121,7 @@ def __init__(self, translator: QTranslator, parent: QObject): file_menu.addAction("Import Configuration", self.importConfiguration) file_menu.addAction("Export Configuration", self.exportConfiguration) file_menu.addAction("Open Configuration Folder", self.openConfigurationFolder) + file_menu.addAction("OCR Training Data Setup", self.openOCRTrainingDataDialog) # Add "Language" menu languageMenu = file_menu.addMenu("Language") @@ -400,10 +402,20 @@ def __init__(self, translator: QTranslator, parent: QObject): partial(self.globalSettingsChanged, "vmix_send_same") ) + self.ui.pushButton_saveOCRTrainingData.clicked.connect(self.saveOCRTrainingData) + self.ui.pushButton_saveOCRTrainingData.setChecked( + fetch_data("scoresight.json", "save_ocr_training_data", False) + ) + self.update_sources.connect(self.updateSources) self.get_sources.connect(self.getSources) self.get_sources.emit() + def saveOCRTrainingData(self): + self.globalSettingsChanged( + "save_ocr_training_data", self.ui.pushButton_saveOCRTrainingData.isChecked() + ) + def openVideoSettings(self): # only allow opening the video settings for an OpenCV type source if not self.image_viewer or self.image_viewer is None: @@ -475,7 +487,10 @@ def changeLanguage(self, locale): if appInstance: logger.info(f"installing translator for {locale}") appInstance.installTranslator(self.translator) - self.ui.retranslateUi(self) + try: + self.ui.retranslateUi(self) + except Exception as e: + logger.error(f"Error retranslating UI: {e}") def addLanguageOption(self, menu: QMenu, language_name: str, locale: str): menu.addAction(language_name, lambda: self.changeLanguage(locale)) @@ -521,6 +536,12 @@ def exportConfiguration(self): # save the configuration to the file self.detectionTargetsStorage.saveBoxesToFile(file) + def openOCRTrainingDataDialog(self): + # open the OCR training data dialog + dialog = OCRTrainingDataDialog() + dialog.setWindowTitle("OCR Training Data Setup") + dialog.exec() + def openConfigurationFolder(self): # open the configuration folder in the file explorer QDesktopServices.openUrl( diff --git a/mainwindow.ui b/mainwindow.ui index d8780b3..7668999 100644 --- a/mainwindow.ui +++ b/mainwindow.ui @@ -6,7 +6,7 @@ 0 0 - 961 + 901 725 @@ -1764,6 +1764,16 @@ + + + + Save OCR Training Data + + + true + + + @@ -2120,7 +2130,7 @@ 0 0 - 961 + 901 20 diff --git a/ocr_training_data.py b/ocr_training_data.py new file mode 100644 index 0000000..4c029bf --- /dev/null +++ b/ocr_training_data.py @@ -0,0 +1,184 @@ +import tempfile +import zipfile +from PySide6.QtWidgets import QDialog, QFileDialog, QMessageBox +import cv2 +from numpy import ndarray +from platformdirs import user_data_dir +import os +import uuid + +from sc_logging import logger +from text_detection_target import TextDetectionResult, TextDetectionTargetWithResult +from ui_ocr_training_data_dialog import Ui_OCRTrainingDataDialog +from storage import fetch_data, store_data, subscribe_to_data + + +class OCRTrainingDataOptions: + def __init__(self): + self.save_ocr_training_data = fetch_data( + "scoresight.json", "save_ocr_training_data", False + ) + subscribe_to_data( + "scoresight.json", "save_ocr_training_data", self.set_save_ocr_training_data + ) + self.ocr_training_data_folder = fetch_data( + "scoresight.json", + "ocr_training_data_folder", + os.path.join(user_data_dir("scoresight"), "ocr_training_data"), + ) + subscribe_to_data( + "scoresight.json", + "ocr_training_data_folder", + self.set_ocr_training_data_folder, + ) + self.ocr_training_data_max_size = fetch_data( + "scoresight.json", "ocr_training_data_max_size", 10 + ) + subscribe_to_data( + "scoresight.json", + "ocr_training_data_max_size", + self.set_ocr_training_data_max_size, + ) + + def set_save_ocr_training_data(self, value): + self.save_ocr_training_data = value + + def set_ocr_training_data_folder(self, value): + self.ocr_training_data_folder = value + + def set_ocr_training_data_max_size(self, value): + self.ocr_training_data_max_size = value + + def save_ocr_result_to_folder( + self, image: ndarray, image_gray: ndarray, result: TextDetectionResult + ): + if self.save_ocr_training_data: + # create the folder if it doesn't exist + if not os.path.exists(self.ocr_training_data_folder): + os.makedirs(self.ocr_training_data_folder) + + if ( + result.state == TextDetectionTargetWithResult.ResultState.SameNoChange + or result.state == TextDetectionTargetWithResult.ResultState.Empty + or result.text == "" + ): + return + + # generate a name for the image and text file using uuid + uuid_for_image = uuid.uuid4() + image_name = f"{uuid_for_image}.png" + image_gray_name = f"{uuid_for_image}_gray.png" + text_name = f"{uuid_for_image}.txt" + + # write the image to the folder + cv2.imwrite(os.path.join(self.ocr_training_data_folder, image_name), image) + cv2.imwrite( + os.path.join(self.ocr_training_data_folder, image_gray_name), image_gray + ) + + # write the text to the folder + with open( + os.path.join(self.ocr_training_data_folder, text_name), "w" + ) as text_file: + text_file.write(result.text) + + +ocr_training_data_options = OCRTrainingDataOptions() + + +def zip_folder(folder_path, zip_path): + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: + for root, _, files in os.walk(folder_path): + for file in files: + file_path = os.path.join(root, file) + arcname = os.path.relpath(file_path, folder_path) + zipf.write(file_path, arcname) + + +class OCRTrainingDataDialog(QDialog): + def __init__(self): + super().__init__() + self.ui = Ui_OCRTrainingDataDialog() + self.ui.setupUi(self) + self.ui.buttonBox.accepted.connect(self.save_settings) + self.ui.buttonBox.rejected.connect(self.close) + self.ui.toolButton_chooseSaveFolder.clicked.connect(self.choose_save_folder) + self.ui.lineEdit_saveFolder.setText( + ocr_training_data_options.ocr_training_data_folder + ) + self.ui.spinBox_maxSize.setValue( + ocr_training_data_options.ocr_training_data_max_size + ) + self.ui.pushButton_openFolder.clicked.connect(self.open_folder) + self.ui.pushButton_saveZipFile.clicked.connect(self.save_zip_file) + + def save_zip_file(self): + logger.debug("Saving OCR training data zip file") + folder = self.ui.lineEdit_saveFolder.text() + if folder: + # Create a temporary file to store the zip + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_zip: + temp_zip_path = temp_zip.name + # Zip the folder + try: + zip_folder(folder, temp_zip_path) + except Exception as e: + QMessageBox.critical( + None, "Error", f"Failed to create zip file: {str(e)}" + ) + os.unlink(temp_zip_path) + return + + # Ask user where to save the zip file + save_path, _ = QFileDialog.getSaveFileName( + None, "Save Zip File", "", "Zip Files (*.zip)" + ) + if save_path: + if not save_path.endswith(".zip"): + save_path += ".zip" + try: + # Copy the temp zip to the chosen location + with open(temp_zip_path, "rb") as src, open(save_path, "wb") as dst: + dst.write(src.read()) + logger.info(f"Zip file saved to {save_path}") + except Exception as e: + QMessageBox.critical( + None, "Error", f"Failed to save zip file: {str(e)}" + ) + else: + logger.info("Zip file was not saved.") + + # Clean up the temporary file + os.unlink(temp_zip_path) + + else: + logger.error("No OCR training data folder set") + + def open_folder(self): + logger.debug("Opening OCR training data save folder") + folder = self.ui.lineEdit_saveFolder.text() + if folder: + os.startfile(folder) + + def choose_save_folder(self): + logger.debug("Choosing OCR training data save folder") + folder = self.ui.lineEdit_saveFolder.text() + folder = QFileDialog.getExistingDirectory( + self, "Choose OCR training data save folder", folder + ) + if folder: + self.ui.lineEdit_saveFolder.setText(folder) + + def save_settings(self): + logger.debug("Saving OCR training data") + store_data( + "scoresight.json", + "ocr_training_data_save_folder", + self.ui.lineEdit_saveFolder.text(), + ) + store_data( + "scoresight.json", + "ocr_training_data_max_size", + self.ui.spinBox_maxSize.value(), + ) + self.close() diff --git a/ocr_training_data_dialog.ui b/ocr_training_data_dialog.ui new file mode 100644 index 0000000..8781bc2 --- /dev/null +++ b/ocr_training_data_dialog.ui @@ -0,0 +1,175 @@ + + + OCRTrainingDataDialog + + + + 0 + 0 + 267 + 132 + + + + + 267 + 0 + + + + + 650 + 16777215 + + + + Dialog + + + + + + Qt::Horizontal + + + QDialogButtonBox::Cancel|QDialogButtonBox::Ok + + + + + + + + + + Save Folder + + + + + + + + 0 + 0 + + + + + 0 + + + 0 + + + 0 + + + 0 + + + + + + + + ... + + + + + + + + + + + 0 + + + 0 + + + 0 + + + 0 + + + + + Open Folder + + + + + + + Save Zip File + + + + + + + + + + Max Size + + + + + + + Mb + + + 1 + + + 10 + + + + + + + + + + + + buttonBox + accepted() + OCRTrainingDataDialog + accept() + + + 248 + 254 + + + 157 + 274 + + + + + buttonBox + rejected() + OCRTrainingDataDialog + reject() + + + 316 + 260 + + + 286 + 274 + + + + + diff --git a/scoresight.spec b/scoresight.spec index 36969c9..52aaaee 100644 --- a/scoresight.spec +++ b/scoresight.spec @@ -61,6 +61,7 @@ sources = [ 'main.py', 'mainwindow.py', 'ndi.py', + 'ocr_training_data.py', 'obs_websocket.py', 'resizable_rect.py', 'sc_logging.py', @@ -74,6 +75,7 @@ sources = [ 'ui_connect_obs.py', 'ui_log_view.py', 'ui_mainwindow.py', + 'ui_ocr_training_data_dialog.py', 'ui_screen_capture.py', 'ui_update_available.py', 'ui_url_source.py', diff --git a/tesseract.py b/tesseract.py index fe25e70..f4fb8fb 100644 --- a/tesseract.py +++ b/tesseract.py @@ -1,15 +1,20 @@ +from defaults import FieldType from os import path -import cv2 -from tesserocr import PyTessBaseAPI, RIL, iterate_level -import numpy as np from PIL import Image -from defaults import FieldType -from storage import fetch_data -from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult -import re from PySide6.QtCore import QRectF +from tesserocr import PyTessBaseAPI, RIL, iterate_level from threading import Lock -from sc_logging import logger +import cv2 +import numpy as np +import re + +from ocr_training_data import ocr_training_data_options +from storage import fetch_data +from text_detection_target import ( + TextDetectionResult, + TextDetectionTarget, + TextDetectionTargetWithResult, +) def autocrop(image_in): @@ -78,14 +83,6 @@ def is_valid_regex(pattern): return False -class TextDetectionResult: - def __init__(self, text, state, rect=None, extra=None): - self.text = text - self.state = state - self.rect = rect - self.extra = extra - - class TextDetector: # model name enum: daktronics=0, scoreboard_general=1 class OcrModelIndex: @@ -570,5 +567,18 @@ def detect_multi_text( if text == "": textstate = TextDetectionTargetWithResult.ResultState.Empty - texts.append(TextDetectionResult(text, textstate, effectiveRect, extras)) + result = TextDetectionResult(text, textstate, effectiveRect, extras) + + if ocr_training_data_options.save_ocr_training_data: + # crop the image from the grayscale image + imagecrop_gray = gray[ + int(rect.y()) : int(rect.y() + rect.height()), + int(rect.x()) : int(rect.x() + rect.width()), + ] + # save the image and the text + ocr_training_data_options.save_ocr_result_to_folder( + imagecrop, imagecrop_gray, result + ) + + texts.append(result) return texts diff --git a/text_detection_target.py b/text_detection_target.py index c626bb7..451f309 100644 --- a/text_detection_target.py +++ b/text_detection_target.py @@ -88,3 +88,17 @@ def to_dict(self): "height": self.height(), }, } + + +class TextDetectionResult: + def __init__( + self, + text: str, + state: TextDetectionTargetWithResult.ResultState, + rect: QRectF | None = None, + extra=None, + ): + self.text = text + self.state = state + self.rect = rect + self.extra = extra diff --git a/ui_mainwindow.py b/ui_mainwindow.py index 096bd83..20c61ac 100644 --- a/ui_mainwindow.py +++ b/ui_mainwindow.py @@ -27,7 +27,7 @@ class Ui_MainWindow(object): def setupUi(self, MainWindow): if not MainWindow.objectName(): MainWindow.setObjectName(u"MainWindow") - MainWindow.resize(961, 725) + MainWindow.resize(901, 725) self.centralwidget = QWidget(MainWindow) self.centralwidget.setObjectName(u"centralwidget") self.formLayout = QFormLayout(self.centralwidget) @@ -916,6 +916,12 @@ def setupUi(self, MainWindow): self.horizontalLayout_4.addItem(self.horizontalSpacer) + self.pushButton_saveOCRTrainingData = QPushButton(self.widget_4) + self.pushButton_saveOCRTrainingData.setObjectName(u"pushButton_saveOCRTrainingData") + self.pushButton_saveOCRTrainingData.setCheckable(True) + + self.horizontalLayout_4.addWidget(self.pushButton_saveOCRTrainingData) + self.verticalLayout_2.addWidget(self.widget_4) @@ -1088,7 +1094,7 @@ def setupUi(self, MainWindow): MainWindow.setCentralWidget(self.centralwidget) self.menubar = QMenuBar(MainWindow) self.menubar.setObjectName(u"menubar") - self.menubar.setGeometry(QRect(0, 0, 961, 20)) + self.menubar.setGeometry(QRect(0, 0, 901, 20)) MainWindow.setMenuBar(self.menubar) self.retranslateUi(MainWindow) @@ -1235,6 +1241,7 @@ def retranslateUi(self, MainWindow): self.pushButton_refresh_sources.setToolTip(QCoreApplication.translate("MainWindow", u"Refresh Sources", None)) #endif // QT_CONFIG(tooltip) self.pushButton_refresh_sources.setText(QCoreApplication.translate("MainWindow", u"Reload", None)) + self.pushButton_saveOCRTrainingData.setText(QCoreApplication.translate("MainWindow", u"Save OCR Training Data", None)) self.pushButton_binary.setText(QCoreApplication.translate("MainWindow", u"Binary View", None)) self.pushButton_fourCorner.setText(QCoreApplication.translate("MainWindow", u"4-corner Correction", None)) #if QT_CONFIG(tooltip) diff --git a/ui_ocr_training_data_dialog.py b/ui_ocr_training_data_dialog.py new file mode 100644 index 0000000..5f51d1a --- /dev/null +++ b/ui_ocr_training_data_dialog.py @@ -0,0 +1,121 @@ +# -*- coding: utf-8 -*- + +################################################################################ +## Form generated from reading UI file 'ocr_training_data_dialog.ui' +## +## Created by: Qt User Interface Compiler version 6.7.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 (QAbstractButton, QApplication, QDialog, QDialogButtonBox, + QFormLayout, QGridLayout, QHBoxLayout, QLabel, + QLineEdit, QPushButton, QSizePolicy, QSpinBox, + QToolButton, QWidget) + +class Ui_OCRTrainingDataDialog(object): + def setupUi(self, OCRTrainingDataDialog): + if not OCRTrainingDataDialog.objectName(): + OCRTrainingDataDialog.setObjectName(u"OCRTrainingDataDialog") + OCRTrainingDataDialog.resize(267, 132) + OCRTrainingDataDialog.setMinimumSize(QSize(267, 0)) + OCRTrainingDataDialog.setMaximumSize(QSize(650, 16777215)) + self.gridLayout = QGridLayout(OCRTrainingDataDialog) + self.gridLayout.setObjectName(u"gridLayout") + self.buttonBox = QDialogButtonBox(OCRTrainingDataDialog) + self.buttonBox.setObjectName(u"buttonBox") + self.buttonBox.setOrientation(Qt.Horizontal) + self.buttonBox.setStandardButtons(QDialogButtonBox.Cancel|QDialogButtonBox.Ok) + + self.gridLayout.addWidget(self.buttonBox, 2, 0, 1, 1) + + self.widget = QWidget(OCRTrainingDataDialog) + self.widget.setObjectName(u"widget") + self.formLayout = QFormLayout(self.widget) + self.formLayout.setObjectName(u"formLayout") + self.label = QLabel(self.widget) + self.label.setObjectName(u"label") + + self.formLayout.setWidget(0, QFormLayout.LabelRole, self.label) + + self.widget_2 = QWidget(self.widget) + self.widget_2.setObjectName(u"widget_2") + sizePolicy = QSizePolicy(QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Fixed) + sizePolicy.setHorizontalStretch(0) + sizePolicy.setVerticalStretch(0) + sizePolicy.setHeightForWidth(self.widget_2.sizePolicy().hasHeightForWidth()) + self.widget_2.setSizePolicy(sizePolicy) + self.horizontalLayout = QHBoxLayout(self.widget_2) + self.horizontalLayout.setObjectName(u"horizontalLayout") + self.horizontalLayout.setContentsMargins(0, 0, 0, 0) + self.lineEdit_saveFolder = QLineEdit(self.widget_2) + self.lineEdit_saveFolder.setObjectName(u"lineEdit_saveFolder") + + self.horizontalLayout.addWidget(self.lineEdit_saveFolder) + + self.toolButton_chooseSaveFolder = QToolButton(self.widget_2) + self.toolButton_chooseSaveFolder.setObjectName(u"toolButton_chooseSaveFolder") + + self.horizontalLayout.addWidget(self.toolButton_chooseSaveFolder) + + + self.formLayout.setWidget(0, QFormLayout.FieldRole, self.widget_2) + + self.widget_3 = QWidget(self.widget) + self.widget_3.setObjectName(u"widget_3") + self.horizontalLayout_2 = QHBoxLayout(self.widget_3) + self.horizontalLayout_2.setObjectName(u"horizontalLayout_2") + self.horizontalLayout_2.setContentsMargins(0, 0, 0, 0) + self.pushButton_openFolder = QPushButton(self.widget_3) + self.pushButton_openFolder.setObjectName(u"pushButton_openFolder") + + self.horizontalLayout_2.addWidget(self.pushButton_openFolder) + + self.pushButton_saveZipFile = QPushButton(self.widget_3) + self.pushButton_saveZipFile.setObjectName(u"pushButton_saveZipFile") + + self.horizontalLayout_2.addWidget(self.pushButton_saveZipFile) + + + self.formLayout.setWidget(1, QFormLayout.FieldRole, self.widget_3) + + self.label_2 = QLabel(self.widget) + self.label_2.setObjectName(u"label_2") + + self.formLayout.setWidget(2, QFormLayout.LabelRole, self.label_2) + + self.spinBox_maxSize = QSpinBox(self.widget) + self.spinBox_maxSize.setObjectName(u"spinBox_maxSize") + self.spinBox_maxSize.setMinimum(1) + self.spinBox_maxSize.setValue(10) + + self.formLayout.setWidget(2, QFormLayout.FieldRole, self.spinBox_maxSize) + + + self.gridLayout.addWidget(self.widget, 1, 0, 1, 1) + + + self.retranslateUi(OCRTrainingDataDialog) + self.buttonBox.accepted.connect(OCRTrainingDataDialog.accept) + self.buttonBox.rejected.connect(OCRTrainingDataDialog.reject) + + QMetaObject.connectSlotsByName(OCRTrainingDataDialog) + # setupUi + + def retranslateUi(self, OCRTrainingDataDialog): + OCRTrainingDataDialog.setWindowTitle(QCoreApplication.translate("OCRTrainingDataDialog", u"Dialog", None)) + self.label.setText(QCoreApplication.translate("OCRTrainingDataDialog", u"Save Folder", None)) + self.toolButton_chooseSaveFolder.setText(QCoreApplication.translate("OCRTrainingDataDialog", u"...", None)) + self.pushButton_openFolder.setText(QCoreApplication.translate("OCRTrainingDataDialog", u"Open Folder", None)) + self.pushButton_saveZipFile.setText(QCoreApplication.translate("OCRTrainingDataDialog", u"Save Zip File", None)) + self.label_2.setText(QCoreApplication.translate("OCRTrainingDataDialog", u"Max Size", None)) + self.spinBox_maxSize.setSuffix(QCoreApplication.translate("OCRTrainingDataDialog", u"Mb", None)) + # retranslateUi + From 446f7540acef8bbe714422d5463e1c5221f72811 Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Mon, 23 Sep 2024 08:40:44 -0400 Subject: [PATCH 6/6] Refactor: Update OCR training data saving logic --- camera_thread.py | 7 +++++ ocr_training_data.py | 70 +++++++++++++++++++++++++++++--------------- tesseract.py | 12 -------- 3 files changed, 54 insertions(+), 35 deletions(-) diff --git a/camera_thread.py b/camera_thread.py index 6356f52..f9dcd62 100644 --- a/camera_thread.py +++ b/camera_thread.py @@ -14,6 +14,7 @@ from text_detection_target import TextDetectionTargetWithResult from sc_logging import logger from frame_stabilizer import FrameStabilizer +from ocr_training_data import ocr_training_data_options # Function to set the resolution @@ -412,6 +413,12 @@ def run(self): # emit the results self.ocr_result_signal.emit(results) + if ocr_training_data_options.save_ocr_training_data: + # save the image and the text + ocr_training_data_options.save_ocr_results_to_folder( + binary, gray, results + ) + # Emit the signal to update the pixmap once per second time_diff_prev = (current_time - self.last_emit_time).total_seconds() * 1000 if time_diff_prev >= self.preview_frame_interval: diff --git a/ocr_training_data.py b/ocr_training_data.py index 4c029bf..40308dd 100644 --- a/ocr_training_data.py +++ b/ocr_training_data.py @@ -49,38 +49,62 @@ def set_ocr_training_data_folder(self, value): def set_ocr_training_data_max_size(self, value): self.ocr_training_data_max_size = value - def save_ocr_result_to_folder( - self, image: ndarray, image_gray: ndarray, result: TextDetectionResult + def save_ocr_results_to_folder( + self, + binary: ndarray, + gray: ndarray, + result: list[TextDetectionTargetWithResult], ): if self.save_ocr_training_data: # create the folder if it doesn't exist if not os.path.exists(self.ocr_training_data_folder): os.makedirs(self.ocr_training_data_folder) - if ( - result.state == TextDetectionTargetWithResult.ResultState.SameNoChange - or result.state == TextDetectionTargetWithResult.ResultState.Empty - or result.text == "" - ): - return + # calculate the size of the folder + folder_size = 0 + for root, _, files in os.walk(self.ocr_training_data_folder): + for file in files: + folder_size += os.path.getsize(os.path.join(root, file)) - # generate a name for the image and text file using uuid - uuid_for_image = uuid.uuid4() - image_name = f"{uuid_for_image}.png" - image_gray_name = f"{uuid_for_image}_gray.png" - text_name = f"{uuid_for_image}.txt" + # check if the folder size is greater than the max size + if folder_size > self.ocr_training_data_max_size * 1024 * 1024: + logger.error( + f"OCR training data folder size exceeds maximum size of {self.ocr_training_data_max_size} MB" + ) + return - # write the image to the folder - cv2.imwrite(os.path.join(self.ocr_training_data_folder, image_name), image) - cv2.imwrite( - os.path.join(self.ocr_training_data_folder, image_gray_name), image_gray - ) + for r in result: + if ( + r.result_state == TextDetectionTargetWithResult.ResultState.SameNoChange + or r.result_state == TextDetectionTargetWithResult.ResultState.Empty + or r.result == "" + ): + continue + + # generate a name for the image and text file using uuid + uuid_for_image = uuid.uuid4() + image_name = f"{uuid_for_image}.png" + image_gray_name = f"{uuid_for_image}_gray.png" + text_name = f"{uuid_for_image}.txt" + + # crop the patch from the image + x, y, w, h = int(r.x()), int(r.y()), int(r.width()), int(r.height()) + binary_patch = binary[y : y + h, x : x + w] + gray_patch = gray[y : y + h, x : x + w] + + # write the image to the folder + cv2.imwrite( + os.path.join(self.ocr_training_data_folder, image_name), binary_patch + ) + cv2.imwrite( + os.path.join(self.ocr_training_data_folder, image_gray_name), gray_patch + ) - # write the text to the folder - with open( - os.path.join(self.ocr_training_data_folder, text_name), "w" - ) as text_file: - text_file.write(result.text) + # write the text to the folder + with open( + os.path.join(self.ocr_training_data_folder, text_name), "w" + ) as text_file: + text_file.write(r.result) ocr_training_data_options = OCRTrainingDataOptions() diff --git a/tesseract.py b/tesseract.py index f4fb8fb..490367a 100644 --- a/tesseract.py +++ b/tesseract.py @@ -8,7 +8,6 @@ import numpy as np import re -from ocr_training_data import ocr_training_data_options from storage import fetch_data from text_detection_target import ( TextDetectionResult, @@ -569,16 +568,5 @@ def detect_multi_text( result = TextDetectionResult(text, textstate, effectiveRect, extras) - if ocr_training_data_options.save_ocr_training_data: - # crop the image from the grayscale image - imagecrop_gray = gray[ - int(rect.y()) : int(rect.y() + rect.height()), - int(rect.x()) : int(rect.x() + rect.width()), - ] - # save the image and the text - ocr_training_data_options.save_ocr_result_to_folder( - imagecrop, imagecrop_gray, result - ) - texts.append(result) return texts