app.py

import os
from rembg import remove
from io import BytesIO
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import StreamingResponse
import gradio as gr

app = FastAPI()

# Background removal function
def remove_background(image_bytes):
    try:
        output_bytes = remove(image_bytes)
        return output_bytes
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Gradio interface
def gradio_interface(input_image):
    if input_image is None:
        return None
    output_bytes = remove_background(input_image)
    return output_bytes

# API endpoint
@app.post("/api/remove-bg")
async def api_remove_bg(file: UploadFile = File(...)):
    if not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    input_image = await file.read()
    output_bytes = remove_background(input_image)
    
    return StreamingResponse(BytesIO(output_bytes), media_type="image/png")

# Gradio app
gr_app = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="filepath", label="Input Image"),
    outputs=gr.Image(type="numpy", label="Output Image"),
    title="Background Remover",
    description="Upload an image to remove the background"
)

# Mount Gradio app
app = gr.mount_gradio_app(app, gr_app, path="/")