diff --git a/src/resizable_rect.py b/src/resizable_rect.py index 37afa4a..cee3031 100644 --- a/src/resizable_rect.py +++ b/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,8 +445,7 @@ 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(), @@ -435,7 +453,18 @@ def mouseReleaseEvent(self, event): 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) diff --git a/src/source_view.py b/src/source_view.py index 344076f..271570b 100644 --- a/src/source_view.py +++ b/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): diff --git a/src/storage.py b/src/storage.py index 617c49d..34ab8c8 100644 --- a/src/storage.py +++ b/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 diff --git a/src/text_detection_target.py b/src/text_detection_target.py index b36a0b0..0afb016 100644 --- a/src/text_detection_target.py +++ b/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):