292 lines
9.3 KiB
Python
292 lines
9.3 KiB
Python
# Modules/data_collector.py
|
|
|
|
import threading
|
|
import time
|
|
import re
|
|
import sys
|
|
import numpy as np
|
|
import cv2
|
|
|
|
# Vision-related Imports
|
|
import pytesseract
|
|
import easyocr
|
|
import torch
|
|
|
|
from PIL import Image, ImageGrab, ImageFilter
|
|
|
|
from PyQt5.QtWidgets import QApplication, QWidget
|
|
from PyQt5.QtCore import QRect, QPoint, Qt, QMutex, QTimer
|
|
from PyQt5.QtGui import QPainter, QPen, QColor, QFont
|
|
|
|
# Initialize EasyOCR with CUDA support
|
|
reader_cpu = easyocr.Reader(['en'], gpu=False)
|
|
reader_gpu = easyocr.Reader(['en'], gpu=True if torch.cuda.is_available() else False)
|
|
|
|
pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
|
|
|
|
DEFAULT_WIDTH = 180
|
|
DEFAULT_HEIGHT = 130
|
|
HANDLE_SIZE = 8
|
|
LABEL_HEIGHT = 20
|
|
|
|
collector_mutex = QMutex()
|
|
regions = {}
|
|
|
|
app_instance = None
|
|
|
|
|
|
def _ensure_qapplication():
|
|
"""
|
|
Ensures that QApplication is initialized before creating widgets.
|
|
Must be called from the main thread.
|
|
"""
|
|
global app_instance
|
|
if app_instance is None:
|
|
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)):
|
|
"""
|
|
Creates an OCR region with a visible, resizable box on the screen.
|
|
The color parameter allows customization (default yellow, blue for overlays).
|
|
"""
|
|
|
|
_ensure_qapplication()
|
|
|
|
collector_mutex.lock()
|
|
if region_id in regions:
|
|
collector_mutex.unlock()
|
|
return
|
|
regions[region_id] = {
|
|
'bbox': [x, y, w, h],
|
|
'raw_text': "",
|
|
'widget': OCRRegionWidget(x, y, w, h, region_id, color)
|
|
}
|
|
collector_mutex.unlock()
|
|
|
|
|
|
def get_raw_text(region_id):
|
|
collector_mutex.lock()
|
|
if region_id not in regions:
|
|
collector_mutex.unlock()
|
|
return ""
|
|
text = regions[region_id]['raw_text']
|
|
collector_mutex.unlock()
|
|
return text
|
|
|
|
|
|
def start_collector():
|
|
t = threading.Thread(target=_update_ocr_loop, daemon=True)
|
|
t.start()
|
|
|
|
|
|
def _update_ocr_loop():
|
|
while True:
|
|
collector_mutex.lock()
|
|
region_ids = list(regions.keys())
|
|
collector_mutex.unlock()
|
|
|
|
for rid in region_ids:
|
|
collector_mutex.lock()
|
|
bbox = regions[rid]['bbox'][:]
|
|
collector_mutex.unlock()
|
|
|
|
x, y, w, h = bbox
|
|
screenshot = ImageGrab.grab(bbox=(x, y, x + w, y + h))
|
|
processed = _preprocess_image(screenshot)
|
|
raw_text = pytesseract.image_to_string(processed, config='--psm 6 --oem 1')
|
|
|
|
collector_mutex.lock()
|
|
if rid in regions:
|
|
regions[rid]['raw_text'] = raw_text
|
|
collector_mutex.unlock()
|
|
|
|
time.sleep(0.7)
|
|
|
|
|
|
def _preprocess_image(image):
|
|
gray = image.convert("L")
|
|
scaled = gray.resize((gray.width * 3, gray.height * 3))
|
|
thresh = scaled.point(lambda p: 255 if p > 200 else 0)
|
|
return thresh.filter(ImageFilter.MedianFilter(3))
|
|
|
|
|
|
def find_word_positions(region_id, word, offset_x=0, offset_y=0, margin=5, ocr_engine="CPU"):
|
|
"""
|
|
Finds positions of a specific word within the OCR region.
|
|
Applies user-defined offset and margin adjustments.
|
|
Uses Tesseract (CPU) or EasyOCR (GPU) depending on the selected engine.
|
|
"""
|
|
collector_mutex.lock()
|
|
if region_id not in regions:
|
|
collector_mutex.unlock()
|
|
return []
|
|
|
|
bbox = regions[region_id]['bbox']
|
|
collector_mutex.unlock()
|
|
|
|
# Extract OCR region position and size
|
|
x, y, w, h = bbox
|
|
left, top, right, bottom = x, y, x + w, y + h
|
|
|
|
if right <= left or bottom <= top:
|
|
print(f"[ERROR] Invalid OCR region bounds: {bbox}")
|
|
return []
|
|
|
|
try:
|
|
image = ImageGrab.grab(bbox=(left, top, right, bottom))
|
|
processed = _preprocess_image(image)
|
|
|
|
# Get original and processed image sizes
|
|
orig_width, orig_height = image.size
|
|
proc_width, proc_height = processed.size
|
|
|
|
# Scale factor between processed image and original screenshot
|
|
scale_x = orig_width / proc_width
|
|
scale_y = orig_height / proc_height
|
|
|
|
word_positions = []
|
|
|
|
if ocr_engine == "CPU":
|
|
# Use Tesseract (CPU)
|
|
data = pytesseract.image_to_data(processed, config='--psm 6 --oem 1', output_type=pytesseract.Output.DICT)
|
|
|
|
for i in range(len(data['text'])):
|
|
if re.search(rf"\b{word}\b", data['text'][i], re.IGNORECASE):
|
|
x_scaled = int(data['left'][i] * scale_x)
|
|
y_scaled = int(data['top'][i] * scale_y)
|
|
w_scaled = int(data['width'][i] * scale_x)
|
|
h_scaled = int(data['height'][i] * scale_y)
|
|
|
|
word_positions.append((x_scaled + offset_x, y_scaled + offset_y, w_scaled + (margin * 2), h_scaled + (margin * 2)))
|
|
|
|
else:
|
|
# Use EasyOCR (GPU) - Convert PIL image to NumPy array
|
|
image_np = np.array(processed)
|
|
results = reader_gpu.readtext(image_np)
|
|
|
|
for (bbox, text, _) in results:
|
|
if re.search(rf"\b{word}\b", text, re.IGNORECASE):
|
|
(x_min, y_min), (x_max, y_max) = bbox[0], bbox[2]
|
|
|
|
x_scaled = int(x_min * scale_x)
|
|
y_scaled = int(y_min * scale_y)
|
|
w_scaled = int((x_max - x_min) * scale_x)
|
|
h_scaled = int((y_max - y_min) * scale_y)
|
|
|
|
word_positions.append((x_scaled + offset_x, y_scaled + offset_y, w_scaled + (margin * 2), h_scaled + (margin * 2)))
|
|
|
|
return word_positions
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to capture OCR region: {e}")
|
|
return []
|
|
|
|
|
|
def draw_identification_boxes(region_id, positions, color=(0, 0, 255)):
|
|
"""
|
|
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)
|
|
collector_mutex.unlock()
|
|
|
|
|
|
class OCRRegionWidget(QWidget):
|
|
def __init__(self, x, y, w, h, region_id, color):
|
|
super().__init__()
|
|
|
|
self.setGeometry(x, y, w, h)
|
|
self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
|
|
self.setAttribute(Qt.WA_TranslucentBackground, True)
|
|
self.setAttribute(Qt.WA_TransparentForMouseEvents, False)
|
|
|
|
self.drag_offset = None
|
|
self.selected_handle = None
|
|
self.region_id = region_id
|
|
self.box_color = QColor(*color)
|
|
self.draw_positions = []
|
|
|
|
self.show()
|
|
|
|
def paintEvent(self, event):
|
|
painter = QPainter(self)
|
|
pen = QPen(self.box_color)
|
|
pen.setWidth(5)
|
|
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):
|
|
"""
|
|
Update the positions where identification boxes should be drawn.
|
|
"""
|
|
self.draw_positions = positions
|
|
self.box_color = QColor(*color)
|
|
self.update()
|
|
|
|
def _resize_handles(self):
|
|
w, h = self.width(), self.height()
|
|
return [
|
|
QRect(0, 0, HANDLE_SIZE, HANDLE_SIZE), # Top-left
|
|
QRect(w - HANDLE_SIZE, h - HANDLE_SIZE, HANDLE_SIZE, HANDLE_SIZE) # Bottom-right
|
|
]
|
|
|
|
def mousePressEvent(self, event):
|
|
if event.button() == Qt.LeftButton:
|
|
for i, handle in enumerate(self._resize_handles()):
|
|
if handle.contains(event.pos()):
|
|
self.selected_handle = i
|
|
return
|
|
|
|
self.drag_offset = event.pos()
|
|
|
|
def mouseMoveEvent(self, event):
|
|
if self.selected_handle is not None:
|
|
w, h = self.width(), self.height()
|
|
if self.selected_handle == 0: # Top-left
|
|
new_w = w + (self.x() - event.globalX())
|
|
new_h = h + (self.y() - event.globalY())
|
|
new_x = event.globalX()
|
|
new_y = event.globalY()
|
|
if new_w < 20: new_w = 20
|
|
if new_h < 20: new_h = 20
|
|
self.setGeometry(new_x, new_y, new_w, new_h)
|
|
elif self.selected_handle == 1: # Bottom-right
|
|
new_w = event.globalX() - self.x()
|
|
new_h = event.globalY() - self.y()
|
|
if new_w < 20: new_w = 20
|
|
if new_h < 20: new_h = 20
|
|
self.setGeometry(self.x(), self.y(), new_w, new_h)
|
|
|
|
collector_mutex.lock()
|
|
if self.region_id in regions:
|
|
regions[self.region_id]['bbox'] = [self.x(), self.y(), self.width(), self.height()]
|
|
collector_mutex.unlock()
|
|
|
|
self.update()
|
|
elif self.drag_offset:
|
|
new_x = event.globalX() - self.drag_offset.x()
|
|
new_y = event.globalY() - self.drag_offset.y()
|
|
self.move(new_x, new_y)
|
|
|
|
collector_mutex.lock()
|
|
if self.region_id in regions:
|
|
regions[self.region_id]['bbox'] = [new_x, new_y, self.width(), self.height()]
|
|
collector_mutex.unlock()
|