Add SAM (Segment Anything) via Replicate API and update Docker

- Implement SAM object selection via Replicate API
  - Click on any object to select it with AI precision
  - Falls back to flood-fill if Replicate API unavailable
- Update Dockerfile for rembg dependencies
  - Add required system libraries (libsm6, libxext6, etc)
  - Pre-download rembg model during build
  - Create data directories for models and patches
This commit is contained in:
Claude
2026-01-25 16:24:08 +00:00
parent 909bb41f8a
commit 4c2574ace4
2 changed files with 90 additions and 7 deletions
+11 -3
View File
@@ -2,10 +2,15 @@ FROM python:3.11-slim
WORKDIR /app
# Install system dependencies
# Install system dependencies for OpenCV, rembg, and image processing
RUN apt-get update && apt-get install -y \
libgl1 \
libglib2.0-0 \
libsm6 \
libxext6 \
libxrender-dev \
libgomp1 \
wget \
&& rm -rf /var/lib/apt/lists/*
# Copy requirements
@@ -14,11 +19,14 @@ COPY requirements.txt .
# Install Python dependencies
RUN pip install --no-cache-dir -r requirements.txt
# Pre-download rembg model (u2net) to avoid first-run delay
RUN python -c "from rembg import remove; print('rembg model downloaded')" || true
# Copy application
COPY . .
# Create data directory
RUN mkdir -p /app/data/projects
# Create data directories
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
# Expose port
EXPOSE 8000
+79 -4
View File
@@ -173,10 +173,85 @@ async def smart_select(
async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
"""Use Segment Anything Model for selection"""
# This would require SAM to be installed and model downloaded
# For now, raise to fall back to flood fill
raise NotImplementedError("SAM integration pending")
"""Use Segment Anything Model for selection via Replicate API"""
import httpx
import base64
import asyncio
from app.config import settings
if not settings.replicate_api_key:
raise ValueError("REPLICATE_API_KEY not configured")
# 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')
# Call SAM via Replicate
async with httpx.AsyncClient(timeout=120.0) as client:
# Use SAM model on Replicate
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", # 1 = foreground point
}
}
headers = {
'Authorization': f'Bearer {settings.replicate_api_key}',
'Content-Type': 'application/json'
}
# Start prediction
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()
# Poll for completion
prediction_url = prediction['urls']['get']
max_attempts = 60
attempt = 0
while attempt < max_attempts:
await asyncio.sleep(2)
status_response = await client.get(prediction_url, headers=headers)
status_data = status_response.json()
if status_data['status'] == 'succeeded':
# Download mask image
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')
# Resize if needed
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 # Normalize to 0-1
elif status_data['status'] == 'failed':
raise Exception(f"SAM prediction failed: {status_data.get('error')}")
attempt += 1
raise Exception("SAM prediction timed out")
def _flood_fill_select(img_array: np.ndarray, x: int, y: int, tolerance: int = 32) -> np.ndarray: