295 lines
9.7 KiB
Python
295 lines
9.7 KiB
Python
"""Retouch service: upload, process via Flux.1-Pro, and retrieve results."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.models.retouch import RetouchResult, RetouchStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ALLOWED_MIME_TYPES = {"image/png", "image/jpeg", "image/jpg", "image/webp"}
|
|
|
|
RETOUCH_SYSTEM_PROMPT = (
|
|
"You are an expert automotive photo retoucher. "
|
|
"Enhance the provided vehicle image with the following improvements: "
|
|
"1. Remove or replace the background with a clean, professional studio backdrop. "
|
|
"2. Correct color balance and enhance saturation for a vibrant, realistic look. "
|
|
"3. Remove reflections and glare from windows and bodywork. "
|
|
"4. Clean up minor blemishes, dust, and scratches on the vehicle surface. "
|
|
"5. Ensure the vehicle is well-lit and centered. "
|
|
"Return the enhanced image."
|
|
)
|
|
|
|
|
|
def validate_mime_type(mime_type: str) -> bool:
|
|
"""Check if the MIME type is an allowed image type."""
|
|
return mime_type.lower() in ALLOWED_MIME_TYPES
|
|
|
|
|
|
def validate_file_size(file_size: int) -> bool:
|
|
"""Check if file size is within the configured limit."""
|
|
max_bytes = settings.MAX_FILE_SIZE_MB * 1024 * 1024
|
|
return file_size <= max_bytes
|
|
|
|
|
|
def _get_file_extension(mime_type: str) -> str:
|
|
"""Map MIME type to file extension."""
|
|
mapping = {
|
|
"image/png": ".png",
|
|
"image/jpeg": ".jpg",
|
|
"image/jpg": ".jpg",
|
|
"image/webp": ".webp",
|
|
}
|
|
return mapping.get(mime_type.lower(), ".png")
|
|
|
|
|
|
def generate_retouch_prompt(vehicle_info: dict[str, Any] | None = None) -> str:
|
|
"""Generate a retouch prompt based on optional vehicle info.
|
|
|
|
Args:
|
|
vehicle_info: Optional dict with make, model, year, color etc.
|
|
|
|
Returns:
|
|
A detailed prompt string for the image retouch model.
|
|
"""
|
|
base_prompt = (
|
|
"Professional automotive photograph retouching: "
|
|
"Remove the original background and replace with a clean white studio backdrop. "
|
|
"Enhance color correction for realistic, vibrant tones. "
|
|
"Remove reflections and glare from windows and paint. "
|
|
"Clean up dust, scratches, and minor blemishes on the bodywork. "
|
|
"Ensure even, professional lighting across the entire vehicle. "
|
|
"Sharpen details on wheels, grille, and badges."
|
|
)
|
|
|
|
if vehicle_info:
|
|
parts = [base_prompt]
|
|
make = vehicle_info.get("make")
|
|
model = vehicle_info.get("model")
|
|
color = vehicle_info.get("color")
|
|
if make and model:
|
|
parts.append(f"The vehicle is a {make} {model}.")
|
|
if color:
|
|
parts.append(f"The vehicle color is {color}; ensure it looks accurate and rich.")
|
|
return " ".join(parts)
|
|
|
|
return base_prompt
|
|
|
|
|
|
async def upload_retouch_file(
|
|
db: AsyncSession,
|
|
file_bytes: bytes,
|
|
file_name: str,
|
|
mime_type: str,
|
|
vehicle_id: uuid.UUID | None = None,
|
|
) -> RetouchResult:
|
|
"""Save uploaded image and create a RetouchResult record.
|
|
|
|
Raises ValueError for invalid MIME type or file size.
|
|
"""
|
|
if not validate_mime_type(mime_type):
|
|
raise ValueError(f"Invalid MIME type: {mime_type}. Only image/* types are allowed.")
|
|
|
|
if not validate_file_size(len(file_bytes)):
|
|
raise ValueError(f"File size exceeds limit of {settings.MAX_FILE_SIZE_MB} MB")
|
|
|
|
upload_dir = settings.UPLOAD_DIR
|
|
os.makedirs(upload_dir, exist_ok=True)
|
|
|
|
ext = _get_file_extension(mime_type)
|
|
stored_filename = f"retouch_{uuid.uuid4().hex}{ext}"
|
|
file_path = os.path.join(upload_dir, stored_filename)
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(file_bytes)
|
|
|
|
result = RetouchResult(
|
|
original_file_path=file_path,
|
|
original_file_name=file_name,
|
|
mime_type=mime_type,
|
|
status=RetouchStatus.pending.value,
|
|
vehicle_id=vehicle_id,
|
|
)
|
|
db.add(result)
|
|
await db.flush()
|
|
await db.refresh(result)
|
|
return result
|
|
|
|
|
|
async def get_result(db: AsyncSession, result_id: uuid.UUID) -> RetouchResult | None:
|
|
"""Get a retouch result by ID."""
|
|
stmt = select(RetouchResult).where(RetouchResult.id == result_id)
|
|
res = await db.execute(stmt)
|
|
return res.scalar_one_or_none()
|
|
|
|
|
|
async def list_results(
|
|
db: AsyncSession,
|
|
vehicle_id: uuid.UUID | None = None,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
) -> tuple[list[RetouchResult], int]:
|
|
"""List retouch results with optional vehicle filter and pagination."""
|
|
count_stmt = select(func.count(RetouchResult.id))
|
|
data_stmt = select(RetouchResult)
|
|
|
|
if vehicle_id is not None:
|
|
count_stmt = count_stmt.where(RetouchResult.vehicle_id == vehicle_id)
|
|
data_stmt = data_stmt.where(RetouchResult.vehicle_id == vehicle_id)
|
|
|
|
total_result = await db.execute(count_stmt)
|
|
total = total_result.scalar_one()
|
|
|
|
offset = (page - 1) * page_size
|
|
data_stmt = data_stmt.order_by(RetouchResult.created_at.desc()).offset(offset).limit(page_size)
|
|
result = await db.execute(data_stmt)
|
|
items = list(result.scalars().all())
|
|
|
|
return items, total
|
|
|
|
|
|
async def process_retouch(
|
|
db: AsyncSession,
|
|
result_id: uuid.UUID,
|
|
vehicle_info: dict[str, Any] | None = None,
|
|
) -> RetouchResult:
|
|
"""Process a retouch result by sending the image to Flux.1-Pro via OpenRouter.
|
|
|
|
Updates the result status to processing, calls the API, saves the
|
|
retouched image, and sets status to completed. On failure, sets
|
|
status to failed with an error message.
|
|
"""
|
|
result = await get_result(db, result_id)
|
|
if result is None:
|
|
raise ValueError(f"Retouch result {result_id} not found")
|
|
|
|
result.status = RetouchStatus.processing.value
|
|
await db.flush()
|
|
|
|
try:
|
|
with open(result.original_file_path, "rb") as f:
|
|
image_bytes = f.read()
|
|
|
|
retouched_bytes = await _call_flux_pro(
|
|
image_bytes=image_bytes,
|
|
mime_type=result.mime_type,
|
|
prompt=generate_retouch_prompt(vehicle_info),
|
|
)
|
|
|
|
upload_dir = settings.UPLOAD_DIR
|
|
ext = _get_file_extension(result.mime_type)
|
|
retouched_filename = f"retouched_{result_id.hex}{ext}"
|
|
retouched_path = os.path.join(upload_dir, retouched_filename)
|
|
|
|
with open(retouched_path, "wb") as f:
|
|
f.write(retouched_bytes)
|
|
|
|
result.retouched_file_path = retouched_path
|
|
result.status = RetouchStatus.completed.value
|
|
result.error_message = None
|
|
await db.flush()
|
|
await db.refresh(result)
|
|
logger.info("Retouch completed for result %s", result_id)
|
|
return result
|
|
|
|
except Exception as exc:
|
|
logger.error("Retouch failed for result %s: %s", result_id, exc)
|
|
result.status = RetouchStatus.failed.value
|
|
result.error_message = str(exc)
|
|
await db.flush()
|
|
await db.refresh(result)
|
|
return result
|
|
|
|
|
|
async def _call_flux_pro(
|
|
image_bytes: bytes,
|
|
mime_type: str,
|
|
prompt: str,
|
|
api_key: str | None = None,
|
|
model: str | None = None,
|
|
) -> bytes:
|
|
"""Send image to Flux.1-Pro via OpenRouter for retouching.
|
|
|
|
Returns the retouched image bytes.
|
|
Raises httpx.HTTPStatusError or ValueError on failure.
|
|
"""
|
|
key = api_key or settings.OPENROUTER_API_KEY
|
|
if not key:
|
|
raise ValueError("OPENROUTER_API_KEY is not configured")
|
|
|
|
model_name = model or "black-forest-labs/flux-1-pro"
|
|
b64_image = base64.b64encode(image_bytes).decode("utf-8")
|
|
data_uri = f"data:{mime_type};base64,{b64_image}"
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload: dict[str, Any] = {
|
|
"model": model_name,
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": RETOUCH_SYSTEM_PROMPT,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": prompt,
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": data_uri},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
"temperature": 0.7,
|
|
"max_tokens": 4096,
|
|
}
|
|
|
|
base_url = settings.OPENROUTER_BASE_URL.rstrip("/")
|
|
url = f"{base_url}/chat/completions"
|
|
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
|
|
response = await client.post(url, headers=headers, json=payload)
|
|
response.raise_for_status()
|
|
|
|
body = response.json()
|
|
content = body.get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
|
|
# The model may return a base64-encoded image or a URL
|
|
if isinstance(content, list):
|
|
for part in content:
|
|
if isinstance(part, dict) and part.get("type") == "image_url":
|
|
url_or_data = part.get("image_url", {}).get("url", "")
|
|
if url_or_data.startswith("data:"):
|
|
b64_data = url_or_data.split(",", 1)[1]
|
|
return base64.b64decode(b64_data)
|
|
elif url_or_data.startswith("http"):
|
|
async with httpx.AsyncClient(timeout=60.0) as dl_client:
|
|
dl_resp = await dl_client.get(url_or_data)
|
|
dl_resp.raise_for_status()
|
|
return dl_resp.content
|
|
|
|
if isinstance(content, str) and content.startswith("data:"):
|
|
b64_data = content.split(",", 1)[1]
|
|
return base64.b64decode(b64_data)
|
|
|
|
# Fallback: if no image returned, return original bytes
|
|
logger.warning("Flux.1-Pro did not return an image, returning original bytes")
|
|
return image_bytes
|