Claude
ed491ab9ba
- Fix brush_select isProcessing flag not resetting after first use (reset in on_leave() when switching tools) - Add onnxruntime dependency for U2net background removal - Use full U2net model (176MB) instead of lightweight for better quality https://claude.ai/code/session_01CLedz6CanT9t46KBvng3vz
950 lines
29 KiB
Python
950 lines
29 KiB
Python
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
|
|
from fastapi.responses import Response
|
|
from sqlalchemy.orm import Session
|
|
from typing import Optional
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import numpy as np
|
|
import json
|
|
import base64
|
|
import cv2
|
|
from pydantic import BaseModel
|
|
|
|
from app.database import get_db
|
|
from app.models.project import Project
|
|
from app.schemas import StatusResponse
|
|
|
|
router = APIRouter(prefix="/tools", tags=["tools"])
|
|
|
|
|
|
# Pydantic models for JSON API
|
|
class SmartSelectRequest(BaseModel):
|
|
image: str # Base64 encoded image
|
|
point_x: int
|
|
point_y: int
|
|
|
|
|
|
class InpaintRequest(BaseModel):
|
|
image: str # Base64 encoded image
|
|
mask: str # Base64 encoded mask
|
|
prompt: str
|
|
negative_prompt: Optional[str] = ""
|
|
strength: Optional[float] = 0.8
|
|
guidance_scale: Optional[float] = 7.5
|
|
|
|
|
|
class RemoveBackgroundRequest(BaseModel):
|
|
image: str # Base64 encoded image
|
|
|
|
|
|
@router.post("/smart-select-base64")
|
|
async def smart_select_base64(request: SmartSelectRequest):
|
|
"""
|
|
Smart select using base64 encoded image (no project required).
|
|
Used by miniPaint frontend.
|
|
"""
|
|
try:
|
|
# Decode base64 image
|
|
image_bytes = base64.b64decode(request.image)
|
|
img = Image.open(BytesIO(image_bytes)).convert('RGB')
|
|
img_array = np.array(img)
|
|
|
|
# Run SAM selection
|
|
try:
|
|
mask = await _sam_select(img_array, request.point_x, request.point_y)
|
|
except Exception as e:
|
|
print(f"SAM not available, using flood fill: {e}")
|
|
mask = _flood_fill_select(img_array, request.point_x, request.point_y)
|
|
|
|
# Convert mask to base64 PNG
|
|
mask_img = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
|
|
buffer = BytesIO()
|
|
mask_img.save(buffer, format='PNG')
|
|
mask_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
# Get polygon and bbox
|
|
polygon, bbox = _mask_to_polygon(mask)
|
|
|
|
return {
|
|
"mask": mask_b64,
|
|
"polygon": polygon,
|
|
"bbox": bbox
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post("/inpaint")
|
|
async def inpaint_base64(request: InpaintRequest):
|
|
"""
|
|
AI inpainting using base64 encoded image and mask.
|
|
Used by miniPaint frontend.
|
|
"""
|
|
try:
|
|
# Decode base64 image and mask
|
|
image_bytes = base64.b64decode(request.image)
|
|
mask_bytes = base64.b64decode(request.mask)
|
|
|
|
img = Image.open(BytesIO(image_bytes)).convert('RGB')
|
|
mask_img = Image.open(BytesIO(mask_bytes)).convert('L')
|
|
|
|
# Resize mask to match image if needed
|
|
if mask_img.size != img.size:
|
|
mask_img = mask_img.resize(img.size, Image.Resampling.LANCZOS)
|
|
|
|
# Get the AI provider
|
|
from app.services.ai_provider import get_ai_provider
|
|
|
|
provider = get_ai_provider()
|
|
|
|
# Convert images to bytes for provider
|
|
img_buffer = BytesIO()
|
|
img.save(img_buffer, format='PNG')
|
|
img_bytes = img_buffer.getvalue()
|
|
|
|
mask_buffer = BytesIO()
|
|
mask_img.save(mask_buffer, format='PNG')
|
|
mask_bytes_png = mask_buffer.getvalue()
|
|
|
|
# Run inpainting using edit_image method
|
|
result_bytes = await provider.edit_image(
|
|
patch_image_bytes=img_bytes,
|
|
mask_image_bytes=mask_bytes_png,
|
|
prompt=request.prompt,
|
|
mode="A" # Patch-only mode
|
|
)
|
|
|
|
# Convert result to base64
|
|
result_b64 = base64.b64encode(result_bytes).decode('utf-8')
|
|
|
|
return {
|
|
"result": result_b64
|
|
}
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
class RemoveBackgroundRequestV2(BaseModel):
|
|
image: str # Base64 encoded image
|
|
model: Optional[str] = "auto" # "auto", "u2net", "rembg", "birefnet"
|
|
|
|
|
|
@router.post("/remove-background-base64")
|
|
async def remove_background_base64(request: RemoveBackgroundRequest):
|
|
"""
|
|
Remove background from a base64 encoded image.
|
|
Tries multiple methods: U2Net (direct), rembg with BiRefNet, rembg default.
|
|
Returns base64 encoded PNG with transparent background.
|
|
Used by miniPaint frontend.
|
|
"""
|
|
try:
|
|
# Decode base64 image
|
|
image_bytes = base64.b64decode(request.image)
|
|
img = Image.open(BytesIO(image_bytes)).convert('RGB')
|
|
|
|
result_bytes = None
|
|
method_used = None
|
|
|
|
# Try U2Net first (direct implementation, no rembg dependency issues)
|
|
try:
|
|
result_bytes = await _remove_background_u2net(img)
|
|
method_used = "u2net"
|
|
except Exception as e:
|
|
print(f"U2Net failed: {e}")
|
|
|
|
# Fall back to rembg if U2Net failed
|
|
if result_bytes is None:
|
|
try:
|
|
from rembg import remove, new_session
|
|
try:
|
|
session = new_session("birefnet-general")
|
|
result_bytes = remove(image_bytes, session=session)
|
|
method_used = "birefnet"
|
|
except Exception:
|
|
result_bytes = remove(image_bytes)
|
|
method_used = "rembg-default"
|
|
except ImportError:
|
|
pass
|
|
except Exception as e:
|
|
print(f"rembg failed: {e}")
|
|
|
|
if result_bytes is None:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="No background removal method available. Install u2net or rembg."
|
|
)
|
|
|
|
# Convert result to base64
|
|
result_b64 = base64.b64encode(result_bytes).decode('utf-8')
|
|
|
|
# Get dimensions
|
|
result_img = Image.open(BytesIO(result_bytes))
|
|
|
|
return {
|
|
"result": result_b64,
|
|
"width": result_img.width,
|
|
"height": result_img.height,
|
|
"method": method_used
|
|
}
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# Global U2Net model cache
|
|
_u2net_model = None
|
|
|
|
|
|
async def _download_u2net_model(models_dir):
|
|
"""Auto-download full U2Net model (~176MB) for best quality background removal"""
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
models_dir = Path(models_dir)
|
|
models_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Download full U2Net model (176MB) for best quality
|
|
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
|
|
dest_path = models_dir / "u2net.onnx"
|
|
|
|
print(f"Downloading full U2Net model from {url} (~176MB)...")
|
|
print("This may take a few minutes...")
|
|
|
|
def download_progress(count, block_size, total_size):
|
|
if total_size > 0:
|
|
percent = min(100, count * block_size * 100 // total_size)
|
|
downloaded_mb = (count * block_size) / (1024 * 1024)
|
|
total_mb = total_size / (1024 * 1024)
|
|
if count % 500 == 0:
|
|
print(f" Download progress: {percent}% ({downloaded_mb:.1f}/{total_mb:.1f} MB)")
|
|
|
|
urllib.request.urlretrieve(url, str(dest_path), download_progress)
|
|
print(f"Full U2Net model downloaded to {dest_path}")
|
|
|
|
return dest_path
|
|
|
|
|
|
async def _remove_background_u2net(img: Image.Image) -> bytes:
|
|
"""
|
|
Remove background using U2Net model directly.
|
|
This avoids rembg dependency issues while providing good quality.
|
|
"""
|
|
global _u2net_model
|
|
|
|
import torch
|
|
from pathlib import Path
|
|
|
|
# Check for U2Net model
|
|
models_dir = Path('/app/data/models')
|
|
u2net_path = models_dir / 'u2net.pth'
|
|
|
|
# Also check alternative names
|
|
if not u2net_path.exists():
|
|
for alt_name in ['u2net.onnx', 'u2netp.pth', 'u2net_human_seg.pth']:
|
|
alt_path = models_dir / alt_name
|
|
if alt_path.exists():
|
|
u2net_path = alt_path
|
|
break
|
|
|
|
if not u2net_path.exists():
|
|
# Try to auto-download the model
|
|
print("U2Net model not found, attempting to download...")
|
|
try:
|
|
await _download_u2net_model(models_dir)
|
|
# Check again
|
|
for alt_name in ['u2net.onnx', 'u2netp.onnx', 'u2net.pth']:
|
|
alt_path = models_dir / alt_name
|
|
if alt_path.exists():
|
|
u2net_path = alt_path
|
|
break
|
|
except Exception as download_error:
|
|
print(f"Auto-download failed: {download_error}")
|
|
|
|
if not u2net_path.exists():
|
|
raise FileNotFoundError(
|
|
"U2Net model not found. To fix this, run:\n"
|
|
" docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py\n"
|
|
"Or manually download from: https://github.com/danielgatis/rembg/releases"
|
|
)
|
|
|
|
# Load model if not cached
|
|
if _u2net_model is None:
|
|
print(f"Loading U2Net model from {u2net_path}")
|
|
|
|
if str(u2net_path).endswith('.onnx'):
|
|
# Use ONNX runtime
|
|
import onnxruntime as ort
|
|
_u2net_model = ort.InferenceSession(str(u2net_path))
|
|
else:
|
|
# Use PyTorch
|
|
from app.services.u2net_model import U2NET
|
|
_u2net_model = U2NET(3, 1)
|
|
_u2net_model.load_state_dict(torch.load(str(u2net_path), map_location='cpu'))
|
|
_u2net_model.eval()
|
|
|
|
print("U2Net model loaded")
|
|
|
|
# Preprocess image
|
|
img_np = np.array(img)
|
|
original_size = img.size
|
|
|
|
# Resize to model input size
|
|
input_size = 320
|
|
img_resized = img.resize((input_size, input_size), Image.Resampling.BILINEAR)
|
|
img_np = np.array(img_resized).astype(np.float32)
|
|
|
|
# Normalize
|
|
img_np = img_np / 255.0
|
|
img_np = (img_np - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
|
|
img_np = img_np.transpose(2, 0, 1) # HWC to CHW
|
|
img_np = np.expand_dims(img_np, 0) # Add batch dimension
|
|
|
|
# Run inference
|
|
if hasattr(_u2net_model, 'run'):
|
|
# ONNX runtime
|
|
input_name = _u2net_model.get_inputs()[0].name
|
|
outputs = _u2net_model.run(None, {input_name: img_np})
|
|
mask = outputs[0][0, 0]
|
|
else:
|
|
# PyTorch
|
|
with torch.no_grad():
|
|
input_tensor = torch.from_numpy(img_np).float()
|
|
d1, d2, d3, d4, d5, d6, d7 = _u2net_model(input_tensor)
|
|
mask = d1[0, 0].numpy()
|
|
|
|
# Post-process mask
|
|
mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
|
|
mask = (mask * 255).astype(np.uint8)
|
|
|
|
# Resize mask back to original size
|
|
mask_img = Image.fromarray(mask).resize(original_size, Image.Resampling.BILINEAR)
|
|
|
|
# Apply mask to original image
|
|
result = img.convert('RGBA')
|
|
result.putalpha(mask_img)
|
|
|
|
# Save to bytes
|
|
buffer = BytesIO()
|
|
result.save(buffer, format='PNG')
|
|
return buffer.getvalue()
|
|
|
|
|
|
@router.post("/remove-background")
|
|
async def remove_background(
|
|
project_id: Optional[int] = Form(None),
|
|
file: Optional[UploadFile] = File(None),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Remove background from an image using rembg with BiRefNet model.
|
|
|
|
Either provide project_id to use current project image,
|
|
or upload a file directly.
|
|
|
|
Returns PNG with transparent background.
|
|
"""
|
|
try:
|
|
from rembg import remove, new_session
|
|
except ImportError:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="rembg not installed. Run: pip install rembg"
|
|
)
|
|
|
|
# Get image bytes
|
|
if file:
|
|
image_bytes = await file.read()
|
|
elif project_id:
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
edit_service = EditService()
|
|
image_path = edit_service.get_current_image_path(project_id)
|
|
|
|
with open(image_path, 'rb') as f:
|
|
image_bytes = f.read()
|
|
else:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Provide either project_id or file"
|
|
)
|
|
|
|
# Remove background using BiRefNet (state-of-the-art)
|
|
try:
|
|
session = new_session("birefnet-general")
|
|
result_bytes = remove(image_bytes, session=session)
|
|
except Exception:
|
|
result_bytes = remove(image_bytes)
|
|
|
|
return Response(
|
|
content=result_bytes,
|
|
media_type="image/png",
|
|
headers={"Content-Disposition": "inline; filename=no-background.png"}
|
|
)
|
|
|
|
|
|
@router.post("/remove-background-to-layer")
|
|
async def remove_background_to_layer(
|
|
project_id: int = Form(...),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Remove background using BiRefNet and save as a new layer in the project.
|
|
Returns layer info that can be added to frontend layer system.
|
|
"""
|
|
try:
|
|
from rembg import remove, new_session
|
|
except ImportError:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail="rembg not installed. Run: pip install rembg"
|
|
)
|
|
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
from pathlib import Path
|
|
|
|
edit_service = EditService()
|
|
image_path = edit_service.get_current_image_path(project_id)
|
|
|
|
with open(image_path, 'rb') as f:
|
|
image_bytes = f.read()
|
|
|
|
# Remove background using BiRefNet (state-of-the-art)
|
|
try:
|
|
session = new_session("birefnet-general")
|
|
result_bytes = remove(image_bytes, session=session)
|
|
except Exception:
|
|
result_bytes = remove(image_bytes)
|
|
|
|
# Save as layer file
|
|
project_dir = edit_service.get_project_dir(project_id)
|
|
layers_dir = project_dir / 'layers'
|
|
layers_dir.mkdir(exist_ok=True)
|
|
|
|
# Find next layer number
|
|
existing_layers = list(layers_dir.glob('layer_*.png'))
|
|
layer_num = len(existing_layers) + 1
|
|
layer_path = layers_dir / f'layer_{layer_num}.png'
|
|
|
|
with open(layer_path, 'wb') as f:
|
|
f.write(result_bytes)
|
|
|
|
# Get dimensions
|
|
img = Image.open(BytesIO(result_bytes))
|
|
|
|
return {
|
|
"status": "success",
|
|
"layer": {
|
|
"id": layer_num,
|
|
"name": f"No Background {layer_num}",
|
|
"path": str(layer_path),
|
|
"width": img.width,
|
|
"height": img.height,
|
|
"type": "background_removed"
|
|
}
|
|
}
|
|
|
|
|
|
@router.post("/smart-select")
|
|
async def smart_select(
|
|
project_id: int = Form(...),
|
|
point_x: int = Form(...),
|
|
point_y: int = Form(...),
|
|
return_format: str = Form("json"), # "json" (default) or "image"
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Use SAM (Segment Anything) to select object at given point.
|
|
Returns mask and polygon data for the selected object.
|
|
|
|
Note: Requires SAM model to be downloaded.
|
|
Falls back to simple flood-fill selection if SAM unavailable.
|
|
"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
edit_service = EditService()
|
|
image_path = edit_service.get_current_image_path(project_id)
|
|
|
|
img = Image.open(image_path).convert('RGB')
|
|
img_array = np.array(img)
|
|
|
|
# Try SAM first, fall back to flood fill
|
|
try:
|
|
mask = await _sam_select(img_array, point_x, point_y)
|
|
except Exception as e:
|
|
print(f"SAM not available, using flood fill: {e}")
|
|
mask = _flood_fill_select(img_array, point_x, point_y)
|
|
|
|
# Convert mask to PNG
|
|
mask_img = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
|
|
|
|
if return_format == "image":
|
|
buffer = BytesIO()
|
|
mask_img.save(buffer, format='PNG')
|
|
return Response(
|
|
content=buffer.getvalue(),
|
|
media_type="image/png"
|
|
)
|
|
|
|
# Return JSON with polygon and bbox
|
|
polygon, bbox = _mask_to_polygon(mask)
|
|
|
|
# Also return mask as base64 for potential use
|
|
buffer = BytesIO()
|
|
mask_img.save(buffer, format='PNG')
|
|
mask_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
return {
|
|
"polygon": polygon,
|
|
"bbox": bbox,
|
|
"mask_base64": mask_b64,
|
|
}
|
|
|
|
|
|
def _mask_to_polygon(mask: np.ndarray) -> tuple:
|
|
"""
|
|
Convert a binary mask to a simplified polygon and bounding box.
|
|
|
|
Returns:
|
|
(polygon, bbox) where:
|
|
- polygon: list of [x, y] points (simplified contour)
|
|
- bbox: dict with x, y, width, height
|
|
"""
|
|
# Ensure mask is binary uint8
|
|
mask_uint8 = (mask * 255).astype(np.uint8) if mask.max() <= 1 else mask.astype(np.uint8)
|
|
|
|
# Find contours
|
|
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
if not contours:
|
|
return [], {"x": 0, "y": 0, "width": 0, "height": 0}
|
|
|
|
# Get largest contour
|
|
largest = max(contours, key=cv2.contourArea)
|
|
|
|
# Get bounding box
|
|
x, y, w, h = cv2.boundingRect(largest)
|
|
bbox = {"x": int(x), "y": int(y), "width": int(w), "height": int(h)}
|
|
|
|
# Simplify contour to reduce points (epsilon = 1% of arc length)
|
|
epsilon = 0.01 * cv2.arcLength(largest, True)
|
|
simplified = cv2.approxPolyDP(largest, epsilon, True)
|
|
|
|
# Convert to list of [x, y] points
|
|
polygon = [[int(pt[0][0]), int(pt[0][1])] for pt in simplified]
|
|
|
|
return polygon, bbox
|
|
|
|
|
|
# Global SAM model cache (loaded once, reused)
|
|
_sam_model = None
|
|
_sam_predictor = None
|
|
|
|
|
|
def _get_sam_model():
|
|
"""Load SAM model from local file (cached after first load)"""
|
|
global _sam_model, _sam_predictor
|
|
|
|
if _sam_predictor is not None:
|
|
return _sam_predictor
|
|
|
|
from pathlib import Path
|
|
|
|
# Check for SAM model in models directory
|
|
models_dir = Path('/app/data/models')
|
|
model_path = models_dir / 'sam_model.pth'
|
|
|
|
# Also check for specific model files
|
|
if not model_path.exists():
|
|
for filename in ['sam_vit_b_01ec64.pth', 'sam_vit_l_0b3195.pth', 'sam_vit_h_4b8939.pth']:
|
|
alt_path = models_dir / filename
|
|
if alt_path.exists():
|
|
model_path = alt_path
|
|
break
|
|
|
|
if not model_path.exists():
|
|
raise FileNotFoundError(
|
|
f"SAM model not found. Download it with:\n"
|
|
f" docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py"
|
|
)
|
|
|
|
# Determine model type from filename
|
|
model_type = 'vit_b' # default
|
|
if 'vit_l' in model_path.name:
|
|
model_type = 'vit_l'
|
|
elif 'vit_h' in model_path.name:
|
|
model_type = 'vit_h'
|
|
|
|
print(f"Loading SAM model: {model_path} (type: {model_type})")
|
|
|
|
import torch
|
|
from segment_anything import sam_model_registry, SamPredictor
|
|
|
|
# Use CPU by default (works everywhere), GPU if available
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
_sam_model = sam_model_registry[model_type](checkpoint=str(model_path))
|
|
_sam_model.to(device)
|
|
_sam_predictor = SamPredictor(_sam_model)
|
|
|
|
print(f"SAM model loaded on {device}")
|
|
return _sam_predictor
|
|
|
|
|
|
def _sam_select_local(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
|
"""Use local SAM model for selection (no API calls, runs offline)"""
|
|
predictor = _get_sam_model()
|
|
|
|
# Set image
|
|
predictor.set_image(img_array)
|
|
|
|
# Point coordinates (x, y) and label (1 = foreground)
|
|
input_point = np.array([[x, y]])
|
|
input_label = np.array([1])
|
|
|
|
# Get mask prediction
|
|
masks, scores, _ = predictor.predict(
|
|
point_coords=input_point,
|
|
point_labels=input_label,
|
|
multimask_output=True, # Get multiple mask options
|
|
)
|
|
|
|
# Use the mask with highest score
|
|
best_mask_idx = np.argmax(scores)
|
|
mask = masks[best_mask_idx]
|
|
|
|
return mask.astype(np.uint8)
|
|
|
|
|
|
async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
|
"""
|
|
Smart object selection using SAM (Segment Anything Model).
|
|
|
|
Priority:
|
|
1. Local SAM model (free, fast, offline)
|
|
2. Replicate API (if local not available and API key set)
|
|
3. Raises exception if neither available
|
|
"""
|
|
# Try local SAM first (free, no API calls)
|
|
try:
|
|
return _sam_select_local(img_array, x, y)
|
|
except FileNotFoundError as e:
|
|
print(f"Local SAM not available: {e}")
|
|
except ImportError as e:
|
|
print(f"SAM dependencies not installed: {e}")
|
|
except Exception as e:
|
|
print(f"Local SAM failed: {e}")
|
|
|
|
# Fall back to Replicate API
|
|
from app.config import settings
|
|
|
|
if not settings.replicate_api_key:
|
|
raise ValueError(
|
|
"SAM model not available. Either:\n"
|
|
" 1. Download local model: docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py\n"
|
|
" 2. Or set REPLICATE_API_KEY in .env for cloud SAM"
|
|
)
|
|
|
|
return await _sam_select_replicate(img_array, x, y)
|
|
|
|
|
|
async def _sam_select_replicate(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
|
"""Fallback: Use SAM via Replicate API (requires API key, costs ~$0.002/call)"""
|
|
import httpx
|
|
import base64
|
|
import asyncio
|
|
from app.config import settings
|
|
|
|
# Convert image to base64
|
|
img = Image.fromarray(img_array)
|
|
buffer = BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
img_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
|
prediction_data = {
|
|
"version": "meta/sam-2-image:fe97b453d6525baeeb530595c74a3c4f567c1f655ee2a0fee11f76bd1d31e495",
|
|
"input": {
|
|
"image": f"data:image/png;base64,{img_b64}",
|
|
"point_coords": f"{x},{y}",
|
|
"point_labels": "1",
|
|
}
|
|
}
|
|
|
|
headers = {
|
|
'Authorization': f'Bearer {settings.replicate_api_key}',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
response = await client.post(
|
|
"https://api.replicate.com/v1/predictions",
|
|
json=prediction_data,
|
|
headers=headers
|
|
)
|
|
|
|
if response.status_code != 201:
|
|
raise Exception(f"Replicate API error: {response.text}")
|
|
|
|
prediction = response.json()
|
|
prediction_url = prediction['urls']['get']
|
|
|
|
# Poll for completion
|
|
for _ in range(60):
|
|
await asyncio.sleep(2)
|
|
status_response = await client.get(prediction_url, headers=headers)
|
|
status_data = status_response.json()
|
|
|
|
if status_data['status'] == 'succeeded':
|
|
mask_url = status_data['output']
|
|
if isinstance(mask_url, list):
|
|
mask_url = mask_url[0]
|
|
|
|
mask_response = await client.get(mask_url)
|
|
mask_img = Image.open(BytesIO(mask_response.content)).convert('L')
|
|
|
|
if mask_img.size != (img_array.shape[1], img_array.shape[0]):
|
|
mask_img = mask_img.resize(
|
|
(img_array.shape[1], img_array.shape[0]),
|
|
Image.Resampling.LANCZOS
|
|
)
|
|
|
|
return np.array(mask_img) // 255
|
|
|
|
elif status_data['status'] == 'failed':
|
|
raise Exception(f"SAM prediction failed: {status_data.get('error')}")
|
|
|
|
raise Exception("SAM prediction timed out")
|
|
|
|
|
|
def _flood_fill_select(img_array: np.ndarray, x: int, y: int, tolerance: int = 32) -> np.ndarray:
|
|
"""Simple flood-fill based selection with color tolerance"""
|
|
import cv2
|
|
|
|
h, w = img_array.shape[:2]
|
|
|
|
# Ensure point is within bounds
|
|
x = max(0, min(x, w - 1))
|
|
y = max(0, min(y, h - 1))
|
|
|
|
# Create mask for flood fill (needs to be 2 pixels larger)
|
|
mask = np.zeros((h + 2, w + 2), np.uint8)
|
|
|
|
# Flood fill
|
|
cv2.floodFill(
|
|
img_array.copy(),
|
|
mask,
|
|
(x, y),
|
|
(255, 255, 255),
|
|
(tolerance, tolerance, tolerance),
|
|
(tolerance, tolerance, tolerance),
|
|
cv2.FLOODFILL_MASK_ONLY
|
|
)
|
|
|
|
# Extract the actual mask (remove padding)
|
|
return mask[1:-1, 1:-1]
|
|
|
|
|
|
@router.post("/color-select")
|
|
async def color_select(
|
|
project_id: int = Form(...),
|
|
color_r: int = Form(...),
|
|
color_g: int = Form(...),
|
|
color_b: int = Form(...),
|
|
tolerance: int = Form(30),
|
|
return_format: str = Form("json"), # "json" (default) or "image"
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Select all pixels similar to the given color.
|
|
Returns a mask and polygon data for selected areas.
|
|
"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
edit_service = EditService()
|
|
image_path = edit_service.get_current_image_path(project_id)
|
|
|
|
img = Image.open(image_path).convert('RGB')
|
|
img_array = np.array(img)
|
|
|
|
# Target color
|
|
target = np.array([color_r, color_g, color_b])
|
|
|
|
# Calculate color distance
|
|
diff = np.abs(img_array.astype(np.int16) - target.astype(np.int16))
|
|
distance = np.sum(diff, axis=2)
|
|
|
|
# Create mask where distance is within tolerance
|
|
mask = (distance <= tolerance * 3).astype(np.uint8)
|
|
|
|
# Convert to PNG
|
|
mask_img = Image.fromarray(mask * 255, mode='L')
|
|
|
|
if return_format == "image":
|
|
buffer = BytesIO()
|
|
mask_img.save(buffer, format='PNG')
|
|
return Response(
|
|
content=buffer.getvalue(),
|
|
media_type="image/png"
|
|
)
|
|
|
|
# Return JSON with polygon and bbox
|
|
polygon, bbox = _mask_to_polygon(mask)
|
|
|
|
buffer = BytesIO()
|
|
mask_img.save(buffer, format='PNG')
|
|
mask_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
return {
|
|
"polygon": polygon,
|
|
"bbox": bbox,
|
|
"mask_base64": mask_b64,
|
|
"color": {"r": color_r, "g": color_g, "b": color_b},
|
|
"tolerance": tolerance,
|
|
}
|
|
|
|
|
|
@router.post("/extract-object")
|
|
async def extract_object(
|
|
project_id: int = Form(...),
|
|
mask: UploadFile = File(...),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Extract object using provided mask.
|
|
Returns PNG with transparent background containing only the masked area.
|
|
"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
edit_service = EditService()
|
|
image_path = edit_service.get_current_image_path(project_id)
|
|
|
|
# Load image and mask
|
|
img = Image.open(image_path).convert('RGBA')
|
|
mask_bytes = await mask.read()
|
|
mask_img = Image.open(BytesIO(mask_bytes)).convert('L')
|
|
|
|
# Resize mask if needed
|
|
if mask_img.size != img.size:
|
|
mask_img = mask_img.resize(img.size, Image.Resampling.LANCZOS)
|
|
|
|
# Apply mask as alpha channel
|
|
img_array = np.array(img)
|
|
mask_array = np.array(mask_img)
|
|
|
|
# Set alpha channel based on mask
|
|
img_array[:, :, 3] = mask_array
|
|
|
|
result = Image.fromarray(img_array, mode='RGBA')
|
|
|
|
buffer = BytesIO()
|
|
result.save(buffer, format='PNG')
|
|
|
|
return Response(
|
|
content=buffer.getvalue(),
|
|
media_type="image/png"
|
|
)
|
|
|
|
|
|
@router.get("/layers/{project_id}")
|
|
async def list_layers(
|
|
project_id: int,
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""List all layers for a project"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
from pathlib import Path
|
|
|
|
edit_service = EditService()
|
|
project_dir = edit_service.get_project_dir(project_id)
|
|
layers_dir = project_dir / 'layers'
|
|
|
|
if not layers_dir.exists():
|
|
return {"layers": []}
|
|
|
|
layers = []
|
|
for layer_file in sorted(layers_dir.glob('layer_*.png')):
|
|
img = Image.open(layer_file)
|
|
layer_num = int(layer_file.stem.split('_')[1])
|
|
layers.append({
|
|
"id": layer_num,
|
|
"name": f"Layer {layer_num}",
|
|
"path": str(layer_file),
|
|
"width": img.width,
|
|
"height": img.height
|
|
})
|
|
|
|
return {"layers": layers}
|
|
|
|
|
|
@router.post("/flatten-layers")
|
|
async def flatten_layers(
|
|
project_id: int = Form(...),
|
|
layer_order: str = Form(...), # JSON array of layer IDs in order
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""
|
|
Flatten all layers into a single image and save as current.
|
|
layer_order is a JSON array like [1, 2, 3] from bottom to top.
|
|
"""
|
|
project = db.query(Project).filter(Project.id == project_id).first()
|
|
if not project:
|
|
raise HTTPException(status_code=404, detail="Project not found")
|
|
|
|
from app.services.edit_service import EditService
|
|
from pathlib import Path
|
|
|
|
edit_service = EditService()
|
|
project_dir = edit_service.get_project_dir(project_id)
|
|
layers_dir = project_dir / 'layers'
|
|
|
|
order = json.loads(layer_order)
|
|
|
|
# Start with original image as base
|
|
base_path = edit_service.get_current_image_path(project_id)
|
|
result = Image.open(base_path).convert('RGBA')
|
|
|
|
# Composite layers in order
|
|
for layer_id in order:
|
|
layer_path = layers_dir / f'layer_{layer_id}.png'
|
|
if layer_path.exists():
|
|
layer = Image.open(layer_path).convert('RGBA')
|
|
# Resize if needed
|
|
if layer.size != result.size:
|
|
layer = layer.resize(result.size, Image.Resampling.LANCZOS)
|
|
result = Image.alpha_composite(result, layer)
|
|
|
|
# Save as current
|
|
result.save(base_path, 'PNG')
|
|
|
|
return StatusResponse(
|
|
status="success",
|
|
message="Layers flattened successfully"
|
|
)
|