""" 图片 OCR API 路由 OCR 识别 + AI 搜索 """ from __future__ import annotations import asyncio import base64 import json import logging import os import time from collections import defaultdict from typing import List, Optional from fastapi import APIRouter, Depends, HTTPException, UploadFile, Form, Query from fastapi.responses import StreamingResponse, Response from pydantic import BaseModel from sqlalchemy import select, func, text, delete from sqlalchemy.ext.asyncio import AsyncSession from app.config import settings from app.database import get_db from app.models.base import OCRImage from app.services.ocr_service import OCRService from app.services.llm_service import LLMService logger = logging.getLogger(__name__) router = APIRouter() # OCR 并发控制信号量（动态读取配置） _ocr_semaphore: asyncio.Semaphore | None = None def _get_ocr_semaphore() -> asyncio.Semaphore: """获取/重建 OCR 并发信号量""" global _ocr_semaphore from app.config import settings concurrency = settings.OCR_CONCURRENCY if _ocr_semaphore is None or _ocr_semaphore._value != concurrency: _ocr_semaphore = asyncio.Semaphore(concurrency) return _ocr_semaphore ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"} def _check_image_ext(filename: str) -> str: suffix = os.path.splitext(filename)[1].lower() if suffix not in ALLOWED_EXTENSIONS: raise HTTPException(status_code=400, detail=f"不支持的图片格式: {suffix}") return suffix def _save_ocr_result(record: OCRImage, ocr_result, db: AsyncSession): """保存 OCR 识别结果到数据库""" record.ocr_text = ocr_result.text record.confidence = ocr_result.confidence record.provider = ocr_result.provider record.status = "completed" async def _check_ocr_duplicate(ocr_text: str, db: AsyncSession) -> Optional[int]: """检查 OCR 文本是否已存在，返回已存在记录的 ID 或 None""" if not ocr_text: return None dup_query = select(OCRImage.id).where( OCRImage.ocr_text == ocr_text, OCRImage.status == "completed", ).limit(1) dup_result = await db.execute(dup_query) return dup_result.scalar_one_or_none() # ─── 上传 ─── @router.post("/upload", summary="上传图片") async def upload_image(file: UploadFile, db: AsyncSession = Depends(get_db)): if not file.filename: raise HTTPException(status_code=400, detail="文件名不能为空") _check_image_ext(file.filename) image_record = OCRImage(file_path=file.filename, status="pending") db.add(image_record) await db.flush() await db.refresh(image_record) return {"id": image_record.id, "file_path": image_record.file_path, "original_filename": file.filename, "status": image_record.status, "created_at": str(image_record.created_at)} # ─── 单文件识别 ─── @router.post("/{image_id}/recognize", summary="识别图片文字") async def recognize_image(image_id: int, db: AsyncSession = Depends(get_db)): result = await db.execute(select(OCRImage).where(OCRImage.id == image_id)) image_record = result.scalar_one_or_none() if image_record is None: raise HTTPException(status_code=404, detail="图片记录不存在") if not os.path.exists(image_record.file_path): raise HTTPException(status_code=404, detail=f"图片文件不存在: {image_record.file_path}") image_record.status = "processing" await db.flush() try: async with _get_ocr_semaphore(): ocr_result = await OCRService.recognize(image_record.file_path) _save_ocr_result(image_record, ocr_result, db) await db.flush() await db.refresh(image_record) return { "id": image_record.id, "file_path": image_record.file_path, "ocr_text": image_record.ocr_text, "confidence": image_record.confidence, "provider": image_record.provider, "status": image_record.status, "block_count": len(ocr_result.blocks), } except Exception as e: image_record.status = "failed" image_record.error_message = str(e) await db.flush() raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}") # ─── 批量识别 ─── @router.post("/batch-recognize", summary="批量上传并识别图片") async def batch_recognize(files: List[UploadFile], db: AsyncSession = Depends(get_db)): results = [] for file in files: if not file.filename: results.append({"filename": "unknown", "status": "error", "message": "文件名为空"}) continue # 只取文件名，去掉文件夹上传时带上的相对路径 filename = os.path.basename(file.filename) suffix = _check_image_ext(filename) dest_path = settings.images_dir / filename counter = 1 stem = os.path.splitext(filename)[0] while dest_path.exists(): dest_path = settings.images_dir / f"{stem}_{counter}{suffix}" counter += 1 content = await file.read() dest_path.write_bytes(content) image_record = OCRImage(file_path=str(dest_path), status="processing") db.add(image_record) await db.flush() try: async with _get_ocr_semaphore(): ocr_result = await OCRService.recognize(str(dest_path)) # OCR 内容去重 existing_id = await _check_ocr_duplicate(ocr_result.text, db) if existing_id is not None: # 已存在相同内容，回滚当前记录，删除图片文件 await db.delete(image_record) await db.flush() try: os.unlink(dest_path) except OSError: pass results.append({ "filename": file.filename, "file_path": str(dest_path), "status": "duplicate", "duplicate_of": existing_id, "ocr_text": ocr_result.text, }) continue _save_ocr_result(image_record, ocr_result, db) await db.flush() await db.refresh(image_record) results.append({ "id": image_record.id, "filename": file.filename, "file_path": str(dest_path), "status": "completed", "ocr_text": ocr_result.text, "confidence": ocr_result.confidence, "provider": ocr_result.provider, "block_count": len(ocr_result.blocks), }) except Exception as e: image_record.status = "failed" image_record.error_message = str(e) await db.flush() results.append({"id": image_record.id, "filename": file.filename, "file_path": str(dest_path), "status": "failed", "message": str(e)}) return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "duplicates": sum(1 for r in results if r["status"] == "duplicate"), "results": results} # ─── 服务器路径导入 ─── class PathImportRequest(BaseModel): paths: List[str] recursive: bool = False @router.post("/import-paths", summary="从服务器路径批量导入图片") async def import_from_paths(data: PathImportRequest, db: AsyncSession = Depends(get_db)): all_paths = [] for p in data.paths: if os.path.isfile(p): if os.path.splitext(p)[1].lower() in ALLOWED_EXTENSIONS: all_paths.append(p) elif os.path.isdir(p): walk = os.walk(p) if data.recursive else [(p, [], os.listdir(p))] for root, dirs, files in walk: for f in sorted(files): fp = os.path.join(root, f) if os.path.isfile(fp) and os.path.splitext(f)[1].lower() in ALLOWED_EXTENSIONS: all_paths.append(fp) if not all_paths: return {"total": 0, "success": 0, "failed": 0, "message": "未找到图片文件", "results": []} results = [] for file_path in all_paths: image_record = OCRImage(file_path=file_path, status="processing") db.add(image_record) await db.flush() try: async with _get_ocr_semaphore(): ocr_result = await OCRService.recognize(file_path) # OCR 内容去重 existing_id = await _check_ocr_duplicate(ocr_result.text, db) if existing_id is not None: await db.delete(image_record) await db.flush() results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "duplicate", "duplicate_of": existing_id, "ocr_text": ocr_result.text}) continue _save_ocr_result(image_record, ocr_result, db) await db.flush() await db.refresh(image_record) results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "completed", "ocr_text": ocr_result.text, "confidence": ocr_result.confidence, "provider": ocr_result.provider}) except Exception as e: image_record.status = "failed" image_record.error_message = str(e) await db.flush() results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "failed", "message": str(e)}) return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "results": results} # ─── 一步上传+识别 ─── @router.post("/recognize-direct", summary="直接上传并识别图片") async def recognize_direct(file: UploadFile, db: AsyncSession = Depends(get_db)): if not file.filename: raise HTTPException(status_code=400, detail="文件名不能为空") # 只取文件名，去掉文件夹上传时带上的相对路径（如 xx/01.jpg -> 01.jpg） filename = os.path.basename(file.filename) suffix = _check_image_ext(filename) dest_path = settings.images_dir / filename counter = 1 stem = os.path.splitext(filename)[0] while dest_path.exists(): dest_path = settings.images_dir / f"{stem}_{counter}{suffix}" counter += 1 content = await file.read() dest_path.write_bytes(content) image_record = OCRImage(file_path=str(dest_path), status="processing") db.add(image_record) await db.flush() try: async with _get_ocr_semaphore(): ocr_result = await OCRService.recognize(str(dest_path)) # OCR 内容去重 existing_id = await _check_ocr_duplicate(ocr_result.text, db) if existing_id is not None: await db.delete(image_record) await db.flush() try: os.unlink(dest_path) except OSError: pass return { "id": image_record.id, "file_path": str(dest_path), "original_filename": file.filename, "status": "duplicate", "duplicate_of": existing_id, "ocr_text": ocr_result.text, } _save_ocr_result(image_record, ocr_result, db) await db.flush() await db.refresh(image_record) return { "id": image_record.id, "file_path": image_record.file_path, "original_filename": file.filename, "ocr_text": image_record.ocr_text, "confidence": image_record.confidence, "provider": image_record.provider, "status": image_record.status, "blocks": [{"text": b.text, "bbox": b.bbox, "confidence": b.confidence} for b in ocr_result.blocks], } except Exception as e: image_record.status = "failed" image_record.error_message = str(e) await db.flush() raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}") # ─── LLM 精排搜索（SSE 流式） ─── @router.get("/search", summary="AI 搜索图片（流式）") async def search_images( keyword: str, limit: int = Query(5, ge=1, le=100), sort: str = Query("time_desc"), db: AsyncSession = Depends(get_db), ): """ AI 搜索图片（SSE 流式返回）。 1. 从 DB 取所有图片，按指定排序 2. 每批 5 条发给 Qwen3-8B 判断 3. 并发池最多 10 个请求 4. 匹配结果实时 SSE 推送 5. 找够 limit 个或遍历完 DB 后结束 """ async def event_stream(): found = 0 offset = 0 batch_size = settings.JUDGE_BATCH_SIZE max_concurrent = 10 semaphore = asyncio.Semaphore(max_concurrent) # 发送开始事件 yield f"data: {json.dumps({'type': 'start', 'keyword': keyword, 'limit': limit}, ensure_ascii=False)}\n\n" # 先获取总记录数 count_result = await db.execute( select(func.count(OCRImage.id)).where(OCRImage.status == "completed") ) total_records = count_result.scalar() or 0 if total_records == 0: yield f"data: {json.dumps({'type': 'done', 'total_found': 0, 'total_checked': 0}, ensure_ascii=False)}\n\n" return # 排序 order_clause = OCRImage.created_at.desc() if sort == "time_asc": order_clause = OCRImage.created_at.asc() elif sort == "random": order_clause = text("RANDOM()") # 用于收集结果的队列 result_queue = asyncio.Queue() checked_count = 0 async def process_batch(batch_items: list, batch_offset: int): """处理一批图片：发给 LLM 判断""" nonlocal checked_count async with semaphore: articles = [{"id": item.id, "text": item.ocr_text or ""} for item in batch_items] matched_ids = await LLMService.judge_batch(keyword, articles) checked_count += len(batch_items) # 发送进度 progress_data = { "type": "progress", "checked": checked_count, "total": total_records, "found": found, } await result_queue.put(progress_data) # 如果有匹配且未达上限，发送结果 for match_id in matched_ids: if found >= limit: break item = next((i for i in batch_items if i.id == match_id), None) if item: file_name = item.file_path.split('/')[-1] if item.file_path else '' # 读取图片并转为 base64（用于发送到企业微信等） image_base64 = "" if file_name: img_path = settings.images_dir / file_name if img_path.exists(): image_base64 = base64.b64encode(img_path.read_bytes()).decode() result_data = { "type": "result", "id": item.id, "file_path": item.file_path, "image_url": file_name, "image_base64": image_base64, "ocr_text": item.ocr_text, "ocr_text_preview": (item.ocr_text[:150] + "...") if item.ocr_text and len(item.ocr_text) > 150 else item.ocr_text, "confidence": item.confidence, "provider": item.provider, "created_at": str(item.created_at), } await result_queue.put(result_data) # 启动所有批次的异步任务 tasks = [] while offset < total_records: query = select(OCRImage).where(OCRImage.status == "completed").order_by(order_clause).offset(offset).limit(batch_size) result = await db.execute(query) batch_items = result.scalars().all() if not batch_items: break tasks.append(asyncio.create_task(process_batch(batch_items, offset))) offset += batch_size # 等待所有任务完成 async def wait_all(): await asyncio.gather(*tasks) await result_queue.put({"type": "all_done"}) asyncio.create_task(wait_all()) # 从队列读取结果并 yield while True: try: data = await asyncio.wait_for(result_queue.get(), timeout=300) except asyncio.TimeoutError: break if data.get("type") == "all_done": break if data.get("type") == "result": found += 1 yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n" if found >= limit: # 找够了，发送结束事件 yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n" return elif data.get("type") == "progress": yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n" # 发送结束事件 yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n" return StreamingResponse( event_stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) # ─── 删除单张图片 ─── @router.delete("/{image_id}", summary="删除单张图片", status_code=204) async def delete_image(image_id: int, db: AsyncSession = Depends(get_db)): result = await db.execute(select(OCRImage).where(OCRImage.id == image_id)) image_record = result.scalar_one_or_none() if image_record is None: raise HTTPException(status_code=404, detail="图片记录不存在") # 删除磁盘上的图片文件 try: os.unlink(image_record.file_path) except OSError: pass # 删除数据库记录 await db.delete(image_record) await db.flush() return Response(status_code=204) # ─── 其他接口 ─── @router.get("/{image_id}", summary="获取图片识别结果") async def get_image_result(image_id: int, db: AsyncSession = Depends(get_db)): result = await db.execute(select(OCRImage).where(OCRImage.id == image_id)) image_record = result.scalar_one_or_none() if image_record is None: raise HTTPException(status_code=404, detail="图片记录不存在") return image_record.to_dict() @router.get("", summary="获取图片列表") async def list_images(page: int = 1, page_size: int = 20, status: Optional[str] = None, db: AsyncSession = Depends(get_db)): query = select(OCRImage) count_query = select(func.count(OCRImage.id)) if status: query = query.where(OCRImage.status == status) count_query = count_query.where(OCRImage.status == status) total_result = await db.execute(count_query) total = total_result.scalar() or 0 offset = (page - 1) * page_size query = query.order_by(OCRImage.created_at.desc()).offset(offset).limit(page_size) result = await db.execute(query) items = result.scalars().all() return {"total": total, "page": page, "page_size": page_size, "items": [item.to_dict() for item in items]} # ─── 一键去重 ─── @router.post("/dedup", summary="一键去重：删除 OCR 内容重复的图片") async def dedup_images(db: AsyncSession = Depends(get_db)): # 查询所有 status='completed' 且 ocr_text 不为空的记录 query = select(OCRImage).where( OCRImage.status == "completed", OCRImage.ocr_text.isnot(None), OCRImage.ocr_text != "", ).order_by(OCRImage.created_at.asc()) result = await db.execute(query) all_records = result.scalars().all() total_checked = len(all_records) # 按 ocr_text 分组 groups: dict[str, list[OCRImage]] = defaultdict(list) for record in all_records: groups[record.ocr_text].append(record) duplicates_found = 0 deleted = 0 kept = 0 for ocr_text, group in groups.items(): if len(group) <= 1: kept += 1 continue # 保留 created_at 最早的一条（已按 asc 排序，第一条即最早） duplicates_found += len(group) - 1 kept += 1 for record in group[1:]: # 删除磁盘上的图片文件 try: os.unlink(record.file_path) except OSError: pass # 删除数据库记录 await db.delete(record) deleted += 1 await db.flush() return { "total_checked": total_checked, "duplicates_found": duplicates_found, "deleted": deleted, "kept": kept, }