pyside与百度飞桨写一个截图提取文字的功能

由于前两篇文章已分开讲这截图与文字提取,这就直接上代码,里面有很多弹窗验证的代码,都可以整理删掉

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())
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值