Update app.py

app.py

@@ -1,41 +1,70 @@
-import gradio as gr
-from loadimg import load_img
-#import spaces
-from transformers import AutoModelForImageSegmentation
+import os
 import torch
-from torchvision import transforms
-from typing import Union, Tuple
 from PIL import Image
+from typing import Union, Tuple
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+import gradio as gr
+from loadimg import load_img
+
+# =========================================================================
+# CONFIGURACIÓN DE DISPOSITIVO (CPU)
+# =========================================================================
+DEVICE = "cpu"
+
+print(f"--- Cargando BiRefNet en {DEVICE.upper()} ---")
 
+# Cargamos el modelo directamente del Hub de Hugging Face
 birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "merve/BiRefNet", low_cpu_mem_usage=False,  trust_remote_code=True, torch_dtype=torch.float32, device_map=None
-)
-birefnet = birefnet.eval()
-#birefnet.to("cuda")
-
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
+    "merve/BiRefNet", 
+    trust_remote_code=True, 
+    torch_dtype=torch.float32
+).to(DEVICE)
 
-def fn(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]:
-    """
-    Remove the background from an image and return both the transparent version and the original.
+birefnet.eval()
+print("Modelo cargado correctamente en CPU.")
 
-    This function performs background removal using a BiRefNet segmentation model. It is intended for use
-    with image input (either uploaded or from a URL). The function returns a transparent PNG version of the image
-    with the background removed, along with the original RGB version for comparison.
+# Transformaciones necesarias para el modelo
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
 
-    Args:
-        image (PIL.Image or str): The input image, either as a PIL object or a filepath/URL string.
+# =========================================================================
+# FUNCIONES DE PROCESAMIENTO
+# =========================================================================
+
+def process(image: Image.Image) -> Image.Image:
+    """
+    Aplica BiRefNet para remover el fondo de la imagen usando CPU.
+    """
+    image_size = image.size
+    
+    # 1. Preparar el tensor para la red
+    input_tensor = transform_image(image).unsqueeze(0).to(DEVICE)
+    
+    # 2. Inferencia (Paso por la red neuronal sin almacenar gradientes)
+    with torch.no_grad():
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+        
+    # 3. Crear la máscara Alfa
+    mask = preds[0].squeeze()
+    mask_pil = transforms.ToPILImage()(mask)
+    
+    # 4. Ajustar máscara al tamaño original con alta calidad (LANCZOS)
+    mask_final = mask_pil.resize(image_size, Image.LANCZOS)
+    
+    # 5. Aplicar transparencia a la imagen original
+    output_image = image.copy()
+    output_image.putalpha(mask_final)
+    
+    return output_image
 
-    Returns:
-        tuple:
-            - origin (PIL.Image): The original RGB image, unchanged.
-            - processed_image (PIL.Image): The input image with the background removed and transparency applied.
+def fn(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]:
+    """
+    Función para las pestañas de Gradio (Subida de Imagen y URL).
+    Devuelve la imagen original y la versión procesada para el ImageSlider.
     """
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
@@ -43,64 +72,50 @@ def fn(image: Union[Image.Image, str]) -> Tuple[Image.Image, Image.Image]:
     processed_image = process(im)
     return (origin, processed_image)
 
-#@spaces.GPU
-def process(image: Image.Image) -> Image.Image:
-    """
-    Apply BiRefNet-based image segmentation to remove the background.
-
-    This function preprocesses the input image, runs it through a BiRefNet segmentation model to obtain a mask,
-    and applies the mask as an alpha (transparency) channel to the original image.
-
-    Args:
-        image (PIL.Image): The input RGB image.
-
-    Returns:
-        PIL.Image: The image with the background removed, using the segmentation mask as transparency.
-    """
-    image_size = image.size
-    input_images = transform_image(image).unsqueeze(0)
-    with torch.inference_mode():
-        preds = birefnet(input_images)[-1].sigmoid().detach().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-    image.putalpha(mask)
-    return image
-
 def process_file(f: str) -> str:
     """
-    Load an image file from disk, remove the background, and save the output as a transparent PNG.
-
-    Args:
-        f (str): Filepath of the image to process.
-
-    Returns:
-        str: Path to the saved PNG image with background removed.
+    Función para la pestaña de archivos. Guarda y devuelve la ruta del PNG.
     """
     name_path = f.rsplit(".", 1)[0] + ".png"
     im = load_img(f, output_type="pil")
     im = im.convert("RGB")
     transparent = process(im)
-    transparent.save(name_path)
+    transparent.save(name_path, "PNG")
     return name_path
 
+# =========================================================================
+# INTERFAZ GRADIO
+# =========================================================================
+
 slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
 slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
+
 image_upload = gr.Image(label="Upload an image")
 image_file_upload = gr.Image(label="Upload an image", type="filepath")
 url_input = gr.Textbox(label="Paste an image URL")
 output_file = gr.File(label="Output PNG File")
 
-# Example images
-chameleon = load_img("butterfly.jpg", output_type="pil")
+# Ejemplos por defecto
+example_image_path = "butterfly.jpg"
 url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
 
-tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
+# Carga segura de la imagen de ejemplo local para evitar crasheos si no se ha subido aún
+try:
+    chameleon = load_img(example_image_path, output_type="pil")
+    examples_img = [chameleon]
+    examples_file = [example_image_path]
+except Exception:
+    examples_img = None
+    examples_file = None
+
+tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=examples_img, api_name="image")
 tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
-tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
+tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=examples_file, api_name="png")
 
 demo = gr.TabbedInterface(
-    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
+    [tab1, tab2, tab3], 
+    ["Image Upload", "URL Input", "File Output"], 
+    title="Background Removal Tool (CPU Edition)"
 )
 
 if __name__ == "__main__":