Update app.py

app.py

@@ -1,36 +1,27 @@
-from transformers import pipeline
 import gradio as gr
+from transformers import pipeline
 from PIL import Image
-import numpy as np
 
-# Initialize the image segmentation pipeline with the DIS-based model
-pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+# Função para remover o background da imagem
+def remove_background(image):
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+
+    # Obter a máscara da imagem
+    pillow_mask = pipe(image, return_mask=True)
+
+    # Aplicar máscara na imagem original
+    pillow_image = pipe(image)
 
-# Function to remove background
-def remove_background(input_image):
-    # Ensure input is a PIL Image
-    if isinstance(input_image, np.ndarray):
-        input_image = Image.fromarray(input_image)
-    
-    # Get the mask and the processed image
-    pillow_mask = pipe(input_image, return_mask=True)  # Returns a binary mask
-    no_bg_image = pipe(input_image)  # Returns the image with background removed
-    
-    # Ensure the output is a PIL Image with transparency
-    if isinstance(no_bg_image, np.ndarray):
-        no_bg_image = Image.fromarray(no_bg_image)
-    
-    return no_bg_image
+    return pillow_image
 
-# Set up Gradio interface
-demo = gr.Interface(
+# Criar uma interface Gradio
+app = gr.Interface(
     fn=remove_background,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Image with Background Removed"),
-    title="Background Remover using DIS (IS-Net)",
-    description="Upload an image to remove its background using the briaai/RMBG-1.4 model."
+    inputs=gr.components.Image(type="pil"),
+    outputs=gr.components.Image(type="pil", format="png"),  # Especificar saída como PNG
+    title="Remoção de Background de Imagens",
+    description="Envie uma imagem e veja o background sendo removido automaticamente. A imagem resultante será no formato PNG."
 )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    app.launch(share=True)