fireedman commited on
Commit
1991283
1 Parent(s): 605391e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -78
app.py CHANGED
@@ -1,96 +1,60 @@
1
- from typing import List
2
- import gradio as gr
3
- import numpy as np
4
  import supervision as sv
5
- import torch
6
- from PIL import Image
7
  from transformers import pipeline
 
8
 
9
- # Global Variables
10
- MARKDOWN = """
11
- # SAM - Softly Activated Masks
12
- """
13
- EXAMPLES = [
14
- ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "dog", 0.5],
15
- ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "building", 0.5],
16
- ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "jacket", 0.5],
17
- ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "coffee", 0.6],
18
- ]
19
-
20
- MIN_AREA_THRESHOLD = 0.01
21
- DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
22
 
23
- # Initialize SAM Generator with exception handling
24
- try:
25
- SAM_GENERATOR = pipeline(
26
- task="mask-generation",
27
- model="facebook/sam-vit-large",
28
- device=DEVICE
29
- )
30
- except Exception as e:
31
- print(f"Error initializing SAM generator: {e}")
32
 
33
- # Mask Annotators
34
- SEMITRANSPARENT_MASK_ANNOTATOR = sv.MaskAnnotator(
35
- color=sv.Color.red(),
36
- color_lookup=sv.ColorLookup.INDEX
 
 
37
  )
38
 
39
- SOLID_MASK_ANNOTATOR = sv.MaskAnnotator(
40
- color=sv.Color.white(),
41
- color_lookup=sv.ColorLookup.INDEX,
42
- opacity=1
43
- )
44
 
45
- # Functions
46
- def run_sam(image_rgb_pil: Image.Image) -> sv.Detections:
47
- try:
48
- outputs = SAM_GENERATOR(image_rgb_pil, points_per_batch=32)
49
- mask = np.array(outputs['masks'], dtype=np.uint8)
50
- return sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
51
- except Exception as e:
52
- print(f"Error running SAM model: {e}")
53
- return sv.Detections(xyxy=[], mask=[])
54
 
55
- def reverse_mask_image(image: np.ndarray, mask: np.ndarray, gray_value=128):
56
- gray_color = np.array([gray_value, gray_value, gray_value], dtype=np.uint8)
57
- return np.where(mask[..., None], image, gray_color)
58
 
59
- def inference(image_rgb_pil: Image.Image) -> List[Image.Image]:
60
- width, height = image_rgb_pil.size
61
- area = width * height
62
 
63
- detections = run_sam(image_rgb_pil)
64
- detections = detections[detections.area / area > MIN_AREA_THRESHOLD]
65
 
66
- blank_image = Image.new("RGB", (width, height), "black")
67
- return [
68
- SEMITRANSPARENT_MASK_ANNOTATOR.annotate(image_rgb_pil, detections),
69
- SOLID_MASK_ANNOTATOR.annotate(blank_image, detections)
70
- ]
71
- #************
72
- #GRADIO CONSTRUCTION
73
  with gr.Blocks() as demo:
74
- gr.Markdown(MARKDOWN)
75
  with gr.Row():
76
  with gr.Column():
77
- input_image = gr.Image(image_mode='RGB', type='pil', height=500)
78
- submit_button = gr.Button("Pruébalo!!!")
79
- gallery = gr.Gallery(label="Result", object_fit="scale-down", preview=True)
80
-
81
- with gr.Row():
82
- gr.Examples(
83
- examples=EXAMPLES,
84
- fn=inference,
85
- inputs=[input_image],
86
- outputs=[gallery],
87
- cache_examples=False,
88
- run_on_click=True
89
- )
90
  submit_button.click(
91
- inference,
92
- inputs=[input_image],
93
- outputs=gallery
94
  )
95
 
96
- demo.launch(debug=False, show_error=True)
 
1
+ import os
2
+ import cv2
 
3
  import supervision as sv
4
+ import numpy as np
5
+ import gradio as gr
6
  from transformers import pipeline
7
+ from PIL import Image
8
 
9
+ # Definición de la clase SamAutomaticMaskGenerator
10
+ class SamAutomaticMaskGenerator:
11
+ def __init__(self, sam_pipeline):
12
+ self.sam_pipeline = sam_pipeline
 
 
 
 
 
 
 
 
 
13
 
14
+ def generate(self, image_rgb):
15
+ outputs = self.sam_pipeline(image_rgb, points_per_batch=32)
16
+ mask = np.array(outputs['masks'], dtype=np.uint8)
17
+ return mask
 
 
 
 
 
18
 
19
+ # Configuración del modelo SAM
20
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
21
+ sam_pipeline = pipeline(
22
+ task="mask-generation",
23
+ model="facebook/sam-vit-large",
24
+ device=DEVICE
25
  )
26
 
27
+ mask_generator = SamAutomaticMaskGenerator(sam_pipeline)
 
 
 
 
28
 
29
+ # Función para procesar y anotar la imagen
30
+ def process_image(image_pil):
31
+ image_rgb = np.array(image_pil)
32
+ image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
 
 
 
 
 
33
 
34
+ sam_result = mask_generator.generate(image_rgb)
35
+ mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
36
+ detections = sv.Detections.from_sam(sam_result=sam_result)
37
 
38
+ annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)
39
+ annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
40
 
41
+ return Image.fromarray(image_rgb), Image.fromarray(annotated_image_rgb)
 
42
 
43
+ # Construcción de la interfaz Gradio
 
 
 
 
 
 
44
  with gr.Blocks() as demo:
45
+ gr.Markdown("# SAM - Segmentación de Imágenes")
46
  with gr.Row():
47
  with gr.Column():
48
+ input_image = gr.Image(type="pil", label="Cargar Imagen")
49
+ submit_button = gr.Button("Segmentar")
50
+ with gr.Column():
51
+ original_image = gr.Image(type="pil", label="Imagen Original")
52
+ segmented_image = gr.Image(type="pil", label="Imagen Segmentada")
53
+
 
 
 
 
 
 
 
54
  submit_button.click(
55
+ process_image,
56
+ inputs=input_image,
57
+ outputs=[original_image, segmented_image]
58
  )
59
 
60
+ demo.launch(debug=True)