Refactor: remove AGPL imgly dependency and migrate background removal to python backend
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
FROM python:3.10-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies
|
||||
# We need ffmpeg for video processing and libsm6 libxext6 for opencv (if needed by backgroundremover)
|
||||
RUN apt-get update && apt-get install -y \
|
||||
ffmpeg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
wget \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Pre-create the directory for u2net models so they persist if mounted,
|
||||
# or at least get downloaded to a known location
|
||||
ENV U2NET_HOME=/root/.u2net
|
||||
RUN mkdir -p /root/.u2net
|
||||
|
||||
COPY requirements.txt .
|
||||
|
||||
RUN pip install --no-cache-dir --upgrade pip && \
|
||||
pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# We explicitly pre-download the model by running backgroundremover on a dummy image if we wanted,
|
||||
# but the script downloads it on first run. To avoid huge first-request times,
|
||||
# you could add a script here to download the u2net model directly.
|
||||
# wget https://github.com/nadermx/backgroundremover/raw/main/models/u2net.pth -O /root/.u2net/u2net.pth (if available)
|
||||
|
||||
COPY . .
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
# Using shm-size is handled in docker-compose, but we set workers to 1
|
||||
# because video processing is extremely heavy and multiprocessing can crash without enough shm
|
||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]
|
||||
@@ -0,0 +1,27 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
background-remover:
|
||||
build: .
|
||||
container_name: background-remover-service
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
# Mount the models folder so it doesn't download the model every time the container is recreated
|
||||
- u2net_models:/root/.u2net
|
||||
# Mount for local development if needed
|
||||
- .:/app
|
||||
# Video processing uses multiprocessing and requires shared memory
|
||||
shm_size: '2g'
|
||||
restart: unless-stopped
|
||||
# If a GPU becomes available, you would uncomment the following lines (and ensure the Dockerfile uses a CUDA base image)
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: 1
|
||||
# capabilities: [gpu]
|
||||
|
||||
volumes:
|
||||
u2net_models:
|
||||
@@ -0,0 +1,192 @@
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import shutil
|
||||
import asyncio
|
||||
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
app = FastAPI(title="Background Remover API")
|
||||
|
||||
# Configure CORS so the React app can call this API
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"], # Adjust in production
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
def cleanup_files(*file_paths):
|
||||
"""Deletes temporary files after the response is sent."""
|
||||
for file_path in file_paths:
|
||||
try:
|
||||
if file_path and os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
print(f"Cleaned up {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Error cleaning up {file_path}: {e}")
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.post("/api/v1/remove-background")
|
||||
async def remove_background(
|
||||
background_tasks: BackgroundTasks,
|
||||
file: UploadFile = File(...)
|
||||
):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="No file uploaded")
|
||||
|
||||
# We will use secure temp files
|
||||
temp_input_fd, temp_input_path = tempfile.mkstemp(suffix=".mp4")
|
||||
temp_output_mov_fd, temp_output_mov_path = tempfile.mkstemp(suffix=".mov")
|
||||
temp_final_webm_fd, temp_final_webm_path = tempfile.mkstemp(suffix=".webm")
|
||||
|
||||
os.close(temp_input_fd)
|
||||
os.close(temp_output_mov_fd)
|
||||
os.close(temp_final_webm_fd)
|
||||
|
||||
try:
|
||||
# 1. Save uploaded video to temp file
|
||||
with open(temp_input_path, "wb") as buffer:
|
||||
shutil.copyfileobj(file.file, buffer)
|
||||
|
||||
print(f"Saved uploaded file to {temp_input_path}")
|
||||
|
||||
# 2. Run backgroundremover
|
||||
# backgroundremover -i input.mp4 -mk -o output.mov
|
||||
cmd_bg_remover = [
|
||||
"backgroundremover",
|
||||
"-i", temp_input_path,
|
||||
"-mk", # This flag tells it to create an alpha matte video (.mov)
|
||||
"-o", temp_output_mov_path
|
||||
]
|
||||
|
||||
print("Starting background removal process...")
|
||||
process_bg = await asyncio.create_subprocess_exec(
|
||||
*cmd_bg_remover,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
stdout, stderr = await process_bg.communicate()
|
||||
|
||||
if process_bg.returncode != 0:
|
||||
print(f"backgroundremover error: {stderr.decode()}")
|
||||
raise HTTPException(status_code=500, detail="Error processing video background removal.")
|
||||
|
||||
print("Background removal finished. Starting WebM conversion...")
|
||||
|
||||
# 3. Convert .mov to .webm with VP9 and yuva420p (alpha channel)
|
||||
# ffmpeg -i output.mov -c:v libvpx-vp9 -pix_fmt yuva420p -auto-alt-ref 0 final.webm
|
||||
cmd_ffmpeg = [
|
||||
"ffmpeg",
|
||||
"-y", # Overwrite output
|
||||
"-i", temp_output_mov_path,
|
||||
"-c:v", "libvpx-vp9",
|
||||
"-pix_fmt", "yuva420p",
|
||||
"-auto-alt-ref", "0",
|
||||
temp_final_webm_path
|
||||
]
|
||||
|
||||
process_ffmpeg = await asyncio.create_subprocess_exec(
|
||||
*cmd_ffmpeg,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
stdout_ff, stderr_ff = await process_ffmpeg.communicate()
|
||||
|
||||
if process_ffmpeg.returncode != 0:
|
||||
print(f"ffmpeg error: {stderr_ff.decode()}")
|
||||
raise HTTPException(status_code=500, detail="Error converting video to WebM.")
|
||||
|
||||
print("WebM conversion finished successfully.")
|
||||
|
||||
# 4. Schedule cleanup of ALL temporary files after response is sent
|
||||
background_tasks.add_task(
|
||||
cleanup_files,
|
||||
temp_input_path,
|
||||
temp_output_mov_path,
|
||||
temp_final_webm_path
|
||||
)
|
||||
|
||||
# 5. Return the file
|
||||
return FileResponse(
|
||||
temp_final_webm_path,
|
||||
media_type="video/webm",
|
||||
filename=f"transparent_{file.filename.split('.')[0]}.webm"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# If an error occurs, clean up immediately
|
||||
cleanup_files(temp_input_path, temp_output_mov_path, temp_final_webm_path)
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/api/v1/remove-image-background")
|
||||
async def remove_image_background(
|
||||
background_tasks: BackgroundTasks,
|
||||
file: UploadFile = File(...)
|
||||
):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="No file uploaded")
|
||||
|
||||
# We will use secure temp files
|
||||
temp_input_fd, temp_input_path = tempfile.mkstemp(suffix=".png")
|
||||
temp_output_png_fd, temp_output_png_path = tempfile.mkstemp(suffix=".png")
|
||||
|
||||
os.close(temp_input_fd)
|
||||
os.close(temp_output_png_fd)
|
||||
|
||||
try:
|
||||
# 1. Save uploaded image to temp file
|
||||
with open(temp_input_path, "wb") as buffer:
|
||||
shutil.copyfileobj(file.file, buffer)
|
||||
|
||||
print(f"Saved uploaded image to {temp_input_path}")
|
||||
|
||||
# 2. Run backgroundremover for image
|
||||
# backgroundremover -i input.jpg -o output.png
|
||||
cmd_bg_remover = [
|
||||
"backgroundremover",
|
||||
"-i", temp_input_path,
|
||||
"-o", temp_output_png_path
|
||||
]
|
||||
|
||||
print("Starting image background removal process...")
|
||||
process_bg = await asyncio.create_subprocess_exec(
|
||||
*cmd_bg_remover,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
stdout, stderr = await process_bg.communicate()
|
||||
|
||||
if process_bg.returncode != 0:
|
||||
print(f"backgroundremover error: {stderr.decode()}")
|
||||
raise HTTPException(status_code=500, detail="Error processing image background removal.")
|
||||
|
||||
print("Image background removal finished successfully.")
|
||||
|
||||
# 3. Schedule cleanup of ALL temporary files after response is sent
|
||||
background_tasks.add_task(
|
||||
cleanup_files,
|
||||
temp_input_path,
|
||||
temp_output_png_path
|
||||
)
|
||||
|
||||
# 4. Return the file
|
||||
return FileResponse(
|
||||
temp_output_png_path,
|
||||
media_type="image/png",
|
||||
filename=f"transparent_{file.filename.split('.')[0]}.png"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# If an error occurs, clean up immediately
|
||||
cleanup_files(temp_input_path, temp_output_png_path)
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,4 @@
|
||||
fastapi==0.111.0
|
||||
uvicorn[standard]==0.29.0
|
||||
python-multipart==0.0.9
|
||||
backgroundremover==0.2.2
|
||||
Reference in New Issue
Block a user