180 lines
5.4 KiB
Python
180 lines
5.4 KiB
Python
"""二维码工具。
|
|
|
|
基于 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]
|