90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
|
import base64
|
||
|
|
import logging
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
|
||
|
|
from fastapi import APIRouter, HTTPException
|
||
|
|
|
||
|
|
from app.schemas.common import TaskResponse, TaskStatus
|
||
|
|
from app.schemas.utilities import ClassificationResponse, IconRequest
|
||
|
|
from app.services import freepik_client, task_tracker
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
router = APIRouter(prefix='/api/v1/util', tags=['utilities'])
|
||
|
|
|
||
|
|
|
||
|
|
@router.post('/remove-background')
|
||
|
|
async def remove_background(request: dict):
|
||
|
|
"""Remove background from an image. Returns processed image as base64."""
|
||
|
|
image = request.get('image')
|
||
|
|
if not image:
|
||
|
|
raise HTTPException(status_code=400, detail='image field is required')
|
||
|
|
|
||
|
|
result_bytes = await freepik_client.remove_background(image)
|
||
|
|
return {
|
||
|
|
'image': base64.b64encode(result_bytes).decode(),
|
||
|
|
'content_type': 'image/png',
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@router.post('/classify', response_model=ClassificationResponse)
|
||
|
|
async def classify_image(request: dict):
|
||
|
|
"""Classify whether an image is AI-generated."""
|
||
|
|
image = request.get('image')
|
||
|
|
if not image:
|
||
|
|
raise HTTPException(status_code=400, detail='image field is required')
|
||
|
|
|
||
|
|
result = await freepik_client.classify_image(image)
|
||
|
|
data = result.get('data', result)
|
||
|
|
return ClassificationResponse(
|
||
|
|
is_ai_generated=data.get('is_ai_generated', False),
|
||
|
|
ai_probability=data.get('ai_probability', 0.0),
|
||
|
|
human_probability=data.get('human_probability', 0.0),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@router.post('/audio-isolate', response_model=TaskResponse)
|
||
|
|
async def audio_isolate(request: dict):
|
||
|
|
"""Isolate audio tracks from an audio file."""
|
||
|
|
audio = request.get('audio')
|
||
|
|
if not audio:
|
||
|
|
raise HTTPException(status_code=400, detail='audio field is required')
|
||
|
|
|
||
|
|
result = await freepik_client.isolate_audio(audio)
|
||
|
|
data = result.get('data', result)
|
||
|
|
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||
|
|
if not freepik_task_id:
|
||
|
|
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||
|
|
|
||
|
|
internal_id = task_tracker.submit(freepik_task_id, {'operation': 'audio-isolate'})
|
||
|
|
return TaskResponse(
|
||
|
|
task_id=internal_id,
|
||
|
|
status=TaskStatus.pending,
|
||
|
|
created_at=datetime.now(timezone.utc),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# Icon generation lives under /api/v1/generate/ but is simple enough to keep here
|
||
|
|
icon_router = APIRouter(prefix='/api/v1/generate', tags=['utilities'])
|
||
|
|
|
||
|
|
|
||
|
|
@icon_router.post('/icon', response_model=TaskResponse)
|
||
|
|
async def generate_icon(request: IconRequest):
|
||
|
|
result = await freepik_client.generate_icon(
|
||
|
|
prompt=request.prompt,
|
||
|
|
color=request.color,
|
||
|
|
shape=request.shape,
|
||
|
|
style=request.style,
|
||
|
|
)
|
||
|
|
data = result.get('data', result)
|
||
|
|
freepik_task_id = str(data.get('task_id') or data.get('id', ''))
|
||
|
|
if not freepik_task_id:
|
||
|
|
raise HTTPException(status_code=502, detail='No task_id in Freepik response')
|
||
|
|
|
||
|
|
internal_id = task_tracker.submit(freepik_task_id, {'operation': 'icon'})
|
||
|
|
return TaskResponse(
|
||
|
|
task_id=internal_id,
|
||
|
|
status=TaskStatus.pending,
|
||
|
|
created_at=datetime.now(timezone.utc),
|
||
|
|
)
|