brandly/services/background-remover/main.py

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))