处理二维码

tool/qr.py

@@ -0,0 +1,179 @@
+"""二维码工具。
+
+基于 OpenCV 的 QRCodeDetector 完成二维码检测、裁剪与解码。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+import cv2
+import numpy as np
+
+
+@dataclass
+class QrCodeResult:
+    """单个二维码识别结果。"""
+
+    text: str
+    points: tuple[tuple[int, int], ...]
+    crop: np.ndarray | None = None
+
+
+@dataclass
+class QrScanResult:
+    """二维码扫描结果。"""
+
+    has_qr: bool
+    items: list[QrCodeResult]
+
+
+def _load_image(image: bytes | np.ndarray | str | Path) -> np.ndarray:
+    """把输入转成 BGR 图像。"""
+    if isinstance(image, np.ndarray):
+        if image.ndim == 2:
+            return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+        return image.copy()
+
+    if isinstance(image, (str, Path)):
+        data = Path(image).read_bytes()
+    else:
+        data = image
+
+    arr = np.frombuffer(data, dtype=np.uint8)
+    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if img is None:
+        raise ValueError("Failed to decode image data.")
+    return img
+
+
+def _order_points(points: np.ndarray) -> np.ndarray:
+    """把四个角点整理成左上、右上、右下、左下。"""
+    pts = np.asarray(points, dtype=np.float32).reshape(4, 2)
+    rect = np.zeros((4, 2), dtype=np.float32)
+
+    s = pts.sum(axis=1)
+    diff = np.diff(pts, axis=1)
+
+    rect[0] = pts[np.argmin(s)]
+    rect[2] = pts[np.argmax(s)]
+    rect[1] = pts[np.argmin(diff)]
+    rect[3] = pts[np.argmax(diff)]
+    return rect
+
+
+def _warp_qr_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
+    """按四边形点做透视矫正，截取二维码区域。"""
+    rect = _order_points(points)
+    (tl, tr, br, bl) = rect
+
+    width_a = np.linalg.norm(br - bl)
+    width_b = np.linalg.norm(tr - tl)
+    height_a = np.linalg.norm(tr - br)
+    height_b = np.linalg.norm(tl - bl)
+
+    width = max(int(round(max(width_a, width_b))), 1)
+    height = max(int(round(max(height_a, height_b))), 1)
+
+    dst = np.array(
+        [
+            [0, 0],
+            [width - 1, 0],
+            [width - 1, height - 1],
+            [0, height - 1],
+        ],
+        dtype=np.float32,
+    )
+
+    matrix = cv2.getPerspectiveTransform(rect, dst)
+    return cv2.warpPerspective(img, matrix, (width, height))
+
+
+def _normalize_points(points: np.ndarray | None) -> list[tuple[tuple[int, int], ...]]:
+    if points is None:
+        return []
+
+    arr = np.asarray(points)
+    if arr.ndim == 2 and arr.shape == (4, 2):
+        arr = arr[None, ...]
+    elif arr.ndim == 3 and arr.shape[-2:] == (4, 2):
+        pass
+    else:
+        return []
+
+    result: list[tuple[tuple[int, int], ...]] = []
+    for item in arr:
+        quad = tuple((int(round(x)), int(round(y))) for x, y in item)
+        result.append(quad)
+    return result
+
+
+def _decode_with_multi_detector(img: np.ndarray) -> tuple[list[str], list[tuple[tuple[int, int], ...]], list[np.ndarray]]:
+    detector = cv2.QRCodeDetector()
+
+    texts: list[str] = []
+    points_list: list[tuple[tuple[int, int], ...]] = []
+    crops: list[np.ndarray] = []
+
+    if hasattr(detector, "detectAndDecodeMulti"):
+        ok, decoded_info, points, _ = detector.detectAndDecodeMulti(img)
+        if ok and points is not None:
+            normalized_points = _normalize_points(points)
+            if isinstance(decoded_info, Iterable) and not isinstance(
+                decoded_info, (str, bytes)
+            ):
+                decoded_iter = [item or "" for item in decoded_info]
+            else:
+                decoded_iter = [decoded_info or ""]
+
+            for idx, quad in enumerate(normalized_points):
+                text = decoded_iter[idx] if idx < len(decoded_iter) else ""
+                crop = _warp_qr_image(img, np.array(quad, dtype=np.float32))
+                if not text:
+                    fallback_text, _, _ = detector.detectAndDecode(crop)
+                    text = fallback_text or ""
+                texts.append(text)
+                points_list.append(quad)
+                crops.append(crop)
+
+            return texts, points_list, crops
+
+    text, points, _ = detector.detectAndDecode(img)
+    normalized_points = _normalize_points(points)
+    if normalized_points:
+        crop = _warp_qr_image(img, np.array(normalized_points[0], dtype=np.float32))
+        if not text:
+            fallback_text, _, _ = detector.detectAndDecode(crop)
+            text = fallback_text or ""
+        return [text or ""], [normalized_points[0]], [crop]
+
+    return [], [], []
+
+
+def scan_qr(image: bytes | np.ndarray | str | Path) -> QrScanResult:
+    """扫描图片中的二维码，返回是否存在、位置和文本。"""
+    img = _load_image(image)
+    texts, points_list, crops = _decode_with_multi_detector(img)
+
+    items = [
+        QrCodeResult(text=text, points=points, crop=crop)
+        for text, points, crop in zip(texts, points_list, crops)
+    ]
+    return QrScanResult(has_qr=bool(items), items=items)
+
+
+def has_qr(image: bytes | np.ndarray | str | Path) -> bool:
+    """判断图片里有没有二维码。"""
+    return scan_qr(image).has_qr
+
+
+def crop_qr(image: bytes | np.ndarray | str | Path) -> list[np.ndarray]:
+    """裁剪出图片中的二维码区域。"""
+    return [item.crop for item in scan_qr(image).items if item.crop is not None]
+
+
+def decode_qr(image: bytes | np.ndarray | str | Path) -> list[str]:
+    """识别图片中的二维码内容。"""
+    return [item.text for item in scan_qr(image).items if item.text]