Update app.py

app.py

@@ -1,28 +1,64 @@
 import gradio as gr
-from PIL import Image
 from rembg import remove
+from PIL import Image
 import io
 
-def remove_background(input_image):
-    # Ensure image is RGBA for transparency
-    input_image = input_image.convert("RGBA")
-    
-    # Remove background
-    result = remove(input_image)
-    
-    return Image.open(io.BytesIO(result))
-
-with gr.Blocks(title="AI Background Remover (rembg)") as demo:
-    gr.Markdown("🧠 **AI Background Remover** - Powered by `rembg` for transparent results")
-    
+# Preserve original resolution & optimize for high-res images
+def process_image(image, background_type, bg_color, bg_image):
+    try:
+        input_image = Image.open(image).convert("RGBA")
+        # Use rembg with bytes instead of loading into RAM completely
+        buffered = io.BytesIO()
+        input_image.save(buffered, format="PNG")
+        result = remove(buffered.getvalue())
+        output = Image.open(io.BytesIO(result)).convert("RGBA")
+
+        if background_type == "Solid Color":
+            bg = Image.new("RGBA", output.size, bg_color)
+            bg.paste(output, (0, 0), output)
+            return bg
+
+        elif background_type == "Custom Image" and bg_image:
+            custom_bg = Image.open(bg_image).convert("RGBA").resize(output.size)
+            custom_bg.paste(output, (0, 0), output)
+            return custom_bg
+
+        else:
+            return output
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# For multiple images
+def batch_process(images, background_type, bg_color, bg_image):
+    return [process_image(img, background_type, bg_color, bg_image) for img in images]
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎯 High-Resolution AI Background Remover with Custom Background")
+
+    with gr.Row():
+        images_input = gr.File(label="📤 Upload High-Res Image(s)", file_types=["image"], file_count="multiple")
+
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Upload Image", type="pil")
-            submit = gr.Button("Remove Background")
-        
-        with gr.Column():
-            output_image = gr.Image(label="Transparent PNG Output")
-
-    submit.click(fn=remove_background, inputs=input_image, outputs=output_image)
+        background_type = gr.Radio(
+            ["Transparent", "Solid Color", "Custom Image"],
+            label="Choose Background Type",
+            value="Transparent"
+        )
+        bg_color = gr.ColorPicker(label="Solid Color", value="#ffffff", visible=False)
+        bg_image = gr.File(label="Custom Background Image", file_types=["image"], visible=False)
+
+    output_gallery = gr.Gallery(label="🖼️ Output Images", columns=2, height=600)
+
+    def toggle_visibility(bg_type):
+        return (
+            gr.update(visible=(bg_type == "Solid Color")),
+            gr.update(visible=(bg_type == "Custom Image"))
+        )
+
+    background_type.change(toggle_visibility, inputs=background_type, outputs=[bg_color, bg_image])
+
+    run_btn = gr.Button("🚀 Start Processing")
+    run_btn.click(fn=batch_process, inputs=[images_input, background_type, bg_color, bg_image], outputs=output_gallery)
 
 demo.launch()