Update app.py

app.py

@@ -1,37 +1,32 @@
-import requests
+# app.py
+import gradio as gr
+from rembg import remove
 from PIL import Image
-from io import BytesIO
-import torch
-from transformers import DPTForDepthEstimation, DPTFeatureExtractor
+import io
 
-# Example: Load an image
-image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/image_classification.png"
-response = requests.get(image_url)
-image = Image.open(BytesIO(response.content)).convert("RGB")
+# Function to remove background
+def remove_background(image):
+    # Convert uploaded image to PIL
+    pil_image = Image.fromarray(image).convert("RGBA")
+    
+    # Remove background
+    result = remove(pil_image)
+    
+    # Convert result to displayable format
+    return result
 
-# Load Hugging Face DPT depth estimation model (can be used to segment background)
-model_name = "Intel/dpt-large"
-feature_extractor = DPTFeatureExtractor.from_pretrained(model_name)
-model = DPTForDepthEstimation.from_pretrained(model_name)
+# Gradio UI
+title = "AI Background Remover"
+description = "Upload an image and remove its background instantly! Made by Viateur Irasubiza."
 
-# Preprocess image
-inputs = feature_extractor(images=image, return_tensors="pt")
-with torch.no_grad():
-    outputs = model(**inputs)
-    predicted_depth = outputs.predicted_depth
+iface = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="numpy", label="Upload Your Image"),
+    outputs=gr.Image(type="pil", label="Result"),
+    title=title,
+    description=description,
+    allow_flagging="never"
+)
 
-# Convert to numpy and normalize
-depth = predicted_depth.squeeze().cpu().numpy()
-depth_min, depth_max = depth.min(), depth.max()
-normalized_depth = (depth - depth_min) / (depth_max - depth_min)
-
-# Simple threshold to create mask (background = 1)
-import numpy as np
-mask = (normalized_depth < 0.6).astype(np.uint8) * 255
-
-# Apply mask to image
-image_np = np.array(image)
-image_rgba = np.dstack([image_np, mask])  # add alpha channel
-result = Image.fromarray(image_rgba)
-result.save("output.png")
-print("Background removed and saved as output.png")
+# Launch the app
+iface.launch()