Update app.py

app.py

@@ -1,96 +1,60 @@
-from typing import List
-import gradio as gr
-import numpy as np
+import os
+import cv2
 import supervision as sv
-import torch
-from PIL import Image
+import numpy as np
+import gradio as gr
 from transformers import pipeline
+from PIL import Image
 
-# Global Variables
-MARKDOWN = """
-# SAM - Softly Activated Masks
-"""
-EXAMPLES = [
-    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "dog", 0.5],
-    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "building", 0.5],
-    ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "jacket", 0.5],
-    ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "coffee", 0.6],
-]
-
-MIN_AREA_THRESHOLD = 0.01
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Definición de la clase SamAutomaticMaskGenerator
+class SamAutomaticMaskGenerator:
+    def __init__(self, sam_pipeline):
+        self.sam_pipeline = sam_pipeline
 
-# Initialize SAM Generator with exception handling
-try:
-    SAM_GENERATOR = pipeline(
-        task="mask-generation",
-        model="facebook/sam-vit-large",
-        device=DEVICE
-    )
-except Exception as e:
-    print(f"Error initializing SAM generator: {e}")
+    def generate(self, image_rgb):
+        outputs = self.sam_pipeline(image_rgb, points_per_batch=32)
+        mask = np.array(outputs['masks'], dtype=np.uint8)
+        return mask
 
-# Mask Annotators
-SEMITRANSPARENT_MASK_ANNOTATOR = sv.MaskAnnotator(
-    color=sv.Color.red(),
-    color_lookup=sv.ColorLookup.INDEX
+# Configuración del modelo SAM
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+sam_pipeline = pipeline(
+    task="mask-generation",
+    model="facebook/sam-vit-large",
+    device=DEVICE
 )
 
-SOLID_MASK_ANNOTATOR = sv.MaskAnnotator(
-    color=sv.Color.white(),
-    color_lookup=sv.ColorLookup.INDEX,
-    opacity=1
-)
+mask_generator = SamAutomaticMaskGenerator(sam_pipeline)
 
-# Functions
-def run_sam(image_rgb_pil: Image.Image) -> sv.Detections:
-    try:
-        outputs = SAM_GENERATOR(image_rgb_pil, points_per_batch=32)
-        mask = np.array(outputs['masks'], dtype=np.uint8)
-        return sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
-    except Exception as e:
-        print(f"Error running SAM model: {e}")
-        return sv.Detections(xyxy=[], mask=[])
+# Función para procesar y anotar la imagen
+def process_image(image_pil):
+    image_rgb = np.array(image_pil)
+    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
 
-def reverse_mask_image(image: np.ndarray, mask: np.ndarray, gray_value=128):
-    gray_color = np.array([gray_value, gray_value, gray_value], dtype=np.uint8)
-    return np.where(mask[..., None], image, gray_color)
+    sam_result = mask_generator.generate(image_rgb)
+    mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
+    detections = sv.Detections.from_sam(sam_result=sam_result)
 
-def inference(image_rgb_pil: Image.Image) -> List[Image.Image]:
-    width, height = image_rgb_pil.size
-    area = width * height
+    annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)
+    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
-    detections = run_sam(image_rgb_pil)
-    detections = detections[detections.area / area > MIN_AREA_THRESHOLD]
+    return Image.fromarray(image_rgb), Image.fromarray(annotated_image_rgb)
 
-    blank_image = Image.new("RGB", (width, height), "black")
-    return [
-        SEMITRANSPARENT_MASK_ANNOTATOR.annotate(image_rgb_pil, detections),
-        SOLID_MASK_ANNOTATOR.annotate(blank_image, detections)
-    ]
-#************    
-#GRADIO CONSTRUCTION
+# Construcción de la interfaz Gradio
 with gr.Blocks() as demo:
-    gr.Markdown(MARKDOWN)
+    gr.Markdown("# SAM - Segmentación de Imágenes")
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(image_mode='RGB', type='pil', height=500)
-            submit_button = gr.Button("Pruébalo!!!")
-        gallery = gr.Gallery(label="Result", object_fit="scale-down", preview=True)
-
-    with gr.Row():
-        gr.Examples(
-            examples=EXAMPLES,
-            fn=inference,
-            inputs=[input_image],
-            outputs=[gallery],
-            cache_examples=False,
-            run_on_click=True
-        )
+            input_image = gr.Image(type="pil", label="Cargar Imagen")
+            submit_button = gr.Button("Segmentar")
+        with gr.Column():
+            original_image = gr.Image(type="pil", label="Imagen Original")
+            segmented_image = gr.Image(type="pil", label="Imagen Segmentada")
+    
     submit_button.click(
-        inference,
-        inputs=[input_image],
-        outputs=gallery
+        process_image, 
+        inputs=input_image, 
+        outputs=[original_image, segmented_image]
     )
 
-demo.launch(debug=False, show_error=True)
+demo.launch(debug=True)