Claude
b2d064573f
Matches the existing vendor/easy-asterisk convention (used by services/asterisk.sh) instead of two one-off top-level directories that cluttered the repo root and didn't look like anything else next to setup.sh, lib/, services/, extras/. Only the two services' own SRC_DIR path resolution and header comments needed updating — nothing else in the repo referenced the old ./ai-stack / ./paintplus paths. Also documents vendor/ in README.md's Layout section.
990 lines
36 KiB
Python
990 lines
36 KiB
Python
"""
|
||
AI tools router — LaMa inpaint, background removal, remote generation, config.
|
||
All endpoints are under /api prefix.
|
||
"""
|
||
|
||
from fastapi import APIRouter, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from pydantic import BaseModel
|
||
from typing import Optional
|
||
import base64
|
||
import asyncio
|
||
import json
|
||
from io import BytesIO
|
||
|
||
from app.services.local_inpaint import (
|
||
lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available,
|
||
)
|
||
|
||
router = APIRouter(prefix="/api", tags=["ai-tools"])
|
||
|
||
|
||
# ─── Request / response models ───────────────────────────────────────────────
|
||
|
||
class EraseRequest(BaseModel):
|
||
image: str # base64
|
||
mask: str # base64
|
||
|
||
|
||
class InpaintRemoteRequest(BaseModel):
|
||
image: str
|
||
mask: str
|
||
prompt: str
|
||
negative_prompt: Optional[str] = ""
|
||
steps: Optional[int] = 30
|
||
cfg_scale: Optional[float] = 7.5
|
||
model: Optional[str] = None
|
||
|
||
|
||
class Txt2ImgRequest(BaseModel):
|
||
prompt: str
|
||
width: Optional[int] = 1024
|
||
height: Optional[int] = 1024
|
||
negative_prompt: Optional[str] = ""
|
||
steps: Optional[int] = 30
|
||
cfg_scale: Optional[float] = 7.5
|
||
model: Optional[str] = None
|
||
seed: Optional[int] = 0
|
||
|
||
|
||
class Img2ImgRequest(BaseModel):
|
||
image: str
|
||
prompt: str
|
||
strength: Optional[float] = 0.75
|
||
negative_prompt: Optional[str] = ""
|
||
steps: Optional[int] = 30
|
||
cfg_scale: Optional[float] = 7.5
|
||
model: Optional[str] = None
|
||
|
||
|
||
class OutpaintRequest(BaseModel):
|
||
image: str
|
||
direction: str # left | right | top | bottom
|
||
size: Optional[int] = 256
|
||
prompt: Optional[str] = ""
|
||
|
||
|
||
class BgRemoveRequest(BaseModel):
|
||
image: str
|
||
|
||
|
||
# ─── Helpers ─────────────────────────────────────────────────────────────────
|
||
|
||
def _decode(b64: str) -> bytes:
|
||
return base64.b64decode(b64)
|
||
|
||
|
||
def _encode(data: bytes) -> str:
|
||
return base64.b64encode(data).decode()
|
||
|
||
|
||
def _require_remote(operation: str = None):
|
||
from app.services.remote_provider import get_remote_provider
|
||
from app.config import settings
|
||
provider = get_remote_provider(operation)
|
||
if provider is None:
|
||
if (settings.ai_provider or "").lower() == "local_gpu":
|
||
raise HTTPException(
|
||
status_code=503,
|
||
detail=(
|
||
"local_gpu provider failed to load — diffusers may be incompatible with "
|
||
"the installed PyTorch version. Check container logs for details. "
|
||
"If you see 'torch has no attribute xpu', rebuild the container from the "
|
||
"correct branch so the pinned diffusers<0.29.0 is installed."
|
||
)
|
||
)
|
||
op_hint = f"AI_PROVIDER_{operation.upper()} or " if operation else ""
|
||
raise HTTPException(
|
||
status_code=503,
|
||
detail=f"No remote AI provider configured for '{operation or 'default'}'. "
|
||
f"Set {op_hint}AI_PROVIDER in .env (openai / invokeai / comfyui)."
|
||
)
|
||
return provider
|
||
|
||
|
||
# ─── Local inpaint endpoints ─────────────────────────────────────────────────
|
||
|
||
@router.post("/erase")
|
||
async def erase(req: EraseRequest):
|
||
"""
|
||
Magic eraser: remove object / fill region using LaMa (local, no API key needed).
|
||
Falls back to OpenCV if LaMa not installed.
|
||
"""
|
||
try:
|
||
image_bytes = _decode(req.image)
|
||
mask_bytes = _decode(req.mask)
|
||
|
||
if lama_available():
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, lama_inpaint, image_bytes, mask_bytes
|
||
)
|
||
method = "lama"
|
||
else:
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, opencv_inpaint, image_bytes, mask_bytes
|
||
)
|
||
method = "opencv"
|
||
|
||
return {"result": _encode(result), "method": method}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/inpaint/lama")
|
||
async def inpaint_lama(req: EraseRequest):
|
||
"""LaMa structural inpainting."""
|
||
if not lama_available():
|
||
raise HTTPException(status_code=503, detail="simple-lama-inpainting not installed.")
|
||
try:
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, lama_inpaint, _decode(req.image), _decode(req.mask)
|
||
)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/inpaint/fast")
|
||
async def inpaint_fast(req: EraseRequest):
|
||
"""OpenCV fast inpainting (CPU, milliseconds)."""
|
||
try:
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, opencv_inpaint, _decode(req.image), _decode(req.mask)
|
||
)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/background/remove")
|
||
async def background_remove(req: BgRemoveRequest):
|
||
"""Remove background — rembg if available, else U2Net."""
|
||
try:
|
||
image_bytes = _decode(req.image)
|
||
|
||
# Try rembg first
|
||
if rembg_available():
|
||
from app.services.local_inpaint import remove_background_rembg
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, remove_background_rembg, image_bytes
|
||
)
|
||
return {"result": _encode(result), "method": "rembg"}
|
||
|
||
# Fall back to U2Net (existing implementation)
|
||
from PIL import Image
|
||
from io import BytesIO as _BytesIO
|
||
img = Image.open(_BytesIO(image_bytes)).convert("RGB")
|
||
from app.routers.tools import _remove_background_u2net
|
||
result = await _remove_background_u2net(img)
|
||
return {"result": _encode(result), "method": "u2net"}
|
||
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# ─── Remote provider endpoints ───────────────────────────────────────────────
|
||
|
||
@router.post("/inpaint/remote")
|
||
async def inpaint_remote(req: InpaintRemoteRequest):
|
||
"""Inpaint via configured remote provider (InvokeAI / ComfyUI / OpenAI)."""
|
||
provider = _require_remote("inpaint")
|
||
try:
|
||
params = {
|
||
"negative_prompt": req.negative_prompt or "",
|
||
"steps": req.steps,
|
||
"cfg_scale": req.cfg_scale,
|
||
}
|
||
if req.model:
|
||
params["model"] = req.model
|
||
result = await provider.inpaint(_decode(req.image), _decode(req.mask), req.prompt, params)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.get("/generate/progress")
|
||
async def generation_progress_stream():
|
||
"""
|
||
SSE stream of local GPU pipeline inference progress.
|
||
Events are JSON arrays of pipeline state objects, emitted every 200 ms.
|
||
Each object: {pipeline, state, step, total_steps, progress, message, model_id, …}
|
||
Clients open this with EventSource before firing a generation POST,
|
||
then close it when the POST resolves.
|
||
"""
|
||
from app.services.local_diffusion import get_all_model_states
|
||
|
||
async def event_gen():
|
||
try:
|
||
while True:
|
||
states = get_all_model_states()
|
||
yield f"data: {json.dumps(states)}\n\n"
|
||
await asyncio.sleep(0.2)
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
return StreamingResponse(
|
||
event_gen(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no",
|
||
"Connection": "keep-alive",
|
||
},
|
||
)
|
||
|
||
|
||
@router.post("/generate/txt2img")
|
||
async def txt2img(req: Txt2ImgRequest):
|
||
"""Text-to-image via configured remote provider."""
|
||
provider = _require_remote("txt2img")
|
||
try:
|
||
params = {
|
||
"negative_prompt": req.negative_prompt or "",
|
||
"steps": req.steps,
|
||
"cfg_scale": req.cfg_scale,
|
||
"seed": req.seed or 0,
|
||
}
|
||
if req.model:
|
||
params["model"] = req.model
|
||
result = await provider.txt2img(req.prompt, req.width, req.height, params)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/generate/img2img")
|
||
async def img2img(req: Img2ImgRequest):
|
||
"""Image-to-image via configured remote provider."""
|
||
provider = _require_remote("img2img")
|
||
try:
|
||
params = {
|
||
"negative_prompt": req.negative_prompt or "",
|
||
"steps": req.steps,
|
||
"cfg_scale": req.cfg_scale,
|
||
}
|
||
if req.model:
|
||
params["model"] = req.model
|
||
result = await provider.img2img(_decode(req.image), req.prompt, req.strength, params)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/generate/outpaint")
|
||
async def outpaint(req: OutpaintRequest):
|
||
"""Expand canvas in given direction via remote provider."""
|
||
provider = _require_remote("outpaint")
|
||
if req.direction not in ("left", "right", "top", "bottom"):
|
||
raise HTTPException(status_code=400, detail="direction must be left/right/top/bottom")
|
||
try:
|
||
result = await provider.outpaint(_decode(req.image), req.direction, req.size, req.prompt or "")
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# ─── Config / capabilities ────────────────────────────────────────────────────
|
||
|
||
class ConfigUpdateRequest(BaseModel):
|
||
ai_provider: Optional[str] = None
|
||
# Per-operation overrides (blank = use default)
|
||
ai_provider_inpaint: Optional[str] = None
|
||
ai_provider_txt2img: Optional[str] = None
|
||
ai_provider_img2img: Optional[str] = None
|
||
ai_provider_outpaint: Optional[str] = None
|
||
# Credentials / URLs
|
||
openai_api_key: Optional[str] = None
|
||
openai_model: Optional[str] = None
|
||
invokeai_url: Optional[str] = None
|
||
invokeai_default_model: Optional[str] = None
|
||
comfyui_url: Optional[str] = None
|
||
comfyui_default_model: Optional[str] = None
|
||
replicate_api_key: Optional[str] = None
|
||
stability_api_key: Optional[str] = None
|
||
|
||
|
||
@router.post("/config")
|
||
async def update_config(req: ConfigUpdateRequest):
|
||
"""
|
||
Apply runtime provider settings (no restart needed).
|
||
Values are applied to the live settings object in-process.
|
||
They do NOT persist across restarts — set them in .env for permanence.
|
||
"""
|
||
from app.config import settings
|
||
|
||
_str_fields = [
|
||
"ai_provider", "ai_provider_inpaint", "ai_provider_txt2img",
|
||
"ai_provider_img2img", "ai_provider_outpaint",
|
||
"openai_api_key", "openai_model",
|
||
"invokeai_url", "invokeai_default_model",
|
||
"comfyui_url", "comfyui_default_model",
|
||
"replicate_api_key", "stability_api_key",
|
||
]
|
||
for field in _str_fields:
|
||
val = getattr(req, field, None)
|
||
if val is not None:
|
||
setattr(settings, field, val)
|
||
|
||
return {
|
||
"status": "ok",
|
||
"ai_provider": settings.ai_provider,
|
||
"overrides": {
|
||
"inpaint": settings.ai_provider_inpaint or None,
|
||
"txt2img": settings.ai_provider_txt2img or None,
|
||
"img2img": settings.ai_provider_img2img or None,
|
||
"outpaint": settings.ai_provider_outpaint or None,
|
||
}
|
||
}
|
||
|
||
|
||
async def _check_provider(operation: str) -> dict:
|
||
"""Health-check the provider for a specific operation."""
|
||
from app.services.remote_provider import get_remote_provider
|
||
try:
|
||
p = get_remote_provider(operation)
|
||
if p is None:
|
||
return {"provider": None, "healthy": False}
|
||
healthy = await asyncio.wait_for(p.health(), timeout=5.0)
|
||
return {"provider": p.__class__.__name__.replace("Provider", "").lower(), "healthy": healthy}
|
||
except Exception:
|
||
return {"provider": None, "healthy": False}
|
||
|
||
|
||
@router.get("/config")
|
||
async def get_config():
|
||
"""
|
||
Return capability flags so the frontend can show/hide tools.
|
||
Includes per-operation provider assignments and health status.
|
||
"""
|
||
from app.config import settings
|
||
|
||
# Run health checks for each operation concurrently
|
||
ops = ["inpaint", "txt2img", "img2img", "outpaint"]
|
||
results = await asyncio.gather(*[_check_provider(op) for op in ops])
|
||
op_status = dict(zip(ops, results))
|
||
|
||
# Default provider for display (used when no per-op override)
|
||
default_name = (settings.ai_provider or "").lower() or None
|
||
|
||
from app.services.gpu_detect import get_cached_gpu_info
|
||
gpu_info = get_cached_gpu_info()
|
||
|
||
return {
|
||
"local": {
|
||
"lama": lama_available(),
|
||
"rembg": rembg_available(),
|
||
"opencv": True,
|
||
"gpu_detected": gpu_available(),
|
||
"gpu_backend": gpu_info.backend,
|
||
"gpu_device": gpu_info.device_name,
|
||
"gpu_vram_total": gpu_info.vram_total_gb,
|
||
"gpu_vram_free": gpu_info.vram_free_gb,
|
||
"gpu_cc": gpu_info.compute_capability,
|
||
"gpu_fp16": gpu_info.fp16,
|
||
"gpu_bf16": gpu_info.bf16,
|
||
"gpu_fp8": gpu_info.fp8,
|
||
"gpu_tensor_cores": gpu_info.tensor_cores,
|
||
"gpu_tier": gpu_info.tier,
|
||
"gpu_eff_vram": gpu_info.effective_vram_gb,
|
||
"local_gpu_available": gpu_info.backend in ("cuda", "mps"),
|
||
"local_gpu_capabilities": gpu_info.capabilities,
|
||
"local_gpu_warnings": gpu_info.warnings,
|
||
},
|
||
"remote": {
|
||
"default_provider": default_name,
|
||
# Legacy field kept for backwards compat with badge/capabilities checks
|
||
"provider": default_name,
|
||
"healthy": any(v["healthy"] for v in op_status.values()),
|
||
"operations": op_status,
|
||
"overrides": {
|
||
"inpaint": settings.ai_provider_inpaint or None,
|
||
"txt2img": settings.ai_provider_txt2img or None,
|
||
"img2img": settings.ai_provider_img2img or None,
|
||
"outpaint": settings.ai_provider_outpaint or None,
|
||
},
|
||
}
|
||
}
|
||
|
||
|
||
# ─── Selection image operations ─────────────────────────────────────────────
|
||
|
||
class ScaleSelectionRequest(BaseModel):
|
||
image: str # base64 full canvas
|
||
mask: str # base64 selection mask (white = object)
|
||
scale_pct: float = 103.0 # 103 = 3% bigger, 95 = 5% smaller
|
||
|
||
|
||
class AiEditRegionRequest(BaseModel):
|
||
image: str
|
||
mask: str
|
||
instruction: str
|
||
negative_prompt: str = ""
|
||
steps: int = 30
|
||
cfg_scale: float = 7.5
|
||
|
||
|
||
class PasteIntoSelectionRequest(BaseModel):
|
||
image: str # base64 target canvas
|
||
mask: str # base64 selection mask
|
||
paste_image: str # base64 image to paste
|
||
|
||
|
||
@router.post("/image/scale-selection")
|
||
async def scale_selection(req: ScaleSelectionRequest):
|
||
"""
|
||
Scale the object selected by mask by scale_pct%, AI-fill the exposed gap.
|
||
Works purely with local tools (LaMa/OpenCV) — no remote provider needed.
|
||
"""
|
||
try:
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter
|
||
except ImportError:
|
||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||
|
||
img = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||
if img.size != mask.size:
|
||
mask = mask.resize(img.size, Image.LANCZOS)
|
||
|
||
mask_arr = np.array(mask)
|
||
ys, xs = np.where(mask_arr > 128)
|
||
if len(xs) == 0:
|
||
raise HTTPException(status_code=400, detail="Empty mask — nothing to scale")
|
||
|
||
minx, maxx = int(xs.min()), int(xs.max())
|
||
miny, maxy = int(ys.min()), int(ys.max())
|
||
cx, cy = (minx + maxx) / 2.0, (miny + maxy) / 2.0
|
||
obj_w, obj_h = maxx - minx + 1, maxy - miny + 1
|
||
|
||
scale = req.scale_pct / 100.0
|
||
new_w = max(1, round(obj_w * scale))
|
||
new_h = max(1, round(obj_h * scale))
|
||
|
||
# Extract masked object crop (RGBA with mask as alpha)
|
||
img_rgba = img.convert("RGBA")
|
||
obj_crop = img_rgba.crop((minx, miny, maxx + 1, maxy + 1))
|
||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||
r, g, b, _ = obj_crop.split()
|
||
obj_masked = Image.merge("RGBA", (r, g, b, mask_crop))
|
||
scaled_obj = obj_masked.resize((new_w, new_h), Image.LANCZOS)
|
||
|
||
# AI-fill the original mask area (gap) with LaMa/OpenCV
|
||
gap_mask = mask.filter(ImageFilter.MaxFilter(9)) # expand ~4px for clean seam
|
||
gap_bytes = BytesIO()
|
||
img.save(gap_bytes, format="PNG")
|
||
gap_mask_bytes = BytesIO()
|
||
gap_mask.save(gap_mask_bytes, format="PNG")
|
||
|
||
try:
|
||
if lama_available():
|
||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||
None, lama_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||
)
|
||
else:
|
||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||
None, opencv_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||
)
|
||
filled = Image.open(BytesIO(filled_bytes)).convert("RGBA")
|
||
except Exception as exc:
|
||
print(f"[scale-selection] fill fallback: {exc}")
|
||
filled = img.convert("RGBA")
|
||
|
||
# Paste scaled object centered on original centroid
|
||
px = round(cx - new_w / 2)
|
||
py = round(cy - new_h / 2)
|
||
result = filled.copy()
|
||
result.paste(scaled_obj, (px, py), scaled_obj.split()[3])
|
||
|
||
out = BytesIO()
|
||
result.convert("RGB").save(out, format="PNG")
|
||
return {"result": _encode(out.getvalue())}
|
||
|
||
|
||
@router.post("/image/ai-edit-region")
|
||
async def ai_edit_region(req: AiEditRegionRequest):
|
||
"""
|
||
AI-edit the selected region using the configured inpaint provider.
|
||
Works with local_gpu, InvokeAI, ComfyUI, or OpenAI.
|
||
"""
|
||
provider = _require_remote("inpaint")
|
||
try:
|
||
result_bytes = await provider.inpaint(
|
||
_decode(req.image),
|
||
_decode(req.mask),
|
||
req.instruction,
|
||
{"negative_prompt": req.negative_prompt, "steps": req.steps, "cfg_scale": req.cfg_scale},
|
||
)
|
||
except Exception as exc:
|
||
import traceback; traceback.print_exc()
|
||
msg = str(exc)
|
||
if "Errno -3" in msg or "Name or service not known" in msg or "ConnectError" in msg:
|
||
raise HTTPException(
|
||
status_code=503,
|
||
detail=(
|
||
"AI model files not yet downloaded — container DNS appears to be blocked. "
|
||
"Fix: sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT on the host, "
|
||
"or pre-download the model: pip install huggingface-hub && "
|
||
"huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 "
|
||
"--cache-dir ./data/hf_cache"
|
||
)
|
||
)
|
||
raise HTTPException(status_code=500, detail=msg)
|
||
return {"result": _encode(result_bytes)}
|
||
|
||
|
||
@router.post("/image/paste-into-selection")
|
||
async def paste_into_selection(req: PasteIntoSelectionRequest):
|
||
"""
|
||
Scale a clipboard image to the selection bounding box, mask it to the
|
||
selection shape, and composite it over the original canvas.
|
||
"""
|
||
try:
|
||
import numpy as np
|
||
from PIL import Image
|
||
except ImportError:
|
||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||
|
||
img = Image.open(BytesIO(_decode(req.image))).convert("RGBA")
|
||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||
paste_img = Image.open(BytesIO(_decode(req.paste_image))).convert("RGBA")
|
||
|
||
if img.size != mask.size:
|
||
mask = mask.resize(img.size, Image.LANCZOS)
|
||
|
||
mask_arr = np.array(mask)
|
||
ys, xs = np.where(mask_arr > 128)
|
||
if len(xs) == 0:
|
||
raise HTTPException(status_code=400, detail="Empty mask")
|
||
|
||
minx, maxx = int(xs.min()), int(xs.max())
|
||
miny, maxy = int(ys.min()), int(ys.max())
|
||
target_w = maxx - minx + 1
|
||
target_h = maxy - miny + 1
|
||
|
||
# Scale clipboard image to fit the selection bounding box
|
||
paste_scaled = paste_img.resize((target_w, target_h), Image.LANCZOS)
|
||
|
||
# Clip paste to selection shape using mask
|
||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||
r, g, b, a = paste_scaled.split()
|
||
mask_np = np.array(mask_crop)
|
||
alpha_np = np.array(a)
|
||
combined = (alpha_np.astype(np.uint16) * mask_np.astype(np.uint16) // 255).astype(np.uint8)
|
||
paste_final = Image.merge("RGBA", (r, g, b, Image.fromarray(combined)))
|
||
|
||
result = img.copy()
|
||
result.paste(paste_final, (minx, miny), paste_final.split()[3])
|
||
|
||
out = BytesIO()
|
||
result.convert("RGB").save(out, format="PNG")
|
||
return {"result": _encode(out.getvalue())}
|
||
|
||
|
||
# ─── SAM (Segment Anything) ──────────────────────────────────────────────────
|
||
|
||
class SegmentPointRequest(BaseModel):
|
||
image: str # base64 PNG/JPEG
|
||
points: list[list[int]] # [[x, y], ...] original image coords
|
||
labels: list[int] # 1=include, 0=exclude — same length as points
|
||
|
||
|
||
@router.post("/segment/point")
|
||
async def segment_point(req: SegmentPointRequest):
|
||
"""
|
||
Run SAM point-prompt segmentation.
|
||
Returns a binary mask PNG (white = selected area).
|
||
Auto-downloads the SAM ViT-B model (~375 MB) on first call.
|
||
"""
|
||
if not req.points:
|
||
raise HTTPException(status_code=400, detail="At least one point required.")
|
||
if len(req.points) != len(req.labels):
|
||
raise HTTPException(status_code=400, detail="points and labels must have the same length.")
|
||
|
||
try:
|
||
image_bytes = base64.b64decode(req.image)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||
|
||
from app.services.sam_service import predict_points, get_install_status
|
||
try:
|
||
mask_bytes = await predict_points(
|
||
image_bytes,
|
||
[tuple(p) for p in req.points],
|
||
req.labels,
|
||
)
|
||
return {
|
||
"mask": base64.b64encode(mask_bytes).decode(),
|
||
"sam_install": get_install_status(),
|
||
}
|
||
except RuntimeError as e:
|
||
raise HTTPException(status_code=503, detail=str(e))
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.get("/segment/install-status")
|
||
def segment_install_status():
|
||
"""Poll SAM model download progress."""
|
||
from app.services.sam_service import get_install_status, sam_model_available
|
||
status = get_install_status()
|
||
status["model_ready"] = sam_model_available()
|
||
return status
|
||
|
||
|
||
@router.post("/segment/install")
|
||
async def segment_install():
|
||
"""Trigger SAM model download explicitly (also auto-triggered on first /segment/point call)."""
|
||
from app.services.sam_service import ensure_sam_installed, get_install_status
|
||
asyncio.create_task(ensure_sam_installed())
|
||
return get_install_status()
|
||
|
||
|
||
# ─── Enhance ─────────────────────────────────────────────────────────────────
|
||
|
||
import io as _io
|
||
import numpy as _np
|
||
import cv2 as _cv2
|
||
from PIL import Image as _Image
|
||
|
||
class EnhanceRequest(BaseModel):
|
||
image: str # base64
|
||
strength: float = 1.0
|
||
|
||
|
||
def _enhance_image(image_bytes: bytes, strength: float) -> bytes:
|
||
"""
|
||
Apply a chain of non-AI image enhancements, each blended with `strength` (0–1).
|
||
|
||
Steps:
|
||
1. Auto white balance (gray-world)
|
||
2. CLAHE on L channel of LAB colorspace
|
||
3. Auto saturation boost in HSV (×1.15, clamped)
|
||
4. Mild unsharp mask (gaussian sigma=1.0, delta weight=0.3)
|
||
"""
|
||
strength = max(0.0, min(1.0, float(strength)))
|
||
|
||
# Decode to RGB numpy array
|
||
pil = _Image.open(_io.BytesIO(image_bytes)).convert("RGB")
|
||
orig = _np.array(pil, dtype=_np.float32) # H×W×3, float [0,255]
|
||
|
||
img = orig.copy()
|
||
|
||
# ── Step 1: Auto white balance (gray-world) ──────────────────────────────
|
||
mean_r = img[:, :, 0].mean()
|
||
mean_g = img[:, :, 1].mean()
|
||
mean_b = img[:, :, 2].mean()
|
||
overall_mean = (mean_r + mean_g + mean_b) / 3.0
|
||
|
||
def _scale(channel, channel_mean):
|
||
if channel_mean == 0:
|
||
return channel
|
||
return channel * (overall_mean / channel_mean)
|
||
|
||
wb = img.copy()
|
||
wb[:, :, 0] = _np.clip(_scale(img[:, :, 0], mean_r), 0, 255)
|
||
wb[:, :, 1] = _np.clip(_scale(img[:, :, 1], mean_g), 0, 255)
|
||
wb[:, :, 2] = _np.clip(_scale(img[:, :, 2], mean_b), 0, 255)
|
||
|
||
img = (orig + strength * (wb - orig)).clip(0, 255)
|
||
|
||
# ── Step 2: CLAHE on L channel (LAB) ────────────────────────────────────
|
||
img_u8 = img.astype(_np.uint8)
|
||
lab = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2LAB)
|
||
clahe = _cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||
l_orig = lab[:, :, 0].copy()
|
||
lab[:, :, 0] = clahe.apply(l_orig)
|
||
# Blend L channel back using strength
|
||
lab_blended = lab.copy()
|
||
lab_blended[:, :, 0] = (l_orig + strength * (lab[:, :, 0].astype(_np.float32) - l_orig.astype(_np.float32))).clip(0, 255).astype(_np.uint8)
|
||
img = _cv2.cvtColor(lab_blended, _cv2.COLOR_LAB2RGB).astype(_np.float32)
|
||
|
||
# ── Step 3: Auto saturation boost (HSV, ×1.15) ──────────────────────────
|
||
img_u8 = img.astype(_np.uint8)
|
||
hsv = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2HSV).astype(_np.float32)
|
||
s_orig = hsv[:, :, 1].copy()
|
||
s_boosted = _np.clip(s_orig * 1.15, 0, 255)
|
||
hsv[:, :, 1] = s_orig + strength * (s_boosted - s_orig)
|
||
hsv = hsv.clip(0, 255).astype(_np.uint8)
|
||
img = _cv2.cvtColor(hsv, _cv2.COLOR_HSV2RGB).astype(_np.float32)
|
||
|
||
# ── Step 4: Mild unsharp mask (sigma=1.0, delta weight=0.3) ─────────────
|
||
img_u8 = img.astype(_np.uint8)
|
||
blurred = _cv2.GaussianBlur(img_u8, (0, 0), sigmaX=1.0)
|
||
sharpness_delta = img_u8.astype(_np.float32) - blurred.astype(_np.float32)
|
||
sharpened = img_u8.astype(_np.float32) + 0.3 * sharpness_delta * strength
|
||
img = sharpened.clip(0, 255)
|
||
|
||
# Encode result as PNG
|
||
result_pil = _Image.fromarray(img.astype(_np.uint8), mode="RGB")
|
||
buf = _io.BytesIO()
|
||
result_pil.save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
|
||
|
||
@router.post("/enhance")
|
||
async def enhance(req: EnhanceRequest):
|
||
"""
|
||
Non-AI image enhancement: auto white balance, CLAHE, saturation boost,
|
||
and unsharp mask. Each step is blended proportionally to `strength` (0–1).
|
||
"""
|
||
try:
|
||
image_bytes = _decode(req.image)
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, _enhance_image, image_bytes, req.strength
|
||
)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# ─── Subject replace ─────────────────────────────────────────────────────────
|
||
|
||
class ExtractSubjectRequest(BaseModel):
|
||
image: str # base64
|
||
|
||
|
||
class ReplaceSubjectRequest(BaseModel):
|
||
background_image: str # base64 — image whose background we keep
|
||
subject_image: str # base64 — image whose subject we extract
|
||
mask: Optional[str] = None # base64 — white = where the subject should land
|
||
match_colors: bool = True # blend subject color stats toward background
|
||
|
||
|
||
def _extract_subject_bytes(image_bytes: bytes) -> bytes:
|
||
"""Remove background from image using rembg; return RGBA PNG bytes."""
|
||
if rembg_available():
|
||
return remove_background_rembg(image_bytes)
|
||
raise RuntimeError(
|
||
"rembg is not installed. Run: pip install rembg (or add it to requirements.txt)"
|
||
)
|
||
|
||
|
||
def _color_transfer_lab(subj_rgba: "Image", bg_rgb: "Image", blend: float = 0.45) -> "Image":
|
||
"""
|
||
Partial LAB color transfer: nudge subject color statistics 'blend' fraction
|
||
toward the background's statistics so it looks like it belongs in the scene.
|
||
"""
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
src_arr = np.array(subj_rgba.convert("RGB"), dtype=np.float32)
|
||
tgt_arr = np.array(bg_rgb.convert("RGB"), dtype=np.float32)
|
||
|
||
alpha = np.array(subj_rgba.split()[3])
|
||
subject_mask = alpha > 10
|
||
|
||
if not subject_mask.any():
|
||
return subj_rgba
|
||
|
||
src_lab = cv2.cvtColor(src_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
tgt_lab = cv2.cvtColor(tgt_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||
|
||
for ch in range(3):
|
||
src_ch = src_lab[:, :, ch]
|
||
src_pixels = src_ch[subject_mask]
|
||
tgt_pixels = tgt_lab[:, :, ch].flatten()
|
||
|
||
src_mean, src_std = float(src_pixels.mean()), float(src_pixels.std()) + 1e-6
|
||
tgt_mean, tgt_std = float(tgt_pixels.mean()), float(tgt_pixels.std()) + 1e-6
|
||
|
||
adjusted_std = src_std + blend * (tgt_std - src_std)
|
||
adjusted = (src_ch - src_mean) * (adjusted_std / src_std) + src_mean + blend * (tgt_mean - src_mean)
|
||
src_lab[:, :, ch] = np.clip(adjusted, 0, 255)
|
||
|
||
result_rgb = cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
||
r, g, b = result_rgb[:, :, 0], result_rgb[:, :, 1], result_rgb[:, :, 2]
|
||
return Image.merge("RGBA", [
|
||
Image.fromarray(r), Image.fromarray(g),
|
||
Image.fromarray(b), Image.fromarray(alpha),
|
||
])
|
||
|
||
|
||
def _do_replace_subject(
|
||
bg_bytes: bytes,
|
||
subj_bytes: bytes,
|
||
mask_bytes: Optional[bytes],
|
||
match_colors: bool,
|
||
) -> bytes:
|
||
"""Core compositing: extract subject → scale → color-match → paste onto background."""
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
bg_img = Image.open(BytesIO(bg_bytes)).convert("RGBA")
|
||
|
||
subj_rgba = Image.open(BytesIO(_extract_subject_bytes(subj_bytes))).convert("RGBA")
|
||
|
||
# Determine target placement bounding box from mask or full canvas
|
||
if mask_bytes:
|
||
mask_img = Image.open(BytesIO(mask_bytes)).convert("L")
|
||
if mask_img.size != bg_img.size:
|
||
mask_img = mask_img.resize(bg_img.size, Image.LANCZOS)
|
||
mask_arr = np.array(mask_img)
|
||
ys, xs = np.where(mask_arr > 128)
|
||
else:
|
||
mask_img = None
|
||
mask_arr = None
|
||
ys, xs = np.array([]), np.array([])
|
||
|
||
if len(xs) > 0:
|
||
minx, maxx = int(xs.min()), int(xs.max())
|
||
miny, maxy = int(ys.min()), int(ys.max())
|
||
else:
|
||
minx, miny = 0, 0
|
||
maxx, maxy = bg_img.width - 1, bg_img.height - 1
|
||
|
||
target_w = maxx - minx + 1
|
||
target_h = maxy - miny + 1
|
||
|
||
# Scale subject to fit target area, preserving aspect ratio
|
||
sw, sh = subj_rgba.size
|
||
scale = min(target_w / sw, target_h / sh)
|
||
new_w = max(1, round(sw * scale))
|
||
new_h = max(1, round(sh * scale))
|
||
subj_scaled = subj_rgba.resize((new_w, new_h), Image.LANCZOS)
|
||
|
||
# Optional color transfer to blend lighting/tone
|
||
if match_colors:
|
||
subj_scaled = _color_transfer_lab(subj_scaled, bg_img.convert("RGB"))
|
||
|
||
# Center in target area
|
||
px = minx + (target_w - new_w) // 2
|
||
py = miny + (target_h - new_h) // 2
|
||
|
||
result = bg_img.copy()
|
||
|
||
if mask_img is not None and len(xs) > 0:
|
||
# Build a full-canvas RGBA layer for the subject
|
||
subj_canvas = Image.new("RGBA", bg_img.size, (0, 0, 0, 0))
|
||
subj_canvas.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||
# Clip subject's alpha to the selection mask
|
||
sc_arr = np.array(subj_canvas)
|
||
sc_arr[:, :, 3] = np.minimum(sc_arr[:, :, 3], mask_arr).astype(np.uint8)
|
||
subj_canvas = Image.fromarray(sc_arr)
|
||
result.paste(subj_canvas, (0, 0), subj_canvas.split()[3])
|
||
else:
|
||
result.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||
|
||
out = BytesIO()
|
||
result.convert("RGB").save(out, format="PNG")
|
||
return out.getvalue()
|
||
|
||
|
||
@router.post("/image/extract-subject")
|
||
async def extract_subject(req: ExtractSubjectRequest):
|
||
"""
|
||
Remove background from an image and return the subject with transparency (RGBA PNG).
|
||
Uses rembg (AI-powered) when available.
|
||
"""
|
||
try:
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None, _extract_subject_bytes, _decode(req.image)
|
||
)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@router.post("/image/replace-subject")
|
||
async def replace_subject(req: ReplaceSubjectRequest):
|
||
"""
|
||
Extract the primary subject from `subject_image` (via rembg background removal),
|
||
scale it to fit the `mask` selection on `background_image`, apply optional LAB
|
||
color transfer for lighting consistency, and composite the result.
|
||
|
||
Returns the composited image as base64 PNG.
|
||
"""
|
||
try:
|
||
result = await asyncio.get_event_loop().run_in_executor(
|
||
None,
|
||
_do_replace_subject,
|
||
_decode(req.background_image),
|
||
_decode(req.subject_image),
|
||
_decode(req.mask) if req.mask else None,
|
||
req.match_colors,
|
||
)
|
||
return {"result": _encode(result)}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||
|
||
class ExtractColorsRequest(BaseModel):
|
||
image: str # base64
|
||
count: int = 6
|
||
|
||
|
||
def _extract_colors(image_bytes: bytes, count: int) -> list[str]:
|
||
"""
|
||
Resize image to 150×150, k-means cluster pixels into `count` groups
|
||
using pure numpy (no sklearn dependency), return hex strings by frequency.
|
||
"""
|
||
import numpy as np
|
||
from PIL import Image
|
||
from io import BytesIO
|
||
|
||
count = max(1, min(count, 32))
|
||
|
||
pil = Image.open(BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||
pixels = np.array(pil, dtype=np.float32).reshape(-1, 3) # (22500, 3)
|
||
n = len(pixels)
|
||
|
||
# Initialise centers with k-means++ seeding
|
||
rng = np.random.default_rng(42)
|
||
centers = [pixels[rng.integers(n)]]
|
||
for _ in range(count - 1):
|
||
dists = np.min([np.sum((pixels - c) ** 2, axis=1) for c in centers], axis=0)
|
||
probs = dists / dists.sum()
|
||
centers.append(pixels[rng.choice(n, p=probs)])
|
||
centers = np.array(centers)
|
||
|
||
labels = np.zeros(n, dtype=np.int32)
|
||
for _ in range(20): # max 20 iterations
|
||
# Assign each pixel to nearest center
|
||
dists = np.sum((pixels[:, None] - centers[None]) ** 2, axis=2) # (n, k)
|
||
new_labels = np.argmin(dists, axis=1)
|
||
if np.all(new_labels == labels):
|
||
break
|
||
labels = new_labels
|
||
# Recompute centers
|
||
for k in range(count):
|
||
mask = labels == k
|
||
if mask.any():
|
||
centers[k] = pixels[mask].mean(axis=0)
|
||
|
||
counts = np.bincount(labels, minlength=count)
|
||
order = np.argsort(-counts)
|
||
|
||
return [
|
||
"#{:02x}{:02x}{:02x}".format(*centers[i].astype(int).clip(0, 255))
|
||
for i in order
|
||
]
|
||
|
||
|
||
@router.post("/extract-colors")
|
||
async def extract_colors(req: ExtractColorsRequest):
|
||
"""
|
||
Extract dominant colors from an image using k-means clustering.
|
||
Returns hex color strings sorted by frequency (most dominant first).
|
||
"""
|
||
try:
|
||
image_bytes = _decode(req.image)
|
||
colors = await asyncio.get_event_loop().run_in_executor(
|
||
None, _extract_colors, image_bytes, req.count
|
||
)
|
||
return {"colors": colors}
|
||
except Exception as e:
|
||
import traceback; traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|