Merge pull request #35 from outis1one/claude/migrate-to-minipaint-eYKWf
Upgrade rembg to use BiRefNet model (state-of-the-art)
This commit is contained in:
+3
-4
@@ -22,7 +22,7 @@ FROM python:3.11-slim
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies for OpenCV, rembg, SAM, and image processing
|
||||
# execstack is needed to fix onnxruntime executable stack issue in containers
|
||||
# Note: onnxruntime 1.17+ fixed executable stack issues, no longer need execstack
|
||||
RUN apt-get update && apt-get install -y \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
@@ -39,9 +39,8 @@ RUN apt-get update && apt-get install -y \
|
||||
COPY backend/requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Try to pre-download rembg model, but don't fail build if onnxruntime has issues
|
||||
# (onnxruntime can have executable stack issues in some Docker environments)
|
||||
# If this fails, Remove Background feature will be disabled
|
||||
# Verify rembg loads correctly (model downloads on first use)
|
||||
# rembg now supports BiRefNet models which are state-of-the-art for background removal
|
||||
RUN python -c "from rembg import remove; print('rembg ready')" || echo "WARNING: rembg not available - Remove Background will be disabled"
|
||||
|
||||
# Copy backend application
|
||||
|
||||
@@ -131,12 +131,13 @@ async def inpaint_base64(request: InpaintRequest):
|
||||
@router.post("/remove-background-base64")
|
||||
async def remove_background_base64(request: RemoveBackgroundRequest):
|
||||
"""
|
||||
Remove background from a base64 encoded image using rembg.
|
||||
Remove background from a base64 encoded image using rembg with BiRefNet.
|
||||
BiRefNet is state-of-the-art for background removal (better than u2net).
|
||||
Returns base64 encoded PNG with transparent background.
|
||||
Used by miniPaint frontend.
|
||||
"""
|
||||
try:
|
||||
from rembg import remove
|
||||
from rembg import remove, new_session
|
||||
except ImportError:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
@@ -147,8 +148,14 @@ async def remove_background_base64(request: RemoveBackgroundRequest):
|
||||
# Decode base64 image
|
||||
image_bytes = base64.b64decode(request.image)
|
||||
|
||||
# Remove background
|
||||
result_bytes = remove(image_bytes)
|
||||
# Use BiRefNet model for best quality (state-of-the-art)
|
||||
# Falls back to default model if BiRefNet not available
|
||||
try:
|
||||
session = new_session("birefnet-general")
|
||||
result_bytes = remove(image_bytes, session=session)
|
||||
except Exception:
|
||||
# Fallback to default model
|
||||
result_bytes = remove(image_bytes)
|
||||
|
||||
# Convert result to base64
|
||||
result_b64 = base64.b64encode(result_bytes).decode('utf-8')
|
||||
@@ -175,7 +182,7 @@ async def remove_background(
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Remove background from an image using rembg.
|
||||
Remove background from an image using rembg with BiRefNet model.
|
||||
|
||||
Either provide project_id to use current project image,
|
||||
or upload a file directly.
|
||||
@@ -183,7 +190,7 @@ async def remove_background(
|
||||
Returns PNG with transparent background.
|
||||
"""
|
||||
try:
|
||||
from rembg import remove
|
||||
from rembg import remove, new_session
|
||||
except ImportError:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
@@ -210,8 +217,12 @@ async def remove_background(
|
||||
detail="Provide either project_id or file"
|
||||
)
|
||||
|
||||
# Remove background
|
||||
result_bytes = remove(image_bytes)
|
||||
# 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,
|
||||
@@ -226,11 +237,11 @@ async def remove_background_to_layer(
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Remove background and save as a new layer in the project.
|
||||
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
|
||||
from rembg import remove, new_session
|
||||
except ImportError:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
@@ -250,8 +261,12 @@ async def remove_background_to_layer(
|
||||
with open(image_path, 'rb') as f:
|
||||
image_bytes = f.read()
|
||||
|
||||
# Remove background
|
||||
result_bytes = remove(image_bytes)
|
||||
# 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)
|
||||
|
||||
@@ -14,9 +14,10 @@ pydantic-settings==2.1.0
|
||||
email-validator==2.1.0
|
||||
opencv-python-headless==4.9.0.80
|
||||
scikit-image==0.22.0
|
||||
# rembg for background removal - use onnxruntime 1.14.1 to avoid executable stack issue in Docker
|
||||
onnxruntime==1.14.1
|
||||
rembg==2.0.50
|
||||
# rembg for background removal with BiRefNet models (state-of-the-art)
|
||||
# Using latest onnxruntime (1.17+ fixed executable stack issues)
|
||||
onnxruntime>=1.17.0
|
||||
rembg>=2.0.70
|
||||
# SAM (Segment Anything) for smart object selection - runs locally, no API needed
|
||||
torch==2.1.2
|
||||
torchvision==0.16.2
|
||||
|
||||
Reference in New Issue
Block a user