由于前两篇文章已分开讲这截图与文字提取,这就直接上代码,里面有很多弹窗验证的代码,都可以整理删掉
import sys
import os
import tempfile
import shutil
from PySide6.QtWidgets import (QApplication, QMainWindow, QLabel,
QPushButton, QVBoxLayout, QWidget,
QFileDialog, QMessageBox)
from PySide6.QtUiTools import QUiLoader
from PySide6.QtGui import (QPixmap, QPainter, QPen, QColor, QClipboard,
QKeySequence, QShortcut)
from PySide6.QtCore import (Qt, QRect, QPoint, Signal, QFile,
QThread, QObject)
# 导入飞桨OCR
try:
from paddleocr import PaddleOCR
except ImportError:
QMessageBox.critical(None, "依赖缺失", "请安装paddleocr: pip install paddleocr")
sys.exit(1)
# OCR处理子线程类
class OCRWorker(QObject):
finished = Signal(str)
error = Signal(str)
def __init__(self, ocr_model, temp_image_path):
super().__init__()
self.ocr_model = ocr_model
self.temp_image_path = temp_image_path
def run(self):
try:
result = self.ocr_model.predict(self.temp_image_path)
extracted_text = []
for line in result:
line_texts = line['rec_texts']
extracted_text.extend(line_texts)
full_text = "\n".join(extracted_text) if extracted_text else ""
self.finished.emit(full_text)
except Exception as e:
self.error.emit(f"OCR识别失败: {str(e)}")
# 截图覆盖窗口类
class OverlayWindow(QWidget):
screenshot_taken = Signal(QPixmap)
cancelled = Signal()
def __init__(self, background):
super().__init__()
self.background = background
self.start_point = QPoint()
self.end_point = QPoint()
self.dragging = False
self.setWindowFlags(
Qt.WindowStaysOnTopHint | Qt.FramelessWindowHint | Qt.X11BypassWindowManagerHint
)
self.setAttribute(Qt.WA_TranslucentBackground)
self.setCursor(Qt.CrossCursor)
def paintEvent(self, event):
painter = QPainter(self)
painter.drawPixmap(0, 0, self.background)
if self.dragging:
painter.setBrush(QColor(0, 0, 0, 120))
painter.setPen(Qt.NoPen)
painter.drawRect(self.rect())
rect = QRect(self.start_point, self.end_point).normalized()
painter.setBrush(Qt.NoBrush)
painter.setPen(QPen(Qt.red, 2))
painter.drawRect(rect)
painter.drawPixmap(rect, self.background, rect)
def mousePressEvent(self, event):
if event.button() == Qt.LeftButton:
self.start_point = event.position().toPoint()
self.end_point = self.start_point
self.dragging = True
def mouseMoveEvent(self, event):
if self.dragging:
self.end_point = event.position().toPoint()
self.update()
def mouseReleaseEvent(self, event):
if event.button() == Qt.LeftButton and self.dragging:
self.dragging = False
rect = QRect(self.start_point, self.end_point).normalized()
if rect.width() > 10 and rect.height() > 10:
screenshot = self.background.copy(rect)
self.screenshot_taken.emit(screenshot)
else:
self.cancelled.emit()
self.close()
def keyPressEvent(self, event):
if event.key() == Qt.Key_Escape:
self.cancelled.emit()
self.close()
# 主窗口类
class ScreenshotOCRTool(QMainWindow):
def __init__(self):
super().__init__()
self.temp_dir = tempfile.mkdtemp()
self.ocr = None
self.ocr_thread = None # 存储OCR子线程
self.ocr_worker = None # 存储OCR工作对象
self.init_ui()
self.init_ocr()
self.init_global_shortcut()
def init_ui(self):
try:
loader = QUiLoader()
self.ui.pushButton_jietushizi.clicked.connect(self.start_capture)
self.setCentralWidget(self.ui)
self.setWindowTitle("截图识字工具(全局快捷键:Ctrl+1)")
except Exception as e:
self.create_fallback_ui()
def create_fallback_ui(self):
central_widget = QWidget()
self.setCentralWidget(central_widget)
layout = QVBoxLayout(central_widget)
self.status_label = QLabel("OCR模型初始化中...(全局快捷键:Ctrl+1)")
self.status_label.setAlignment(Qt.AlignCenter)
self.status_label.setStyleSheet("font-size: 14px; margin: 20px 0;")
self.pushButton_jietushizi = QPushButton("截图识别文字")
self.pushButton_jietushizi.setStyleSheet("padding: 10px; font-size: 14px;")
self.pushButton_jietushizi.clicked.connect(self.start_capture)
layout.addWidget(self.status_label)
layout.addWidget(self.pushButton_jietushizi)
self.resize(380, 180)
self.setWindowTitle("截图识字工具(全局快捷键:Ctrl+1)")
def init_ocr(self):
try:
if hasattr(self, 'status_label'):
self.status_label.setText("正在初始化OCR模型...(全局快捷键:Ctrl+Alt+O)")
QApplication.processEvents()
self.ocr = PaddleOCR(use_textline_orientation=True, lang="ch")
if hasattr(self, 'status_label'):
self.status_label.setText("OCR就绪,点击按钮或按Ctrl+1开始截图")
except Exception as e:
err_msg = f"OCR初始化失败: {str(e)}"
QMessageBox.critical(self, "错误", err_msg)
if hasattr(self, 'status_label'):
self.status_label.setText("OCR初始化失败,请重启程序(全局快捷键:Ctrl+1)")
self.ocr = None
def init_global_shortcut(self):
self.shortcut = QShortcut(QKeySequence("Ctrl+1"), self)
self.shortcut.activated.connect(self.start_capture)
def start_capture(self):
if not self.ocr:
QMessageBox.warning(self, "OCR未就绪", "OCR模型未初始化,无法识别")
return
self.hide()
QApplication.processEvents()
screen = QApplication.primaryScreen()
full_screenshot = screen.grabWindow(0)
self.overlay = OverlayWindow(full_screenshot)
self.overlay.screenshot_taken.connect(self.process_screenshot)
self.overlay.cancelled.connect(self.show)
self.overlay.showFullScreen()
def process_screenshot(self, screenshot):
try:
temp_image_path = os.path.join(self.temp_dir, "temp_screenshot.png")
screenshot.save(temp_image_path)
self.show()
QApplication.processEvents()
# 确保在创建新线程前重置旧线程引用
self.ocr_thread = None
self.ocr_worker = None
self.ocr_thread = QThread()
self.ocr_worker = OCRWorker(self.ocr, temp_image_path)
self.ocr_worker.moveToThread(self.ocr_thread)
self.ocr_thread.started.connect(self.ocr_worker.run)
self.ocr_worker.finished.connect(self.on_ocr_success)
self.ocr_worker.error.connect(self.on_ocr_error)
self.ocr_worker.finished.connect(self.ocr_thread.quit)
self.ocr_worker.finished.connect(self.ocr_worker.deleteLater)
self.ocr_thread.finished.connect(self.ocr_thread.deleteLater)
# 线程结束后重置引用
self.ocr_thread.finished.connect(self.reset_thread_reference)
self.ocr_thread.start()
except Exception as e:
self.show()
QMessageBox.critical(self, "处理失败", f"截图保存失败: {str(e)}")
def reset_thread_reference(self):
"""线程结束后重置引用,避免访问已删除对象"""
self.ocr_thread = None
self.ocr_worker = None
def on_ocr_success(self, full_text):
if full_text:
clipboard = QApplication.clipboard()
clipboard.setText(full_text)
QMessageBox.information(self, "识别结果", full_text)
def on_ocr_error(self, err_msg):
QMessageBox.critical(self, "识别错误", err_msg)
def closeEvent(self, event):
"""安全关闭线程,检查对象是否存在"""
# 先检查线程对象是否存在且有效
if self.ocr_thread is not None:
try:
if self.ocr_thread.isRunning():
self.ocr_thread.quit()
self.ocr_thread.wait()
except RuntimeError:
# 线程对象已被删除,直接忽略
pass
# 清理临时文件
if os.path.exists(self.temp_dir):
try:
shutil.rmtree(self.temp_dir)
except Exception as e:
print(f"临时文件清理失败: {e}")
event.accept()
if __name__ == "__main__":
app = QApplication(sys.argv)
window = ScreenshotOCRTool()
window.show()
sys.exit(app.exec())
598

被折叠的 条评论
为什么被折叠?



