Update app.py

app.py

@@ -2,16 +2,25 @@ import gradio as gr
 from PIL import Image
 import os
 import time
+import json
+from fastapi import FastAPI
+from pydantic import BaseModel
 
 from src.utils import change_background, matte
 
+app = FastAPI()
+
+class ImageRequest(BaseModel):
+    image: str  # Base64 encoded image
+    background_color: str
+
 def process_images(uploaded_files, background_color):
     # Add debugging statement
     print(f"background_color received: {background_color}")
 
     # Default value check
     if background_color is None:
-        return "Please select a background color."
+        background_color = "Transparent (PNG)"
 
     hexmap = {
         "Transparent (PNG)": "#000000",
@@ -38,6 +47,33 @@ def process_images(uploaded_files, background_color):
     
     return results, "\n".join(times)
 
+@app.post("/process")
+def api_process_images(request: ImageRequest):
+    image_data = request.image
+    background_color = request.background_color
+    
+    # Convert base64 to PIL image
+    img_input = Image.open(BytesIO(base64.b64decode(image_data)))
+    
+    hexmap = {
+        "Transparent (PNG)": "#000000",
+        "Black": "#000000",
+        "White": "#FFFFFF",
+        "Green": "#22EE22",
+        "Red": "#EE2222",
+        "Blue": "#2222EE",
+    }
+    alpha = 0.0 if background_color == "Transparent (PNG)" else 1.0
+
+    img_matte = matte(img_input)
+    img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[background_color])
+    
+    buffered = BytesIO()
+    img_output.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    return {"image": img_str}
+
 with gr.Blocks() as demo:
     gr.Markdown("""
     # AI Photo Background Removal
@@ -48,7 +84,7 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         folder_input = gr.Files(type="filepath", label="Upload your photos here")
-        bg_color = gr.Dropdown(choices=["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"], label="Choose background color")
+        bg_color = gr.Dropdown(choices=["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"], label="Choose background color", value="Transparent (PNG)")
     
     output_images = gr.Gallery(label="Processed Images")
     processing_times = gr.Textbox(label="Processing Times")
@@ -63,4 +99,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()