vmoras commited on
Commit
93d65f4
1 Parent(s): 2c3c995

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -5,18 +5,22 @@ import traceback
5
  import numpy as np
6
  import gradio as gr
7
  from itertools import chain
 
8
  from segment_anything import SamPredictor, sam_model_registry
9
 
10
 
11
- sam_checkpoint = "sam_vit_h_4b8939.pth"
 
 
 
12
  model_type = "vit_h"
 
13
 
14
 
15
  def set_predictor(image):
16
  """
17
  Creates a Sam predictor object based on a given image and model.
18
  """
19
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
20
  sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
21
  sam.to(device=device)
22
 
@@ -84,19 +88,18 @@ with gr.Blocks() as demo:
84
  4. Click on the image where the lower right corner of the bbox should be
85
  5. Check the coordinates using the 'bbox' box
86
  6. Click on 'Send bounding box'
87
- 7. On the right side you will see the binary mask*
88
- 8. on the lower side you will see the points that made up the polygon*
89
  9. Click on 'Clear bbox' to send another bounding box and repeat the steps from the thrid point
90
  10. Repeat steps 3 to 9 until all the segments for this image are done
91
  11. Click on the right corner of the image to remove it and repeat all the steps with the next
92
  image
93
 
94
- * If the binary mask is all black and the polygon is an empty list, it means the program did
95
  not find any segment in the bbox. Make the bbox a little big bigger if that happens.
96
  """)
97
 
98
 
99
-
100
  image = gr.State()
101
  embedding = gr.State()
102
  bbox = gr.State([[0, 0], [0, 0]])
 
5
  import numpy as np
6
  import gradio as gr
7
  from itertools import chain
8
+ from huggingface_hub import hf_hub_download
9
  from segment_anything import SamPredictor, sam_model_registry
10
 
11
 
12
+ hf_hub_download(repo_id="vmoras/sam_api", filename="sam_vit_h.pth")
13
+
14
+
15
+ sam_checkpoint = "sam_vit_h.pth"
16
  model_type = "vit_h"
17
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
18
 
19
 
20
  def set_predictor(image):
21
  """
22
  Creates a Sam predictor object based on a given image and model.
23
  """
 
24
  sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
25
  sam.to(device=device)
26
 
 
88
  4. Click on the image where the lower right corner of the bbox should be
89
  5. Check the coordinates using the 'bbox' box
90
  6. Click on 'Send bounding box'
91
+ 7. On the right side you will see the binary mask ·
92
+ 8. on the lower side you will see the points that made up the polygon ·
93
  9. Click on 'Clear bbox' to send another bounding box and repeat the steps from the thrid point
94
  10. Repeat steps 3 to 9 until all the segments for this image are done
95
  11. Click on the right corner of the image to remove it and repeat all the steps with the next
96
  image
97
 
98
+ · If the binary mask is all black and the polygon is an empty list, it means the program did
99
  not find any segment in the bbox. Make the bbox a little big bigger if that happens.
100
  """)
101
 
102
 
 
103
  image = gr.State()
104
  embedding = gr.State()
105
  bbox = gr.State([[0, 0], [0, 0]])