Merge pull request #48 from locaal-ai/roy.mini_rect_storage

Refactor: Update ResizableRect and ResizableRectWithNameTypeAndResult…

src/resizable_rect.py

@@ -1,3 +1,4 @@
+from typing import Callable
 from PySide6.QtCore import QPointF, QRectF, Qt
 from PySide6.QtGui import QBrush, QColor, QFont, QPen
 from PySide6.QtWidgets import (
@@ -6,7 +7,7 @@
     QGraphicsSimpleTextItem,
 )
 
-from text_detection_target import TextDetectionTargetWithResult
+from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult
 
 
 class ResizableRect(QGraphicsRectItem):
@@ -151,12 +152,16 @@ def hoverMoveEvent(self, event):
 
 
 class MiniRect(ResizableRect):
-    def __init__(self, x, y, width, height, parent=None):
+    def __init__(self, x, y, width, height, boxChangedCallback, parent=None):
         super().__init__(x, y, width, height)
         self.setPen(QPen(QColor(255, 0, 0)))
         self.setBrush(QBrush(QColor(255, 0, 0, 50)))
         self.setParentItem(parent)
-        self.setFlags(QGraphicsItem.ItemIsMovable | QGraphicsItem.ItemIsSelectable)
+        self.setFlags(
+            QGraphicsItem.GraphicsItemFlag.ItemIsMovable
+            | QGraphicsItem.GraphicsItemFlag.ItemIsSelectable
+        )
+        self.boxChangedCallback = boxChangedCallback
 
     def mousePressEvent(self, event):
         super().mousePressEvent(event)
@@ -165,27 +170,30 @@ def mousePressEvent(self, event):
     def mouseReleaseEvent(self, event):
         super().mouseReleaseEvent(event)
         self.unsetCursor()
+        self.boxChangedCallback()
 
 
 class ResizableRectWithNameTypeAndResult(ResizableRect):
     def __init__(
         self,
-        x,
-        y,
-        width,
-        height,
-        name,
+        detectionTarget: TextDetectionTarget,
         image_size,
         result="",
         onCenter=False,
-        boxChangedCallback=None,
+        boxChangedCallback: Callable[[str, QRectF, list[QRectF]], None] = None,
         itemSelectedCallback=None,
         boxDisplayStyle: int = 1,
     ):
-        super().__init__(x, y, width, height, onCenter)
+        super().__init__(
+            detectionTarget.x(),
+            detectionTarget.y(),
+            detectionTarget.width(),
+            detectionTarget.height(),
+            onCenter,
+        )
         self.setAcceptedMouseButtons(Qt.MouseButton.LeftButton)
         self.setAcceptHoverEvents(True)
-        self.name = name
+        self.name = detectionTarget.name
         self.result = result
         self.boxChangedCallback = boxChangedCallback
         self.itemSelectedCallback = itemSelectedCallback
@@ -198,7 +206,17 @@ def __init__(
         self.bgItem = QGraphicsRectItem(self.posItem.boundingRect(), parent=self)
 
         # Mini-rect related attributes
-        self.mini_rects = []
+        self.mini_rects = [
+            MiniRect(
+                r.x(),
+                r.y(),
+                r.width(),
+                r.height(),
+                self.sendBoxChangedCallback,
+                parent=self,
+            )
+            for r in detectionTarget.mini_rects
+        ]
         self.mini_rect_mode = False
         self.add_button = None
         self.setupAddButton()
@@ -210,15 +228,15 @@ def __init__(
         self.updateCornerBoxes()
 
     def setupAddButton(self):
-        self.add_button = QGraphicsRectItem(0, 0, 60, 30, parent=self)
+        self.add_button = QGraphicsRectItem(0, 0, 25, 30, parent=self)
         self.add_button.setBrush(QBrush(QColor(0, 255, 0)))
         self.add_button.setPen(QPen(Qt.black))
-        self.add_button.setPos(self.rect().topLeft() + QPointF(5, 5))
+        self.add_button.setPos(self.rect().topLeft() + QPointF(2, 2))
         self.add_button.setZValue(4)
         self.add_button.setVisible(False)
 
         # Add a "+" text to the button
-        text = QGraphicsSimpleTextItem("Add", self.add_button)
+        text = QGraphicsSimpleTextItem("+", self.add_button)
         text.setPos(5, 0)
         text.setFont(QFont("Arial", 20))
 
@@ -414,6 +432,7 @@ def startCreateMiniRect(self, rect: QRectF):
             rect.y(),
             rect.width(),
             rect.height(),
+            self.sendBoxChangedCallback,
             parent=self,
         )
         self.mini_rects.append(new_mini_rect)
@@ -426,16 +445,26 @@ def clearMiniRects(self):
     def getMiniRects(self):
         return [rect.rect() for rect in self.mini_rects]
 
-    def mouseReleaseEvent(self, event):
-        super().mouseReleaseEvent(event)
+    def sendBoxChangedCallback(self):
         origRect = self.getRect()
         boxRect = QRectF(
             origRect.x() + self.x(),
             origRect.y() + self.y(),
             origRect.width(),
             origRect.height(),
         )
-        self.boxChangedCallback(self.name, boxRect)
+        self.boxChangedCallback(
+            self.name,
+            boxRect,
+            [
+                QRectF(r.x(), r.y(), r.rect().width(), r.rect().height())
+                for r in self.mini_rects
+            ],
+        )
+
+    def mouseReleaseEvent(self, event):
+        super().mouseReleaseEvent(event)
+        self.sendBoxChangedCallback()
 
     def mousePressEvent(self, event):
         if self.mini_rect_mode:
@@ -449,6 +478,9 @@ def mousePressEvent(self, event):
                     )
                 )
             else:
+                # deselect all mini rects
+                for mini_rect in self.mini_rects:
+                    mini_rect.setSelected(False)
                 super().mousePressEvent(event)
         else:
             super().mousePressEvent(event)

src/source_view.py

@@ -1,5 +1,5 @@
 import math
-from PySide6.QtCore import QPointF, Qt, QTimer
+from PySide6.QtCore import QPointF, Qt, QTimer, QRectF
 from PySide6.QtGui import QBrush, QColor, QMouseEvent, QPen, QPolygonF
 from PySide6.QtWidgets import (
     QGraphicsPolygonItem,
@@ -123,11 +123,7 @@ def detectionTargetsChanged(self):
             boxFound = self.findBox(detectionTarget.name)
             if boxFound is None:
                 boxFound = ResizableRectWithNameTypeAndResult(
-                    detectionTarget.x(),
-                    detectionTarget.y(),
-                    detectionTarget.width(),
-                    detectionTarget.height(),
-                    detectionTarget.name,
+                    detectionTarget,
                     # image size
                     self.scene.sceneRect().width(),
                     onCenter=False,
@@ -151,7 +147,7 @@ def detectionTargetsChanged(self):
                 if item.name not in done_targets:
                     self.scene.removeItem(item)
 
-    def boxChanged(self, name, rect):
+    def boxChanged(self, name: str, rect: QRectF, mini_rects: list[QRectF]):
         # update the detection target in the storage
         detectionTargets: list[TextDetectionTarget] = (
             self.detectionTargetsStorage.get_data()
@@ -165,6 +161,7 @@ def boxChanged(self, name, rect):
                 self.detectionTargetsStorage.edit_item(
                     detectionTarget.name, detectionTarget
                 )
+                detectionTarget.mini_rects = mini_rects
                 break
 
     def findBox(self, name):

src/storage.py

@@ -1,6 +1,6 @@
 import json
 import os
-from PySide6.QtCore import QObject, Signal
+from PySide6.QtCore import QObject, Signal, QRectF
 from platformdirs import user_data_dir
 from defaults import default_info_for_box_name, normalize_settings_dict
 
@@ -218,6 +218,17 @@ def loadBoxesFromDict(self, boxes) -> bool:
                     box["settings"] = {}
 
                 default_box_info = default_info_for_box_name(box["name"])
+                mini_rects: list[QRectF] = []
+                if "mini_rects" in box:
+                    for mini_rect in box["mini_rects"]:
+                        mini_rects.append(
+                            QRectF(
+                                mini_rect["x"],
+                                mini_rect["y"],
+                                mini_rect["width"],
+                                mini_rect["height"],
+                            )
+                        )
                 # set the position of the box
                 self._data.append(
                     TextDetectionTarget(
@@ -227,6 +238,7 @@ def loadBoxesFromDict(self, boxes) -> bool:
                         box["rect"]["height"],
                         box["name"],
                         normalize_settings_dict(box["settings"], default_box_info),
+                        mini_rects,
                     )
                 )
                 if "is_custom" in box["settings"] and box["settings"]["is_custom"]:
@@ -239,7 +251,7 @@ def loadBoxesFromDict(self, boxes) -> bool:
             return False
         return True
 
-    def getBoxesForStorage(self):
+    def getBoxesForStorage(self) -> list[dict]:
         # save all the boxes to scoresight.json
         boxes = []
         for detectionTarget in self._data:
@@ -296,6 +308,15 @@ def getBoxesForStorage(self):
                         ),
                         "composite_box": detectionTarget.settings.get("composite_box"),
                     },
+                    "mini_rects": [
+                        {
+                            "x": mini_rect.x(),
+                            "y": mini_rect.y(),
+                            "width": mini_rect.width(),
+                            "height": mini_rect.height(),
+                        }
+                        for mini_rect in detectionTarget.mini_rects
+                    ],
                 }
             )
         return boxes

src/text_detection_target.py

@@ -39,13 +39,23 @@ def clear(self):
 
 
 class TextDetectionTarget(QRectF):
-    def __init__(self, x, y, width, height, name: str, settings: dict = {}):
+    def __init__(
+        self,
+        x,
+        y,
+        width,
+        height,
+        name: str,
+        settings: dict = {},
+        mini_rects: list[QRectF] = [],
+    ):
         super().__init__(x, y, width, height)
         self.name = name
         self.settings = settings
         self.ocrResultPerCharacterSmoother = OCRResultPerCharacterSmoother()
         self.last_image = None
         self.last_text = None
+        self.mini_rects: list[QRectF] = mini_rects
 
 
 class TextDetectionTargetWithResult(TextDetectionTarget):