95 lines
2.6 KiB
Python
95 lines
2.6 KiB
Python
"""图片上传与 OCR 提取触发。"""
|
|
from fastapi import (
|
|
APIRouter,
|
|
BackgroundTasks,
|
|
Depends,
|
|
HTTPException,
|
|
Request,
|
|
UploadFile,
|
|
status,
|
|
)
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.db import get_db
|
|
from app.deps import get_current_user
|
|
from app.models.image import Image
|
|
from app.models.job import ProcessingJob
|
|
from app.models.user import User
|
|
from app.ratelimit import AI_LIMIT, UPLOAD_LIMIT, limiter, user_key
|
|
from app.schemas.job import ImageOut, JobOut
|
|
from app.services import storage
|
|
from app.services.extract_pipeline import run_extract_job
|
|
|
|
router = APIRouter(prefix="/api/images", tags=["images"])
|
|
|
|
|
|
@router.post("", response_model=ImageOut, status_code=status.HTTP_201_CREATED)
|
|
@limiter.limit(UPLOAD_LIMIT, key_func=user_key)
|
|
async def upload(
|
|
request: Request,
|
|
file: UploadFile,
|
|
db: AsyncSession = Depends(get_db),
|
|
current: User = Depends(get_current_user),
|
|
):
|
|
data = await file.read()
|
|
try:
|
|
mime, ext = storage.validate_image(data)
|
|
except storage.ImageValidationError as e:
|
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(e))
|
|
|
|
sha = storage.sha256_hex(data)
|
|
# 去重：同用户同图直接复用
|
|
existing = await db.scalar(
|
|
select(Image).where(Image.user_id == current.id, Image.sha256 == sha)
|
|
)
|
|
if existing:
|
|
return existing
|
|
|
|
try:
|
|
key = await storage.upload_image_async(current.id, data, ext, mime)
|
|
except storage.StorageError as e:
|
|
raise HTTPException(status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(e))
|
|
|
|
image = Image(
|
|
user_id=current.id,
|
|
object_key=key,
|
|
mime=mime,
|
|
size_bytes=len(data),
|
|
sha256=sha,
|
|
)
|
|
db.add(image)
|
|
await db.commit()
|
|
await db.refresh(image)
|
|
return image
|
|
|
|
|
|
@router.post(
|
|
"/{image_id}/extract",
|
|
response_model=JobOut,
|
|
status_code=status.HTTP_202_ACCEPTED,
|
|
)
|
|
@limiter.limit(AI_LIMIT, key_func=user_key)
|
|
async def extract(
|
|
request: Request,
|
|
image_id: str,
|
|
background: BackgroundTasks,
|
|
db: AsyncSession = Depends(get_db),
|
|
current: User = Depends(get_current_user),
|
|
):
|
|
image = await db.scalar(
|
|
select(Image).where(Image.id == image_id, Image.user_id == current.id)
|
|
)
|
|
if image is None:
|
|
raise HTTPException(status.HTTP_404_NOT_FOUND, detail="图片不存在")
|
|
|
|
job = ProcessingJob(user_id=current.id, image_id=image.id, kind="ocr_extract")
|
|
db.add(job)
|
|
await db.commit()
|
|
await db.refresh(job)
|
|
|
|
background.add_task(
|
|
run_extract_job, job.id, current.id, image.object_key, image.mime
|
|
)
|
|
return job
|