From 0d6b341945ccf8ac32e824878f694f061e8c9cad Mon Sep 17 00:00:00 2001 From: ChenJunheng Date: Sun, 12 Jul 2026 18:13:50 +0800 Subject: [PATCH] =?UTF-8?q?fix(watchlist):=20=E9=87=87=E7=BA=B3=20#94=20?= =?UTF-8?q?=E5=AE=A1=E6=9F=A5=20=E2=80=94=20=E9=98=B2=E5=A4=A7=E5=9B=BE=20?= =?UTF-8?q?OOM=E3=80=81OCR=20=E4=B8=8D=E9=98=BB=E5=A1=9E=E4=BA=8B=E4=BB=B6?= =?UTF-8?q?=E5=BE=AA=E7=8E=AF=E3=80=81=E6=94=B6=E7=B4=A7=E5=9B=BE=E7=89=87?= =?UTF-8?q?=E7=B1=BB=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为截图预处理增加像素上限与长边降采样；import-image 将 OCR 放入线程池； 去掉任意 image/* 放行，仅允许白名单 MIME/扩展名。 --- backend/app/api/watchlist.py | 14 ++++++-- .../app/services/watchlist_ocr/provider.py | 33 ++++++++++++++++--- backend/tests/test_watchlist_ocr.py | 23 ++++++++++++- 3 files changed, 61 insertions(+), 9 deletions(-) diff --git a/backend/app/api/watchlist.py b/backend/app/api/watchlist.py index 50a9458..004c3ef 100644 --- a/backend/app/api/watchlist.py +++ b/backend/app/api/watchlist.py @@ -81,12 +81,15 @@ def ocr_status(): @router.post("/import-image") async def import_from_image(request: Request, file: UploadFile = File(...)): """从自选截图识别股票代码，返回候选列表（不自动写入自选）。""" + import anyio + content_type = (file.content_type or "").split(";")[0].strip().lower() filename = (file.filename or "").lower() - ok_type = content_type in _IMPORT_IMAGE_TYPES or content_type.startswith("image/") + # 严格白名单：不接受任意 image/*（如 image/svg+xml） + ok_type = content_type in _IMPORT_IMAGE_TYPES ok_ext = filename.endswith((".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif")) if not ok_type and not ok_ext: - raise HTTPException(400, "仅支持 JPG / PNG / WebP 等图片") + raise HTTPException(400, "仅支持 JPG / PNG / WebP / BMP / GIF 图片") data = await file.read() if not data: @@ -97,7 +100,12 @@ async def import_from_image(request: Request, file: UploadFile = File(...)): existing = {r["symbol"] for r in watchlist.list_symbols()} data_dir = request.app.state.repo.store.data_dir try: - result = import_watchlist_image(data, data_dir, existing_symbols=existing) + # OCR 为同步 CPU/子进程；丢进线程池，避免卡住事件循环（行情 SSE 等） + result = await anyio.to_thread.run_sync( + lambda: import_watchlist_image(data, data_dir, existing_symbols=existing), + ) + except ValueError as e: + raise HTTPException(400, str(e)) from e except RuntimeError as e: raise HTTPException(503, str(e)) from e except Exception as e: # noqa: BLE001 diff --git a/backend/app/services/watchlist_ocr/provider.py b/backend/app/services/watchlist_ocr/provider.py index 4d33148..f7a6f3e 100644 --- a/backend/app/services/watchlist_ocr/provider.py +++ b/backend/app/services/watchlist_ocr/provider.py @@ -10,6 +10,13 @@ from PIL import Image, ImageEnhance, ImageOps logger = logging.getLogger(__name__) +# 像素上限：防 JPEG 高分辨率解码后 OOM（字节上限挡不住高压缩大图） +_MAX_PIXELS = 12_000_000 +# 长边上限：过大则降采样，减内存并加快 OCR +_MAX_EDGE = 2000 +# 过窄时适度放大，提升小字识别率 +_MIN_WIDTH = 1400 + class OcrProvider(ABC): """OCR 引擎接口。实现类只负责「图 → 文本」，代码抽取与证券匹配在 pipeline 中完成。""" @@ -26,19 +33,35 @@ class OcrProvider(ABC): def preprocess_for_ocr(image_bytes: bytes) -> Image.Image: - """暗色券商截图预处理：灰度、反相、增强对比，提升数字识别率。""" + """暗色券商截图预处理：像素上限、降采样、灰度、反相、增强对比。 + + Raises: + ValueError: 图片分辨率过高（像素数超限）。 + """ img = Image.open(BytesIO(image_bytes)) + img.load() # 强制解码，便于尽早失败 / 量尺寸 + + pixels = img.width * img.height + if pixels > _MAX_PIXELS: + raise ValueError("图片分辨率过高,请裁剪后重试") + + # 大图降采样：既减内存又提速 OCR（只放大不缩小的旧逻辑已去掉） + if max(img.size) > _MAX_EDGE: + img.thumbnail((_MAX_EDGE, _MAX_EDGE), Image.Resampling.LANCZOS) + if img.mode not in ("RGB", "L"): img = img.convert("RGB") gray = ImageOps.grayscale(img) # 暗底白字 → 白底黑字，Tesseract 更稳 inverted = ImageOps.invert(gray) contrasted = ImageEnhance.Contrast(inverted).enhance(1.8) - # 小字放大一档 + # 小字适度放大 w, h = contrasted.size - if w < 1400: - scale = 1400 / w - contrasted = contrasted.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS) + if w < _MIN_WIDTH: + scale = _MIN_WIDTH / w + contrasted = contrasted.resize( + (int(w * scale), int(h * scale)), Image.Resampling.LANCZOS + ) return contrasted diff --git a/backend/tests/test_watchlist_ocr.py b/backend/tests/test_watchlist_ocr.py index c91eead..dd266b4 100644 --- a/backend/tests/test_watchlist_ocr.py +++ b/backend/tests/test_watchlist_ocr.py @@ -1,16 +1,19 @@ """自选截图 OCR：代码抽取与 instruments 匹配（不依赖本机 tesseract）。""" from __future__ import annotations +from io import BytesIO from pathlib import Path import polars as pl +import pytest +from PIL import Image from app.services.watchlist_ocr.pipeline import ( extract_codes, import_watchlist_image, resolve_candidates, ) -from app.services.watchlist_ocr.provider import OcrProvider +from app.services.watchlist_ocr.provider import OcrProvider, preprocess_for_ocr class _FakeOcr(OcrProvider): @@ -26,6 +29,12 @@ class _FakeOcr(OcrProvider): return self._text +def _png_bytes(width: int, height: int) -> bytes: + buf = BytesIO() + Image.new("RGB", (width, height), color=(20, 20, 20)).save(buf, format="PNG") + return buf.getvalue() + + def test_extract_codes_order_and_dedupe(): text = "563230 融\n515880\n价格 1.340\n563230 重复\nXAUUSD\n601636" assert extract_codes(text) == ["563230", "515880", "601636"] @@ -73,3 +82,15 @@ def test_import_watchlist_image_with_fake_ocr(tmp_path: Path): assert result["codes"] == ["600036", "515880"] assert result["matched_count"] == 2 assert result["unmatched_count"] == 0 + + +def test_preprocess_rejects_excessive_pixels(): + # 4000×8000 = 32M 像素 > 12M 上限 + with pytest.raises(ValueError, match="分辨率过高"): + preprocess_for_ocr(_png_bytes(4000, 8000)) + + +def test_preprocess_downsamples_large_edge(): + # 2500×1000 = 2.5M 像素未超限，但长边 > 2000，应降采样 + out = preprocess_for_ocr(_png_bytes(2500, 1000)) + assert max(out.size) <= 2000