Fully Customizable GPU Accelerated Identification Nodes

This commit is contained in:
Nicole Rappe 2025-02-26 02:23:38 -07:00
parent ce392d7a04
commit 5c23653d59
6 changed files with 36 additions and 24 deletions

View File

@ -45,10 +45,10 @@ def _ensure_qapplication():
app_instance = QApplication(sys.argv) # Start in main thread
def create_ocr_region(region_id, x=250, y=50, w=DEFAULT_WIDTH, h=DEFAULT_HEIGHT, color=(255, 255, 0)):
def create_ocr_region(region_id, x=250, y=50, w=DEFAULT_WIDTH, h=DEFAULT_HEIGHT, color=(255, 255, 0), thickness=2):
"""
Creates an OCR region with a visible, resizable box on the screen.
The color parameter allows customization (default yellow, blue for overlays).
Allows setting custom color (RGB) and line thickness.
"""
_ensure_qapplication()
@ -60,7 +60,7 @@ def create_ocr_region(region_id, x=250, y=50, w=DEFAULT_WIDTH, h=DEFAULT_HEIGHT,
regions[region_id] = {
'bbox': [x, y, w, h],
'raw_text': "",
'widget': OCRRegionWidget(x, y, w, h, region_id, color)
'widget': OCRRegionWidget(x, y, w, h, region_id, color, thickness)
}
collector_mutex.unlock()
@ -179,19 +179,19 @@ def find_word_positions(region_id, word, offset_x=0, offset_y=0, margin=5, ocr_e
return []
def draw_identification_boxes(region_id, positions, color=(0, 0, 255)):
def draw_identification_boxes(region_id, positions, color=(0, 0, 255), thickness=2):
"""
Draws non-interactive rectangles at specified positions within the given OCR region.
"""
collector_mutex.lock()
if region_id in regions and 'widget' in regions[region_id]:
widget = regions[region_id]['widget']
widget.set_draw_positions(positions, color)
widget.set_draw_positions(positions, color, thickness)
collector_mutex.unlock()
class OCRRegionWidget(QWidget):
def __init__(self, x, y, w, h, region_id, color):
def __init__(self, x, y, w, h, region_id, color, thickness):
super().__init__()
self.setGeometry(x, y, w, h)
@ -203,6 +203,7 @@ class OCRRegionWidget(QWidget):
self.selected_handle = None
self.region_id = region_id
self.box_color = QColor(*color)
self.line_thickness = thickness
self.draw_positions = []
self.show()
@ -210,31 +211,23 @@ class OCRRegionWidget(QWidget):
def paintEvent(self, event):
painter = QPainter(self)
pen = QPen(self.box_color)
pen.setWidth(5)
pen.setWidth(self.line_thickness)
painter.setPen(pen)
# Draw main rectangle
painter.drawRect(0, 0, self.width(), self.height())
# Draw detected word overlays
pen.setWidth(2)
pen.setColor(QColor(0, 0, 255))
painter.setPen(pen)
for x, y, w, h in self.draw_positions:
painter.drawRect(x, y, w, h)
# Draw resize handles
painter.setBrush(self.box_color)
for handle in self._resize_handles():
painter.drawRect(handle)
def set_draw_positions(self, positions, color):
def set_draw_positions(self, positions, color, thickness):
"""
Update the positions where identification boxes should be drawn.
Updates the overlay positions and visual settings.
"""
self.draw_positions = positions
self.box_color = QColor(*color)
self.line_thickness = thickness
self.update()
def _resize_handles(self):

View File

@ -29,7 +29,7 @@ class FlyffCharacterStatusNode(BaseNode):
self.add_text_input("exp", "EXP", text="EXP: 0%")
self.region_id = "character_status"
data_collector.create_ocr_region(self.region_id, x=250, y=50, w=180, h=130)
data_collector.create_ocr_region(self.region_id, x=250, y=50, w=180, h=130, color=(255, 255, 0), thickness=2)
data_collector.start_collector()
self.set_name("Flyff - Character Status")

View File

@ -3,12 +3,13 @@
Identification Overlay Node:
- Creates an OCR region in data_collector with a blue overlay.
- Detects instances of a specified word and draws adjustable overlays.
- Users can configure offset, margin, polling frequency, and select OCR engine.
- Users can configure offset, margin, polling frequency, overlay color, and thickness.
"""
import re
from OdenGraphQt import BaseNode
from PyQt5.QtCore import QTimer
from PyQt5.QtGui import QColor
from Modules import data_collector
@ -25,10 +26,14 @@ class IdentificationOverlayNode(BaseNode):
self.add_text_input("margin", "Margin", text="5") # Box Margin
self.add_text_input("polling_freq", "Polling Frequency (ms)", text="500") # Polling Rate
self.add_combo_menu("ocr_engine", "Type", items=["CPU", "GPU"])
self.set_property("ocr_engine", "CPU") # Set default value after adding the menu
self.set_property("ocr_engine", "CPU") # Default to CPU mode
# Custom overlay options
self.add_text_input("overlay_color", "Overlay Color (RGB)", text="0,0,255") # Default blue
self.add_text_input("thickness", "Line Thickness", text="2") # Default 2px
self.region_id = "identification_overlay"
data_collector.create_ocr_region(self.region_id, x=250, y=50, w=300, h=200, color=(0, 0, 255))
data_collector.create_ocr_region(self.region_id, x=250, y=50, w=300, h=200, color=(0, 0, 255), thickness=2)
data_collector.start_collector()
self.set_name("Identification Overlay")
@ -76,6 +81,20 @@ class IdentificationOverlayNode(BaseNode):
except ValueError:
margin = 5 # Default margin if invalid input
# Parse overlay color
color_text = self.get_property("overlay_color")
try:
color = tuple(map(int, color_text.split(","))) # Convert "255,0,0" -> (255,0,0)
except ValueError:
color = (0, 0, 255) # Default to blue if invalid input
# Parse thickness
thickness_text = self.get_property("thickness")
try:
thickness = max(1, int(thickness_text)) # Ensure at least 1px thickness
except ValueError:
thickness = 2 # Default thickness
if not search_term:
return
@ -84,5 +103,5 @@ class IdentificationOverlayNode(BaseNode):
self.region_id, search_term, offset_x, offset_y, margin, ocr_engine
)
# Draw detected word boxes
data_collector.draw_identification_boxes(self.region_id, detected_positions, color=(0, 0, 255))
# Draw detected word boxes with custom color & thickness
data_collector.draw_identification_boxes(self.region_id, detected_positions, color=color, thickness=thickness)