Merge pull request #38 from locaal-ai/roy.ocr_training_dojo

Refactor: Add training dojo functionality and speed control
locaal-ai · Sep 25, 2024 · 3d8f89f · 3d8f89f
2 parents 84d950a + c807336
commit 3d8f89f
Show file tree

Hide file tree

Showing 11 changed files with 681 additions and 27 deletions.
diff --git a/camera_thread.py b/camera_thread.py
@@ -206,6 +206,13 @@ def __init__(
         self.fps_alpha = 0.1  # Smoothing factor
         self.updateOnChange = True
         self.crop = FrameCropAndRotation()
+        self.speed = 1
+
+    def getSpeed(self):
+        return self.speed
+
+    def setSpeed(self, speed):
+        self.speed = speed
 
     def setUpdateFrameInterval(self, cadence):
         self.update_frame_interval = 1000 / cadence
@@ -303,6 +310,13 @@ def run(self):
                 self.sleep_fps_target()
                 continue
 
+            if self.camera_info.type == CameraInfo.CameraType.FILE:
+                if self.speed != 1:
+                    self.video_capture.set(
+                        cv2.CAP_PROP_POS_FRAMES,
+                        self.video_capture.get(cv2.CAP_PROP_POS_FRAMES) + self.speed,
+                    )
+
             if not ret:
                 self.retry_count += 1
                 if self.camera_info.type == CameraInfo.CameraType.FILE:

diff --git a/mainwindow.py b/mainwindow.py
@@ -407,10 +407,34 @@ def __init__(self, translator: QTranslator, parent: QObject):
             fetch_data("scoresight.json", "save_ocr_training_data", False)
         )
 
+        self.ui.toolButton_speed.clicked.connect(self.toggleSpeed)
+
         self.update_sources.connect(self.updateSources)
         self.get_sources.connect(self.getSources)
         self.get_sources.emit()
 
+    def toggleSpeed(self):
+        # check the current speed and toggle it
+        # possible speeds are x2, x4, x8, x16, x32 and back to x1
+        # change the button text to the current speed
+        # change the speed of the image viewer
+        if self.image_viewer:
+            speed = self.image_viewer.timerThread.getSpeed()
+            if speed == 1:
+                speed = 2
+            elif speed == 2:
+                speed = 4
+            elif speed == 4:
+                speed = 8
+            elif speed == 8:
+                speed = 16
+            elif speed == 16:
+                speed = 32
+            else:
+                speed = 1
+            self.image_viewer.timerThread.setSpeed(speed)
+            self.ui.toolButton_speed.setText(f"x{speed}")
+
     def saveOCRTrainingData(self):
         self.globalSettingsChanged(
             "save_ocr_training_data", self.ui.pushButton_saveOCRTrainingData.isChecked()

diff --git a/mainwindow.ui b/mainwindow.ui
@@ -2058,6 +2058,20 @@
             </property>
            </widget>
           </item>
+          <item>
+           <widget class="QLabel" name="label_22">
+            <property name="text">
+             <string>Speed</string>
+            </property>
+           </widget>
+          </item>
+          <item>
+           <widget class="QToolButton" name="toolButton_speed">
+            <property name="text">
+             <string>x1</string>
+            </property>
+           </widget>
+          </item>
           <item>
            <spacer name="horizontalSpacer_4">
             <property name="orientation">

diff --git a/ocr_training_data.py b/ocr_training_data.py
@@ -135,6 +135,18 @@ def __init__(self):
         )
         self.ui.pushButton_openFolder.clicked.connect(self.open_folder)
         self.ui.pushButton_saveZipFile.clicked.connect(self.save_zip_file)
+        self.ui.pushButton_openTrainingDojo.clicked.connect(self.open_training_dojo)
+
+    def open_training_dojo(self):
+        logger.debug("Opening OCR training dojo")
+        folder = self.ui.lineEdit_saveFolder.text()
+        if folder:
+            from training_dojo import TrainingDojo
+
+            training_dojo = TrainingDojo(folder)
+            training_dojo.exec_()
+        else:
+            logger.error("No OCR training data folder set")
 
     def save_zip_file(self):
         logger.debug("Saving OCR training data zip file")

diff --git a/ocr_training_data_dialog.ui b/ocr_training_data_dialog.ui
@@ -6,8 +6,8 @@
    <rect>
     <x>0</x>
     <y>0</y>
-    <width>267</width>
-    <height>132</height>
+    <width>377</width>
+    <height>164</height>
    </rect>
   </property>
   <property name="minimumSize">
@@ -19,25 +19,21 @@
   <property name="maximumSize">
    <size>
     <width>650</width>
-    <height>16777215</height>
+    <height>300</height>
    </size>
   </property>
   <property name="windowTitle">
    <string>Dialog</string>
   </property>
   <layout class="QGridLayout" name="gridLayout">
-   <item row="2" column="0">
-    <widget class="QDialogButtonBox" name="buttonBox">
-     <property name="orientation">
-      <enum>Qt::Horizontal</enum>
-     </property>
-     <property name="standardButtons">
-      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>
-     </property>
-    </widget>
-   </item>
    <item row="1" column="0">
     <widget class="QWidget" name="widget" native="true">
+     <property name="sizePolicy">
+      <sizepolicy hsizetype="Preferred" vsizetype="Fixed">
+       <horstretch>0</horstretch>
+       <verstretch>0</verstretch>
+      </sizepolicy>
+     </property>
      <layout class="QFormLayout" name="formLayout">
       <item row="0" column="0">
        <widget class="QLabel" name="label">
@@ -132,9 +128,39 @@
         </property>
        </widget>
       </item>
+      <item row="3" column="1">
+       <widget class="QPushButton" name="pushButton_openTrainingDojo">
+        <property name="text">
+         <string>Open OCR Training Dojo</string>
+        </property>
+       </widget>
+      </item>
      </layout>
     </widget>
    </item>
+   <item row="3" column="0">
+    <widget class="QDialogButtonBox" name="buttonBox">
+     <property name="orientation">
+      <enum>Qt::Horizontal</enum>
+     </property>
+     <property name="standardButtons">
+      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>
+     </property>
+    </widget>
+   </item>
+   <item row="2" column="0">
+    <spacer name="verticalSpacer">
+     <property name="orientation">
+      <enum>Qt::Vertical</enum>
+     </property>
+     <property name="sizeHint" stdset="0">
+      <size>
+       <width>20</width>
+       <height>40</height>
+      </size>
+     </property>
+    </spacer>
+   </item>
   </layout>
  </widget>
  <resources/>

diff --git a/scoresight.spec b/scoresight.spec
@@ -71,13 +71,15 @@ sources = [
     'storage.py',
     'tesseract.py',
     'text_detection_target.py',
+    'training_dojo.py',
     'video_settings.py',
     'ui_about.py',
     'ui_connect_obs.py',
     'ui_log_view.py',
     'ui_mainwindow.py',
     'ui_ocr_training_data_dialog.py',
     'ui_screen_capture.py',
+    'ui_training_dojo.py',
     'ui_update_available.py',
     'ui_url_source.py',
     'ui_video_settings.py',