Refactor: remove AGPL imgly dependency and migrate background removal to python backend
This commit is contained in:
@@ -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))
|
||||
Reference in New Issue
Block a user