Update app.py

app.py

@@ -1,115 +1,29 @@
 import gradio as gr
-import torch
-import torch.nn.functional as F
-import numpy as np
+from transformers import pipeline
 from PIL import Image
-import cv2
-import os
 
-# Ensure models directory is accessible
-try:
-    from models.isnet import ISNetGT
-except ImportError:
-    raise ImportError("Could not import ISNetGT from models.isnet. Ensure models/isnet.py is in the Space.")
-
-# Define model loading function
-def load_model(model_path="isnet-general-use.pth"):
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model file {model_path} not found. Upload it to the Space root directory.")
-    
-    model = ISNetGT()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.to(device).eval()
-    return model, device
-
-# Image preprocessing function
-def preprocess_image(image, target_size=(1024, 1024)):
-    # Convert PIL Image to numpy array
-    image = np.array(image)
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+# Função para remover o background da imagem
+def remove_background(image):
+    # Inicializar o pipeline de segmentação de imagem
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
     
-    # Resize image while preserving aspect ratio
-    h, w = image.shape[:2]
-    scale = min(target_size[0] / h, target_size[1] / w)
-    new_h, new_w = int(h * scale), int(w * scale)
-    image_resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
+    # Obter a máscara da imagem
+    pillow_mask = pipe(image, return_mask=True)
     
-    # Pad to target size
-    padded_image = np.zeros((target_size[0], target_size[1], 3), dtype=np.uint8)
-    padded_image[:new_h, :new_w] = image_resized
+    # Aplicar máscara na imagem original
+    pillow_image = pipe(image)
     
-    # Normalize and convert to tensor
-    image_tensor = torch.from_numpy(padded_image).permute(2, 0, 1).float() / 255.0
-    image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
-    
-    return image_tensor, (new_h, new_w), (h, w)
-
-# Inference function
-def inference(model, image_tensor, device):
-    image_tensor = image_tensor.to(device)
-    with torch.no_grad():
-        output = model(image_tensor)[0]  # Get segmentation output
-        output = F.interpolate(output, size=image_tensor.shape[2:], mode='bilinear', align_corners=True)
-        output = torch.sigmoid(output).cpu().numpy()[0, 0]  # Convert to probability map
-    return output
-
-# Post-processing function
-def postprocess_output(output, original_size, resized_size):
-    # Resize mask to resized image size, then to original size
-    mask = cv2.resize(output, resized_size[::-1], interpolation=cv2.INTER_LANCZOS4)
-    mask = cv2.resize(mask, original_size[::-1], interpolation=cv2.INTER_LANCZOS4)
-    mask = (mask > 0.5).astype(np.uint8) * 255  # Binarize mask
-    return mask
+    return pillow_image
 
-# Background removal function
-def remove_background(input_image):
-    if input_image is None:
-        return None
-    
-    try:
-        # Load model
-        model, device = load_model()
-        
-        # Preprocess image
-        image_tensor, resized_size, original_size = preprocess_image(input_image)
-        
-        # Run inference
-        mask = inference(model, image_tensor, device)
-        
-        # Post-process mask
-        mask = postprocess_output(mask, original_size, resized_size)
-        
-        # Apply mask to create transparent image
-        input_array = np.array(input_image)
-        alpha = mask
-        rgba = np.zeros((input_array.shape[0], input_array.shape[1], 4), dtype=np.uint8)
-        rgba[..., :3] = input_array
-        rgba[..., 3] = alpha
-        
-        # Convert to PIL Image
-        output_image = Image.fromarray(rgba, mode='RGBA')
-        return output_image
-    
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Set up Gradio Blocks interface
-with gr.Blocks(title="DIS Background Remover") as demo:
-    gr.Markdown("## DIS Background Remover")
-    gr.Markdown("Upload an image to remove its background using the IS-Net model from xuebinqin/DIS.")
-    
-    with gr.Row():
-        input_image = gr.Image(type="pil", label="Upload Image")
-        output_image = gr.Image(type="pil", label="Image with Background Removed")
-    
-    submit_btn = gr.Button("Remove Background")
-    submit_btn.click(
-        fn=remove_background,
-        inputs=input_image,
-        outputs=output_image
-    )
+# Criar uma interface Gradio
+app = gr.Interface(
+    fn=remove_background,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.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
+# Iniciar a interface
 if __name__ == "__main__":
-    demo.launch()
+    app.launch(share=True)