534 lines
21 KiB
Python
534 lines
21 KiB
Python
"""
|
||
图片 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, async_session_factory
|
||
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"),
|
||
):
|
||
"""
|
||
AI 搜索图片（SSE 流式返回）。
|
||
注意：不使用 Depends(get_db)，因为 StreamingResponse 的生命周期
|
||
超出 request scope，需要在生成器内部自行管理 session。
|
||
"""
|
||
|
||
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"
|
||
|
||
# 在生成器内部自行管理 session，避免 get_db 提前关闭
|
||
async with async_session_factory() as db:
|
||
# 先获取总记录数
|
||
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,
|
||
}
|