Update app.py

app.py

@@ -1,33 +1,36 @@
 import gradio as gr
+from rembg import remove
 from PIL import Image
-import random
+import io
 import os
 
-# Dummy classifier function (replace with real model logic)
-def classify_parrot(image):
-    # Dummy label list
-    labels = ["Hyacinth Macaw", "Eclectus Parrot", "Cockatiel", "Kakapo", "Budgerigar", "Palm Cockatoo", "Sulphur-crested Cockatoo"]
-    prediction = random.choice(labels)
-    output_img_path = "output_result.png"
-    image.save(output_img_path)
-    return output_img_path, f"Predicted Species: {prediction}"
+# Background removal function
+def remove_background(image):
+    input_bytes = io.BytesIO()
+    image.save(input_bytes, format='PNG')
+    input_bytes = input_bytes.getvalue()
 
-# Example images (replace with your actual trial images)
+    output_bytes = remove(input_bytes)
+    output_img = Image.open(io.BytesIO(output_bytes))
+
+    output_path = "output.png"
+    output_img.save(output_path, format="PNG")
+
+    return output_path, "Background removed successfully!", output_path
+
+# Example images (replace with actual sample image paths)
 example_images = [
-    "examples/hyacinth.jpg",
-    "examples/eclectus.jpg",
-    "examples/cockatiel.jpg",
-    "examples/kakapo.jpg",
-    "examples/macaw.jpg",
-    "examples/sulphur.jpg",
-    "examples/green_parrot.jpg",
-    "examples/palm_cockatoo.jpg",
-    "examples/yellow_parrot.jpg",
+    "examples/person1.jpg",
+    "examples/product1.jpg",
+    "examples/cat.jpg",
+    "examples/dog.jpg",
+    "examples/portrait.jpg",
+    "examples/bird.jpg"
 ]
 
-with gr.Blocks(title="Parrot60 Classifier") as demo:
-    gr.Markdown("## 🦜 Parrot60 Classifier")
-    gr.Markdown("This classifier can classify 60 parrot species with 93% accuracy.")
+with gr.Blocks(title="Background Remover") as demo:
+    gr.Markdown("# 🖼️ Background Remover")
+    gr.Markdown("Upload an image or use an example to remove its background.")
 
     with gr.Row():
         with gr.Column():
@@ -36,11 +39,11 @@ with gr.Blocks(title="Parrot60 Classifier") as demo:
             clear_btn = gr.Button("Clear")
         with gr.Column():
             output_image = gr.Image(label="Output", type="filepath")
-            output_text = gr.Textbox(label="Prediction", lines=1)
+            output_text = gr.Textbox(label="Status", lines=1)
             download_output = gr.File(label="Download Output")
 
     with gr.Row():
-        gr.Markdown("### 🧪 Examples")
+        gr.Markdown("### 🧪 Example Images")
     with gr.Row():
         gr.Examples(
             examples=example_images,
@@ -48,12 +51,8 @@ with gr.Blocks(title="Parrot60 Classifier") as demo:
             label="Click to try an example",
         )
 
-    def run_pipeline(img):
-        result_path, label = classify_parrot(img)
-        return result_path, label, result_path
-
-    submit_btn.click(fn=run_pipeline, inputs=image_input, outputs=[output_image, output_text, download_output])
+    # Hook logic
+    submit_btn.click(fn=remove_background, inputs=image_input, outputs=[output_image, output_text, download_output])
     clear_btn.click(fn=lambda: (None, "", None), outputs=[image_input, output_text, download_output])
 
-# Launch app
 demo.launch()