65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
from PIL import Image, UnidentifiedImageError
|
|
from fastapi import APIRouter, UploadFile, File, Depends, HTTPException
|
|
|
|
from app.core.dependencies import get_image_moderation_service
|
|
from app.models.response.image_moderation_response import \
|
|
ImageModerationResponse
|
|
from app.services.image_moderation_service import ImageModerationService
|
|
|
|
router = APIRouter(
|
|
prefix="/api/v1/moderation",
|
|
tags=["Moderation"]
|
|
)
|
|
|
|
MAX_SIZE = 5 * 1024 * 1024 # например, 5 МБ
|
|
|
|
ALLOWED_CONTENT_TYPES = (
|
|
"image/jpeg",
|
|
"image/png",
|
|
"image/webp",
|
|
)
|
|
|
|
|
|
@router.post("/image", response_model=ImageModerationResponse)
|
|
def moderate_image(
|
|
file: UploadFile = File(...),
|
|
service: ImageModerationService = Depends(get_image_moderation_service)
|
|
):
|
|
# 1. Проверка content-type — самая дешёвая, делаем её первой
|
|
if file.content_type not in ALLOWED_CONTENT_TYPES:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Unsupported image format"
|
|
)
|
|
|
|
# 2. Проверка размера — читаем файл целиком
|
|
contents = file.file.read()
|
|
|
|
if len(contents) > MAX_SIZE:
|
|
raise HTTPException(
|
|
status_code=413,
|
|
detail="Image too large"
|
|
)
|
|
|
|
file.file.seek(0) # обязательно вернуть указатель в начало!
|
|
|
|
# 3. Только теперь пытаемся открыть изображение
|
|
try:
|
|
image = Image.open(file.file)
|
|
image = image.convert("RGB")
|
|
except UnidentifiedImageError:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Invalid image"
|
|
)
|
|
|
|
prediction = service.moderate(image)
|
|
|
|
return ImageModerationResponse(
|
|
approved=prediction.approved,
|
|
score=prediction.score,
|
|
reason=prediction.reason,
|
|
label=prediction.label,
|
|
detected_labels=prediction.detected_labels
|
|
)
|