104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Identification Overlay Node:
|
|
- Users can configure threads/slices for parallel processing.
|
|
"""
|
|
|
|
import re
|
|
from OdenGraphQt import BaseNode
|
|
from PyQt5.QtCore import QTimer
|
|
from PyQt5.QtGui import QColor
|
|
from Modules import data_collector
|
|
|
|
|
|
class IdentificationOverlayNode(BaseNode):
|
|
__identifier__ = "bunny-lab.io.identification_overlay_node"
|
|
NODE_NAME = "Identification Overlay"
|
|
|
|
def __init__(self):
|
|
super(IdentificationOverlayNode, self).__init__()
|
|
|
|
# User-configurable options
|
|
self.add_text_input("search_term", "Search Term", text="Aibatt")
|
|
self.add_text_input("offset_value", "Offset Value (X,Y)", text="0,0") # X,Y Offset
|
|
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") # 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.add_text_input("threads_slices", "Threads / Slices", text="8") # Default 8 threads/slices
|
|
|
|
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), thickness=2)
|
|
|
|
data_collector.start_collector()
|
|
self.set_name("Identification Overlay")
|
|
|
|
# Timer for updating overlays
|
|
self.timer = QTimer()
|
|
self.timer.timeout.connect(self.update_overlay)
|
|
|
|
# Set initial polling frequency
|
|
self.update_polling_frequency()
|
|
|
|
def update_polling_frequency(self):
|
|
polling_text = self.get_property("polling_freq")
|
|
try:
|
|
polling_interval = max(50, int(polling_text))
|
|
except ValueError:
|
|
polling_interval = 500
|
|
|
|
self.timer.start(polling_interval)
|
|
|
|
def update_overlay(self):
|
|
search_term = self.get_property("search_term")
|
|
offset_text = self.get_property("offset_value")
|
|
margin_text = self.get_property("margin")
|
|
ocr_engine = self.get_property("ocr_engine")
|
|
threads_slices_text = self.get_property("threads_slices")
|
|
|
|
self.update_polling_frequency()
|
|
|
|
try:
|
|
offset_x, offset_y = map(int, offset_text.split(","))
|
|
except ValueError:
|
|
offset_x, offset_y = 0, 0
|
|
|
|
try:
|
|
margin = int(margin_text)
|
|
except ValueError:
|
|
margin = 5
|
|
|
|
color_text = self.get_property("overlay_color")
|
|
try:
|
|
color = tuple(map(int, color_text.split(",")))
|
|
except ValueError:
|
|
color = (0, 0, 255)
|
|
|
|
thickness_text = self.get_property("thickness")
|
|
try:
|
|
thickness = max(1, int(thickness_text))
|
|
except ValueError:
|
|
thickness = 2
|
|
|
|
try:
|
|
num_slices = max(1, int(threads_slices_text)) # Ensure at least 1 slice
|
|
except ValueError:
|
|
num_slices = 1
|
|
|
|
if not search_term:
|
|
return
|
|
|
|
detected_positions = data_collector.find_word_positions(
|
|
self.region_id, search_term, offset_x, offset_y, margin, ocr_engine, num_slices
|
|
)
|
|
|
|
# Ensure slice count is updated visually in the region widget
|
|
data_collector.update_region_slices(self.region_id, num_slices)
|
|
|
|
data_collector.draw_identification_boxes(self.region_id, detected_positions, color=color, thickness=thickness)
|
|
|