添加Ai 能力，重构项目结构
This commit is contained in:
zk
@@ -0,0 +1 @@
|
||||
"""爬虫工具集。"""
|
||||
@@ -0,0 +1,274 @@
|
||||
"""Playwright 页面抓取工具。
|
||||
|
||||
对外仅暴露两个动作：
|
||||
- open_page(url, wait_ms): 打开页面，返回 page_id
|
||||
- query(page_id, selector): 选中节点并提取文字、图片 URL 与属性
|
||||
|
||||
实现上内部常驻一个 Browser；每次 open_page 创建独立 BrowserContext 和 Page。
|
||||
Playwright 只在一个后台线程里运行，但外层 API 保持同步。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import queue
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
from playwright.sync_api import Browser, BrowserContext, Page, sync_playwright
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeSnapshot:
|
||||
"""页面节点快照。"""
|
||||
|
||||
tag: str
|
||||
text: str
|
||||
image_urls: list[str]
|
||||
attrs: dict[str, str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PageSession:
|
||||
context: BrowserContext
|
||||
page: Page
|
||||
|
||||
|
||||
class _BrowserRuntime:
|
||||
"""在单独线程中管理 Playwright 生命周期。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._thread: threading.Thread | None = None
|
||||
self._ready = threading.Event()
|
||||
self._closed = threading.Event()
|
||||
self._job_queue: queue.Queue[
|
||||
tuple[Callable[..., Any], tuple[Any, ...], dict[str, Any], Future[Any]]
|
||||
] = queue.Queue()
|
||||
self._lock = threading.RLock()
|
||||
self._playwright = None
|
||||
self._browser: Browser | None = None
|
||||
self._pages: dict[str, _PageSession] = {}
|
||||
atexit.register(self.close)
|
||||
|
||||
def _thread_main(self) -> None:
|
||||
with sync_playwright() as playwright:
|
||||
self._playwright = playwright
|
||||
self._browser = playwright.chromium.launch(headless=True)
|
||||
self._ready.set()
|
||||
|
||||
while True:
|
||||
job = self._job_queue.get()
|
||||
if job is None:
|
||||
break
|
||||
|
||||
fn, args, kwargs, fut = job
|
||||
try:
|
||||
result = fn(*args, **kwargs)
|
||||
except BaseException as exc: # pragma: no cover - propagate failures
|
||||
fut.set_exception(exc)
|
||||
else:
|
||||
fut.set_result(result)
|
||||
|
||||
for session in self._pages.values():
|
||||
try:
|
||||
session.page.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
session.context.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._pages.clear()
|
||||
|
||||
if self._browser is not None:
|
||||
try:
|
||||
self._browser.close()
|
||||
finally:
|
||||
self._browser = None
|
||||
|
||||
self._playwright = None
|
||||
self._closed.set()
|
||||
|
||||
def _ensure_thread(self) -> None:
|
||||
with self._lock:
|
||||
if self._thread and self._thread.is_alive():
|
||||
return
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=self._thread_main,
|
||||
name="offerpie-playwright",
|
||||
daemon=True,
|
||||
)
|
||||
self._thread.start()
|
||||
self._ready.wait()
|
||||
|
||||
def _run(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
||||
self._ensure_thread()
|
||||
fut: Future[Any] = Future()
|
||||
self._job_queue.put((fn, args, kwargs, fut))
|
||||
return fut.result()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
if self._closed.is_set():
|
||||
return
|
||||
if self._thread is None:
|
||||
return
|
||||
self._job_queue.put(None)
|
||||
self._closed.wait(timeout=10)
|
||||
|
||||
def _ensure_browser(self) -> Browser:
|
||||
if self._browser is None:
|
||||
raise RuntimeError("Browser runtime is not ready.")
|
||||
return self._browser
|
||||
|
||||
def _open_page_impl(self, url: str, wait_ms: int) -> str:
|
||||
browser = self._ensure_browser()
|
||||
context = browser.new_context()
|
||||
page = context.new_page()
|
||||
page.goto(url, wait_until="domcontentloaded", timeout=30000)
|
||||
if wait_ms > 0:
|
||||
page.wait_for_timeout(wait_ms)
|
||||
|
||||
page_id = uuid.uuid4().hex
|
||||
self._pages[page_id] = _PageSession(context=context, page=page)
|
||||
return page_id
|
||||
|
||||
def _query_impl(self, page_id: str, selector: str) -> list[NodeSnapshot]:
|
||||
session = self._pages.get(page_id)
|
||||
if session is None:
|
||||
raise KeyError(f"Unknown page_id: {page_id}")
|
||||
|
||||
locator = session.page.locator(selector)
|
||||
payload = locator.evaluate_all(
|
||||
r"""(elements) => elements.map((el) => {
|
||||
const attrs = {};
|
||||
for (const attr of Array.from(el.attributes || [])) {
|
||||
attrs[attr.name] = attr.value;
|
||||
}
|
||||
|
||||
const imageUrls = [];
|
||||
const addImageUrl = (url) => {
|
||||
if (!url) return;
|
||||
try {
|
||||
const absUrl = new URL(url, location.href).href;
|
||||
if (
|
||||
absUrl.startsWith('data:') ||
|
||||
absUrl.startsWith('blob:') ||
|
||||
absUrl === 'about:blank'
|
||||
) {
|
||||
return;
|
||||
}
|
||||
imageUrls.push(absUrl);
|
||||
} catch {
|
||||
if (
|
||||
url.startsWith('data:') ||
|
||||
url.startsWith('blob:') ||
|
||||
url === 'about:blank'
|
||||
) {
|
||||
return;
|
||||
}
|
||||
imageUrls.push(url);
|
||||
}
|
||||
};
|
||||
|
||||
const resolveImageUrl = (img) => {
|
||||
const candidates = [
|
||||
img.getAttribute('data-src'),
|
||||
img.getAttribute('data-original'),
|
||||
img.getAttribute('data-croporisrc'),
|
||||
img.currentSrc,
|
||||
img.src,
|
||||
img.getAttribute('src'),
|
||||
];
|
||||
|
||||
const srcset = img.getAttribute('srcset');
|
||||
if (srcset) {
|
||||
for (const part of srcset.split(',')) {
|
||||
const candidate = part.trim().split(/\s+/)[0];
|
||||
candidates.push(candidate);
|
||||
}
|
||||
}
|
||||
|
||||
for (const candidate of candidates) {
|
||||
if (!candidate) continue;
|
||||
if (
|
||||
candidate.startsWith('data:') ||
|
||||
candidate.startsWith('blob:') ||
|
||||
candidate === 'about:blank'
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
const absUrl = new URL(candidate, location.href).href;
|
||||
if (
|
||||
absUrl.startsWith('data:') ||
|
||||
absUrl.startsWith('blob:') ||
|
||||
absUrl === 'about:blank'
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
return absUrl;
|
||||
} catch {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
};
|
||||
|
||||
if (el.tagName === 'IMG') {
|
||||
addImageUrl(resolveImageUrl(el));
|
||||
}
|
||||
|
||||
for (const img of Array.from(el.querySelectorAll('img'))) {
|
||||
addImageUrl(resolveImageUrl(img));
|
||||
}
|
||||
|
||||
const text = (el.innerText || el.textContent || '').trim();
|
||||
|
||||
return {
|
||||
tag: (el.tagName || '').toLowerCase(),
|
||||
text,
|
||||
image_urls: Array.from(new Set(imageUrls)),
|
||||
attrs,
|
||||
};
|
||||
})"""
|
||||
)
|
||||
return [
|
||||
NodeSnapshot(
|
||||
tag=item.get("tag", ""),
|
||||
text=item.get("text", ""),
|
||||
image_urls=list(item.get("image_urls", [])),
|
||||
attrs=dict(item.get("attrs", {})),
|
||||
)
|
||||
for item in payload
|
||||
]
|
||||
|
||||
|
||||
_runtime = _BrowserRuntime()
|
||||
|
||||
|
||||
def open_page(url: str, wait_ms: int = 3000) -> str:
|
||||
"""打开页面并返回 page_id。"""
|
||||
return _runtime._run(_runtime._open_page_impl, url, wait_ms)
|
||||
|
||||
|
||||
def query(page_id: str, selector: str) -> list[NodeSnapshot]:
|
||||
"""按选择器抓取节点，提取文本、图片 URL 与属性。"""
|
||||
items = _runtime._run(_runtime._query_impl, page_id, selector)
|
||||
seen_urls: set[str] = set()
|
||||
|
||||
for item in items:
|
||||
unique_urls: list[str] = []
|
||||
for url in item.image_urls:
|
||||
if url in seen_urls:
|
||||
continue
|
||||
seen_urls.add(url)
|
||||
unique_urls.append(url)
|
||||
item.image_urls = unique_urls
|
||||
|
||||
return items
|
||||
+146
@@ -0,0 +1,146 @@
|
||||
"""二维码检测与裁剪工具。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class QrRegion:
|
||||
"""单个二维码区域。"""
|
||||
|
||||
points: tuple[tuple[int, int], ...]
|
||||
crop: np.ndarray | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class QrDetectResult:
|
||||
"""二维码检测结果。"""
|
||||
|
||||
has_qr: bool
|
||||
items: list[QrRegion]
|
||||
|
||||
|
||||
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 _detect_points(img: np.ndarray) -> list[tuple[tuple[int, int], ...]]:
|
||||
detector = cv2.QRCodeDetector()
|
||||
|
||||
if hasattr(detector, "detectMulti"):
|
||||
ok, points = detector.detectMulti(img)
|
||||
if ok:
|
||||
normalized_points = _normalize_points(points)
|
||||
if normalized_points:
|
||||
return normalized_points
|
||||
|
||||
ok, points = detector.detect(img)
|
||||
if ok:
|
||||
normalized_points = _normalize_points(points)
|
||||
if normalized_points:
|
||||
return normalized_points
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def scan_qr(image: bytes | np.ndarray | str | Path) -> QrDetectResult:
|
||||
"""扫描图片中是否存在二维码，并返回二维码区域。"""
|
||||
img = _load_image(image)
|
||||
points_list = _detect_points(img)
|
||||
|
||||
items = [
|
||||
QrRegion(points=points, crop=_warp_qr_image(img, np.array(points, dtype=np.float32)))
|
||||
for points in points_list
|
||||
]
|
||||
return QrDetectResult(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]
|
||||
@@ -0,0 +1,36 @@
|
||||
"""OCR 工具
|
||||
|
||||
基于 RapidOCR(ONNXRuntime 后端)的图片文字提取工具。
|
||||
模块级持有单例引擎,导入时完成初始化,后续直接复用。
|
||||
"""
|
||||
|
||||
from rapidocr import RapidOCR
|
||||
|
||||
# 模块级单例:导入时初始化一次,全局复用
|
||||
# 初始化配置(RapidOCR 3.x 用 params 字典,key 为配置项点路径):
|
||||
# - 关闭方向分类(Cls):公告图基本都是正的,省一道计算
|
||||
# - 其余检测/识别/后端配置沿用 RapidOCR 默认值,避免和新版枚举参数冲突
|
||||
_engine = RapidOCR(
|
||||
params={
|
||||
"Global.use_cls": False,
|
||||
"Global.text_score": 0.5,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def ocr(image: bytes) -> str:
|
||||
"""识别图片中的文字。
|
||||
|
||||
Args:
|
||||
image: 图片的字节数据(爬虫下载下来的图片流,无需落盘)。
|
||||
|
||||
Returns:
|
||||
识别出的文字,多段以换行拼接为一段文本;没有文字或识别失败时返回空字符串。
|
||||
"""
|
||||
result = _engine(image)
|
||||
|
||||
# .txts 是识别出的多段文字(元组),没识别到时可能为 None 或空
|
||||
if not result or not result.txts:
|
||||
return ""
|
||||
|
||||
return "\n".join(result.txts)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""二维码识别工具。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .cv import _load_image, _normalize_points, _warp_qr_image
|
||||
|
||||
|
||||
def _decode_with_detector(img: np.ndarray) -> list[str]:
|
||||
detector = cv2.QRCodeDetector()
|
||||
texts: list[str] = []
|
||||
|
||||
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, (list, tuple)):
|
||||
decoded_iter = [item or "" for item in decoded_info]
|
||||
else:
|
||||
decoded_iter = [decoded_info or ""]
|
||||
|
||||
for idx, points_item in enumerate(normalized_points):
|
||||
text = decoded_iter[idx] if idx < len(decoded_iter) else ""
|
||||
if not text:
|
||||
crop = _warp_qr_image(img, np.array(points_item, dtype=np.float32))
|
||||
fallback_text, _, _ = detector.detectAndDecode(crop)
|
||||
text = fallback_text or ""
|
||||
if text:
|
||||
texts.append(text)
|
||||
if texts:
|
||||
return texts
|
||||
|
||||
text, points, _ = detector.detectAndDecode(img)
|
||||
if text:
|
||||
return [text]
|
||||
|
||||
normalized_points = _normalize_points(points)
|
||||
if normalized_points:
|
||||
crop = _warp_qr_image(img, np.array(normalized_points[0], dtype=np.float32))
|
||||
fallback_text, _, _ = detector.detectAndDecode(crop)
|
||||
if fallback_text:
|
||||
return [fallback_text]
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def decode_qr(image: bytes | np.ndarray | str | Path) -> list[str]:
|
||||
"""识别图片中的二维码内容。"""
|
||||
img = _load_image(image)
|
||||
return _decode_with_detector(img)
|
||||
Reference in New Issue
Block a user