ocr_engine.py

"""
OCR识别引擎
负责图片预处理和文字识别
"""
import os
from paddleocr import PaddleOCR
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np


class OCREngine:
    def __init__(self, lang='ch'):
        """
        初始化OCR引擎
        :param lang: 识别语言，'ch'为中文，'en'为英文
        """
        # 新版 paddleocr（3.x）已移除 show_log / use_gpu 参数
        self.ocr = PaddleOCR(
            lang=lang
        )

    def preprocess_image(self, image_path):
        """
        图片预处理：提高OCR识别准确率
        """
        img = Image.open(image_path)

        if img.mode != 'L':
            img = img.convert('L')

        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(2.0)

        threshold = 140
        img = img.point(lambda p: 255 if p > threshold else 0)

        img = img.filter(ImageFilter.MedianFilter(size=3))

        return img

    def recognize(self, image_path, preprocess=True):
        """
        识别图片中的文字
        :param image_path: 图片路径
        :param preprocess: 是否预处理
        :return: 识别结果列表
        """
        if preprocess:
            img = self.preprocess_image(image_path)
            temp_path = image_path + '.temp.jpg'
            img.save(temp_path)
            result = self.ocr.predict(temp_path)
            os.remove(temp_path)
        else:
            result = self.ocr.predict(image_path)

        texts = []
        for res in result:
            rec_texts = res.get('rec_texts', [])
            rec_scores = res.get('rec_scores', [])
            rec_polys = res.get('rec_polys', [])

            for i, text in enumerate(rec_texts):
                confidence = rec_scores[i] if i < len(rec_scores) else 0.0
                position = rec_polys[i].tolist() if i < len(rec_polys) else []
                texts.append({
                    'text': text,
                    'confidence': confidence,
                    'position': position
                })

        return texts

    def recognize_batch(self, image_paths):
        """
        批量识别
        :param image_paths: 图片路径列表
        :return: {图片路径: 识别结果}
        """
        results = {}
        for path in image_paths:
            try:
                results[path] = self.recognize(path)
            except Exception as e:
                results[path] = {'error': str(e)}
        return results