GuangYu/HuiBrain
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix docker issue

Browse Source
This commit is contained in:
EduBrain Dev
2026-04-14 17:25:06 +08:00
parent 4c6a20e5fc
commit c5f96056fa
9 changed files with 170 additions and 537 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
app/api/v1/images.py
View File

@@ -1,6 +1,6 @@
"""
图片 OCR API 路由
OCR 识别 + BM25 倒排索引搜索
OCR 识别 + AI 搜索
"""
from __future__ import annotations
@@ -24,7 +24,6 @@ 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.search_engine import BM25Index, tokenize
from app.services.llm_service import LLMService
logger = logging.getLogger(__name__)
@@ -45,9 +44,6 @@ def _get_ocr_semaphore() -> asyncio.Semaphore:
ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
# 全局 BM25 索引实例（在 main.py 中初始化）
bm25_index: Optional[BM25Index] = None
def _check_image_ext(filename: str) -> str:
suffix = os.path.splitext(filename)[1].lower()
@@ -62,41 +58,8 @@ def _save_ocr_result(record: OCRImage, ocr_result, db: AsyncSession):
record.confidence = ocr_result.confidence
record.provider = ocr_result.provider
record.status = "completed"
# 保留 keywords（如果有）
if hasattr(ocr_result, 'keywords') and ocr_result.keywords:
record.tags = json.dumps(ocr_result.keywords, ensure_ascii=False)
def _add_to_index(record: OCRImage):
"""将 OCR 结果添加到 BM25 索引"""
if bm25_index is None or not record.ocr_text:
return
# 索引内容 = OCR 文本（权重最高）
index_text = record.ocr_text
# 如果有 tags/story_summary 也加入索引
if record.tags:
try:
tags = json.loads(record.tags)
if tags:
index_text += "\n" + " ".join(tags)
except Exception:
pass
if record.story_summary:
index_text += "\n" + record.story_summary
bm25_index.add_document(
doc_id=record.id,
text=index_text,
metadata={
"file_path": record.file_path,
"tags": json.loads(record.tags) if record.tags else [],
"story_summary": record.story_summary or "",
"confidence": record.confidence,
"provider": record.provider,
"created_at": str(record.created_at),
},
)
async def _check_ocr_duplicate(ocr_text: str, db: AsyncSession) -> Optional[int]:
"""检查 OCR 文本是否已存在，返回已存在记录的 ID 或 None"""
@@ -144,13 +107,11 @@ async def recognize_image(image_id: int, db: AsyncSession = Depends(get_db)):
_save_ocr_result(image_record, ocr_result, db)
await db.flush()
await db.refresh(image_record)
_add_to_index(image_record)
return {
"id": image_record.id,
"file_path": image_record.file_path,
"ocr_text": image_record.ocr_text,
"tags": json.loads(image_record.tags) if image_record.tags else [],
"confidence": image_record.confidence,
"provider": image_record.provider,
"status": image_record.status,
@@ -208,12 +169,10 @@ async def batch_recognize(files: List[UploadFile], db: AsyncSession = Depends(ge
_save_ocr_result(image_record, ocr_result, db)
await db.flush()
await db.refresh(image_record)
_add_to_index(image_record)
results.append({
"id": image_record.id, "filename": file.filename,
"file_path": str(dest_path), "status": "completed",
"ocr_text": ocr_result.text,
"tags": json.loads(image_record.tags) if image_record.tags else [],
"confidence": ocr_result.confidence, "provider": ocr_result.provider,
"block_count": len(ocr_result.blocks),
})
@@ -266,8 +225,7 @@ async def import_from_paths(data: PathImportRequest, db: AsyncSession = Depends(
_save_ocr_result(image_record, ocr_result, db)
await db.flush()
await db.refresh(image_record)
_add_to_index(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, "tags": json.loads(image_record.tags) if image_record.tags else [], "confidence": ocr_result.confidence, "provider": ocr_result.provider})
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)
@@ -315,12 +273,10 @@ async def recognize_direct(file: UploadFile, db: AsyncSession = Depends(get_db))
_save_ocr_result(image_record, ocr_result, db)
await db.flush()
await db.refresh(image_record)
_add_to_index(image_record)
return {
"id": image_record.id, "file_path": image_record.file_path,
"original_filename": file.filename,
"ocr_text": image_record.ocr_text,
"tags": json.loads(image_record.tags) if image_record.tags else [],
"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],
@@ -420,7 +376,6 @@ async def search_images(
"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,
"tags": json.loads(item.tags) if item.tags else [],
"confidence": item.confidence,
"provider": item.provider,
"created_at": str(item.created_at),
@@ -489,9 +444,6 @@ async def delete_image(image_id: int, db: AsyncSession = Depends(get_db)):
image_record = result.scalar_one_or_none()
if image_record is None:
raise HTTPException(status_code=404, detail="图片记录不存在")
# 从 BM25 索引移除
if bm25_index is not None:
bm25_index.remove_document(image_record.id)
# 删除磁盘上的图片文件
try:
os.unlink(image_record.file_path)
@@ -561,9 +513,6 @@ async def dedup_images(db: AsyncSession = Depends(get_db)):
duplicates_found += len(group) - 1
kept += 1
for record in group[1:]:
# 从 BM25 索引移除
if bm25_index is not None:
bm25_index.remove_document(record.id)
# 删除磁盘上的图片文件
try:
os.unlink(record.file_path)