From 3dcdde647c6584c1e42d504843ca60ab634a4d00 Mon Sep 17 00:00:00 2001 From: Roy Shilkrot Date: Sat, 19 Oct 2024 23:23:21 -0400 Subject: [PATCH] Refactor: Update ResizableRect and ResizableRectWithNameTypeAndResult classes This commit refactors the ResizableRect and ResizableRectWithNameTypeAndResult classes in the src/resizable_rect.py file. - Import statements have been updated to include the TextDetectionTarget and TextDetectionTargetWithResult classes from the text_detection_target module. - The setFlags method in the MiniRect class has been updated to use the QGraphicsItem.GraphicsItemFlag enum for improved readability. - The ResizableRectWithNameTypeAndResult class now accepts a TextDetectionTarget object as a parameter instead of individual x, y, width, height, and name parameters. - The setupAddButton method in the ResizableRectWithNameTypeAndResult class has been modified to use a smaller add button size and display a "+" symbol instead of "Add". - The mini_rects attribute in the ResizableRectWithNameTypeAndResult class is now initialized with MiniRect objects created from the mini_rects attribute of the TextDetectionTarget object. These changes improve the readability and maintainability of the code. --- src/resizable_rect.py | 36 ++++++++++++++++++++++-------------- src/source_view.py | 5 +++-- src/storage.py | 25 +++++++++++++++++++++++-- src/text_detection_target.py | 12 +++++++++++- 4 files changed, 59 insertions(+), 19 deletions(-) diff --git a/src/resizable_rect.py b/src/resizable_rect.py index 37afa4a..5057618 100644 --- a/src/resizable_rect.py +++ b/src/resizable_rect.py @@ -6,7 +6,7 @@ QGraphicsSimpleTextItem, ) -from text_detection_target import TextDetectionTargetWithResult +from text_detection_target import TextDetectionTarget, TextDetectionTargetWithResult class ResizableRect(QGraphicsRectItem): @@ -156,7 +156,10 @@ def __init__(self, x, y, width, height, parent=None): 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 + ) def mousePressEvent(self, event): super().mousePressEvent(event) @@ -170,11 +173,7 @@ def mouseReleaseEvent(self, event): class ResizableRectWithNameTypeAndResult(ResizableRect): def __init__( self, - x, - y, - width, - height, - name, + detectionTarget: TextDetectionTarget, image_size, result="", onCenter=False, @@ -182,10 +181,16 @@ def __init__( 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 +203,10 @@ 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(), parent=self) + for r in detectionTarget.mini_rects + ] self.mini_rect_mode = False self.add_button = None self.setupAddButton() @@ -210,15 +218,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)) @@ -435,7 +443,7 @@ def mouseReleaseEvent(self, event): origRect.width(), origRect.height(), ) - self.boxChangedCallback(self.name, boxRect) + self.boxChangedCallback(self.name, boxRect, self.mini_rects) def mousePressEvent(self, event): if self.mini_rect_mode: diff --git a/src/source_view.py b/src/source_view.py index 344076f..615fd46 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, @@ -151,7 +151,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 +165,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):