Update app.py

app.py

@@ -1,18 +1,14 @@
 from typing import List
-
 import gradio as gr
 import numpy as np
 import supervision as sv
 import torch
-
 from PIL import Image
-from transformers import pipeline, CLIPProcessor, CLIPModel
-
+from transformers import pipeline
 
-#************    
-#Variables globales
+# Global Variables
 MARKDOWN = """
-#SAM
+# SAM - Softly Activated Masks
 """
 EXAMPLES = [
     ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "dog", 0.5],
@@ -22,130 +18,79 @@ EXAMPLES = [
 ]
 
 MIN_AREA_THRESHOLD = 0.01
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 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}")
 
-DEVICE = "cuda" if  torch.cuda.is_available() else "cpu"
-
-SAM_GENERATOR = pipeline(
-    task = "mask-generation",
-    model = "facebook/sam-vit-large",
-    device = DEVICE
-)
-
+# Mask Annotators
 SEMITRANSPARENT_MASK_ANNOTATOR = sv.MaskAnnotator(
-    color = sv.Color.red(),
-    color_lookup = sv.ColorLookup.INDEX
+    color=sv.Color.red(),
+    color_lookup=sv.ColorLookup.INDEX
 )
 
 SOLID_MASK_ANNOTATOR = sv.MaskAnnotator(
-    color = sv.Color.white(),
-    color_lookup = sv.ColorLookup.INDEX,
-    opacity = 1
+    color=sv.Color.white(),
+    color_lookup=sv.ColorLookup.INDEX,
+    opacity=1
 )
 
-
-#************    
-#funciones de trabajo
-
-def run_sam(image_rgb_pil : Image.Image ) -> sv.Detections:
-    outputs = SAM_GENERATOR(image_rgb_pil, points_per_batch = 32)
-    mask = np.array(outputs['masks'])    
-    return sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
-
+# 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=[])
 
 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)
+    gray_color = np.array([gray_value, gray_value, gray_value], dtype=np.uint8)
     return np.where(mask[..., None], image, gray_color)
 
-
-"""
-def filter_detections(image_rgb_pil: Image.Image, detections: sv.Detections) -> sv.Detections:
-    img_rgb_numpy = np.array(image_rgb_pil)
-    filtering_mask = []
-    for xyxy, mask in zip(detections.xyxy, detections.mask):
-        crop = sv.crop_image(
-            image = img_rgb_numpy,
-            xyxy =xyxy
-        )
-        mask_crop = sv.crop_image(
-            image=mask,
-            xyxy=xyxy
-        )
-        masked_crop = reverse_mask_image(
-            image=crop,
-            mask=mask_crop
-        )
-        
-    filtering_mask = np.array(
-        filtering_mask
-    )
-    return detections[filtering_mask]
-"""
-        
-def inference (image_rgb_pil: Image.Image) -> List[Image.Image]:
+def inference(image_rgb_pil: Image.Image) -> List[Image.Image]:
     width, height = image_rgb_pil.size
     area = width * height
 
-    detections = run_sam(
-        image_rgb_pil
-    )
-    detections = detections[ detections.area /area > MIN_AREA_THRESHOLD ]
-    
-    #detections = filter_detections(
-    #    image_rgb_pil=image_rgb_pil,
-    #    detections=detections,
-    #)
-    
+    detections = run_sam(image_rgb_pil)
+    detections = detections[detections.area / area > MIN_AREA_THRESHOLD]
+
     blank_image = Image.new("RGB", (width, height), "black")
     return [
-        annotate(
-            image_rgb_pil=image_rgb_pil,
-            detections=detections,
-            annotator=SEMITRANSPARENT_MASK_ANNOTATOR),
-        annotate(
-            image_rgb_pil=blank_image,
-            detections=detections,
-            annotator=SOLID_MASK_ANNOTATOR)
+        SEMITRANSPARENT_MASK_ANNOTATOR.annotate(image_rgb_pil, detections),
+        SOLID_MASK_ANNOTATOR.annotate(blank_image, detections)
     ]
-
-    
 #************    
 #GRADIO CONSTRUCTION
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                image_mode = 'RGB',
-                type = 'pil',
-                height = 500
-            )
+            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
-        )
+        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
+            examples=EXAMPLES,
+            fn=inference,
+            inputs=[input_image],
+            outputs=[gallery],
+            cache_examples=False,
+            run_on_click=True
         )
     submit_button.click(
         inference,
-        inputs = [
-            input_image
-        ],
-        outputs = gallery
+        inputs=[input_image],
+        outputs=gallery
     )
 
-demo.launch( debug = True, show_error = True )   
+demo.launch(debug=False, show_error=True)