新增代码

difyPlugin/app/services/workcase/qrcode/QrCode.py

@@ -0,0 +1,298 @@
+# -*- coding: utf-8 -*-
+"""二维码处理核心类 - 基于 OpenCV QRCodeDetector
+
+本模块使用 OpenCV 的 QRCodeDetector 进行二维码识别，
+配合多种图像预处理策略，确保高识别率和跨平台兼容性。
+"""
+import base64
+import io
+from typing import Optional, Callable, Tuple
+
+import cv2
+import httpx
+import numpy as np
+import qrcode
+from PIL import Image
+
+
+class QRCodeProcessor:
+    """二维码处理器 - 负责二维码的生成、解析和图像预处理"""
+
+    # 预处理策略映射
+    PREPROCESSING_STRATEGIES = {
+        "original": ("原图", lambda img, gray: img),
+        "grayscale": ("灰度图", lambda img, gray: cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)),
+        "clahe": ("CLAHE增强", lambda img, gray: cv2.cvtColor(
+            cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(gray),
+            cv2.COLOR_GRAY2BGR
+        )),
+        "adaptive_threshold": ("自适应二值化", lambda img, gray: cv2.cvtColor(
+            cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2),
+            cv2.COLOR_GRAY2BGR
+        )),
+        "otsu": ("Otsu二值化", lambda img, gray: cv2.cvtColor(
+            cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1],
+            cv2.COLOR_GRAY2BGR
+        )),
+        "denoise": ("去噪+二值化", lambda img, gray: cv2.cvtColor(
+            cv2.threshold(
+                cv2.fastNlMeansDenoising(gray, None, 10, 7, 21),
+                0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
+            )[1],
+            cv2.COLOR_GRAY2BGR
+        )),
+        "sharpen": ("锐化", lambda img, gray: cv2.cvtColor(
+            cv2.filter2D(gray, -1, np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])),
+            cv2.COLOR_GRAY2BGR
+        )),
+        "morphology": ("形态学处理", lambda img, gray: cv2.cvtColor(
+            cv2.morphologyEx(
+                cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2),
+                cv2.MORPH_CLOSE,
+                cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+            ),
+            cv2.COLOR_GRAY2BGR
+        )),
+        "scale_0.5": ("0.5x缩放", lambda img, gray: cv2.resize(
+            img, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA
+        )),
+        "scale_1.5": ("1.5x缩放", lambda img, gray: cv2.resize(
+            img, None, fx=1.5, fy=1.5, interpolation=cv2.INTER_CUBIC
+        )),
+        "scale_2.0": ("2.0x缩放", lambda img, gray: cv2.resize(
+            img, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC
+        )),
+    }
+
+    # 策略组合映射
+    STRATEGY_MAP = {
+        "basic": ["original", "grayscale"],
+        "enhanced": ["original", "grayscale", "clahe", "adaptive_threshold", "otsu", "denoise", "sharpen", "morphology"],
+        "all": ["original", "grayscale", "clahe", "adaptive_threshold", "otsu", "denoise", "sharpen", "morphology", "scale_0.5", "scale_1.5", "scale_2.0"],
+    }
+
+    @staticmethod
+    def generate(
+        content: str,
+        size: int = 300,
+        error_correction: str = "H",
+        box_size: int = 10,
+        border: int = 4
+    ) -> str:
+        """
+        生成二维码
+
+        Args:
+            content: 二维码内容
+            size: 图片大小（像素）
+            error_correction: 纠错级别 (L/M/Q/H)
+            box_size: 每个格子的像素大小
+            border: 边框大小（格子数）
+
+        Returns:
+            base64编码的图片数据 (data:image/png;base64,...)
+
+        Raises:
+            ValueError: 参数错误时抛出异常
+        """
+        # 验证纠错级别
+        error_levels = {
+            "L": qrcode.constants.ERROR_CORRECT_L,  # 7% 容错
+            "M": qrcode.constants.ERROR_CORRECT_M,  # 15% 容错
+            "Q": qrcode.constants.ERROR_CORRECT_Q,  # 25% 容错
+            "H": qrcode.constants.ERROR_CORRECT_H,  # 30% 容错
+        }
+
+        if error_correction not in error_levels:
+            raise ValueError(f"无效的纠错级别: {error_correction}，支持: L/M/Q/H")
+
+        # 创建二维码对象
+        qr = qrcode.QRCode(
+            version=None,  # 自动确定版本
+            error_correction=error_levels[error_correction],
+            box_size=box_size,
+            border=border,
+        )
+
+        # 添加数据并生成
+        qr.add_data(content)
+        qr.make(fit=True)
+
+        # 生成图片
+        img = qr.make_image(fill_color="black", back_color="white")
+
+        # 调整到指定大小
+        img = img.resize((size, size), Image.Resampling.LANCZOS)
+
+        # 转换为base64
+        buffer = io.BytesIO()
+        img.save(buffer, format="PNG")
+        img_base64 = base64.b64encode(buffer.getvalue()).decode()
+
+        return f"data:image/png;base64,{img_base64}"
+
+    @staticmethod
+    async def load_image(image_source: str) -> np.ndarray:
+        """
+        加载图片（支持URL、base64、本地路径）
+
+        Args:
+            image_source: 图片来源
+                - URL: http://... 或 https://...
+                - base64: data:image/...;base64,...
+                - 本地路径: /path/to/image.png
+
+        Returns:
+            OpenCV图片对象 (BGR格式)
+
+        Raises:
+            ValueError: 图片加载失败时抛出异常
+        """
+        try:
+            # 检查是否为base64
+            if image_source.startswith("data:image"):
+                # 提取base64数据
+                base64_data = image_source.split(",")[1]
+                img_data = base64.b64decode(base64_data)
+                img_array = np.frombuffer(img_data, np.uint8)
+                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+
+            elif image_source.startswith("http://") or image_source.startswith("https://"):
+                # 下载图片
+                async with httpx.AsyncClient(timeout=30.0) as client:
+                    response = await client.get(image_source)
+                    response.raise_for_status()
+                    img_array = np.frombuffer(response.content, np.uint8)
+                    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+
+            else:
+                # 本地文件
+                img = cv2.imread(image_source)
+
+            if img is None:
+                raise ValueError("无法解析图片数据")
+
+            return img
+
+        except Exception as e:
+            raise ValueError(f"图片加载失败: {str(e)}")
+
+    @staticmethod
+    async def search_qrcode(img: np.ndarray)-> list[np.ndarray]:
+        """
+        搜索二维码，只搜索不解析
+
+        Args:
+            img: OpenCV图像对象
+
+        Returns:
+            二维码列表
+        """
+        detector = cv2.QRCodeDetector()
+        imgs = detector.detect(img)
+
+    @staticmethod
+    def decode(img: np.ndarray) -> Optional[str]:
+        """
+        使用 OpenCV QRCodeDetector 解码二维码
+
+        Args:
+            img: OpenCV图像对象
+
+        Returns:
+            二维码内容，如果没有检测到返回None
+        """
+        detector = cv2.QRCodeDetector()
+        data, bbox, _ = detector.detectAndDecode(img)
+
+        if data:
+            return data
+
+        return None
+
+    @staticmethod
+    async def parse(
+        image_source: str,
+        strategy: str = "auto"
+    ) -> dict:
+        """
+        解析二维码（使用 OpenCV + 按需预处理策略）
+
+        解码策略：
+        1. 根据 strategy 参数选择预处理步骤列表
+        2. 按需应用每种预处理算法并立即尝试解码
+        3. 一旦成功立即返回，避免不必要的计算
+
+        Args:
+            image_source: 图片来源（URL/base64/本地路径）
+            strategy: 预处理策略
+                - basic: 原图 + 灰度图（2种）
+                - enhanced: basic + 6种增强算法（8种）
+                - all: auto + 多尺度（11种）
+
+        Returns:
+            解析结果字典:
+            {
+                "success": bool,
+                "content": str or None,
+                "strategy_used": str,  # 使用的预处理策略名称
+                "preprocessing_index": int,  # 预处理索引
+                "total_attempts": int,
+                "message": str  # 仅失败时有
+            }
+        """
+        # 验证策略参数
+        if strategy not in QRCodeProcessor.STRATEGY_MAP:
+            raise ValueError(f"无效的策略: {strategy}，支持: {list(QRCodeProcessor.STRATEGY_MAP.keys())}")
+
+        # 加载原始图片
+        img = await QRCodeProcessor.load_image(image_source)
+
+        # 预先生成灰度图（很多预处理都需要）
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+        # 获取该策略对应的预处理步骤列表
+        preprocessing_steps = QRCodeProcessor.STRATEGY_MAP[strategy]
+
+        # 依次应用每种预处理并尝试解码
+        for idx, step_key in enumerate(preprocessing_steps):
+            strategy_name, preprocess_func = QRCodeProcessor.PREPROCESSING_STRATEGIES[step_key]
+
+            try:
+                # 按需处理图像
+                processed_img = preprocess_func(img, gray)
+
+                # 立即尝试解码
+                result = QRCodeProcessor.decode(processed_img)
+
+                if result:
+                    # 解码成功，立即返回
+                    return {
+                        "success": True,
+                        "content": result,
+                        "strategy_used": f"opencv_{strategy_name}",
+                        "preprocessing_index": idx,
+                        "total_attempts": idx + 1
+                    }
+
+            except Exception as e:
+                # 某个预处理步骤失败，继续尝试下一个
+                continue
+
+        # 所有预处理方法都失败
+        return {
+            "success": False,
+            "content": None,
+            "message": f"未检测到二维码或二维码损坏（已尝试 {len(preprocessing_steps)} 种预处理）",
+            "total_attempts": len(preprocessing_steps)
+        }
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def main():
+        # 示例用法
+        result = await QRCodeProcessor.parse("F:/Project/urbanLifeline/docs/qrcode.png", "enhanced")
+        print(result)
+
+    asyncio.run(main())