lemonaddie commited on
Commit
a3aa5f5
·
verified ·
1 Parent(s): 5a419d7

Update app_recon.py

Browse files
Files changed (1) hide show
  1. app_recon.py +8 -1
app_recon.py CHANGED
@@ -88,10 +88,17 @@ def scale_img(img):
88
 
89
  return img
90
 
 
 
91
  def sam_init():
92
  #sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_l_0b3195.pth")
93
  #model_type = "vit_l"
94
- sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
 
 
 
 
 
95
  model_type = "vit_h"
96
 
97
  sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda")
 
88
 
89
  return img
90
 
91
+
92
+
93
  def sam_init():
94
  #sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_l_0b3195.pth")
95
  #model_type = "vit_l"
96
+ sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_vit_h_4b8939.pth")
97
+
98
+ if not os.path.exists(sam_checkpoint):
99
+ print("Downloading SAM-H checkpoint to " + os.path.exists(sam_checkpoint))
100
+ os.system('wget -P' + os.path.dirname(__file__) + ' -nv https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth')
101
+
102
  model_type = "vit_h"
103
 
104
  sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda")