haritsahm
commited on
Commit
·
563d5aa
1
Parent(s):
24b1db2
hotfix remove unused args
Browse files- app.py +1 -1
- utils/utils.py +1 -1
app.py
CHANGED
@@ -10,7 +10,7 @@ from PIL import Image
|
|
10 |
from streamlit_drawable_canvas import st_canvas
|
11 |
from utils import utils
|
12 |
|
13 |
-
PREDICTOR_MODEL, AUTOMASK_MODEL = utils.get_model('checkpoint/medsam_vit_b.pth')
|
14 |
|
15 |
|
16 |
def process_box(predictor_model, show_mask, radius_width):
|
|
|
10 |
from streamlit_drawable_canvas import st_canvas
|
11 |
from utils import utils
|
12 |
|
13 |
+
PREDICTOR_MODEL, AUTOMASK_MODEL = utils.get_model(checkpoint='checkpoint/medsam_vit_b.pth')
|
14 |
|
15 |
|
16 |
def process_box(predictor_model, show_mask, radius_width):
|
utils/utils.py
CHANGED
@@ -11,7 +11,7 @@ def get_color():
|
|
11 |
|
12 |
|
13 |
@st.cache_resource
|
14 |
-
def get_model(
|
15 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
16 |
build_sam = sam_model_registry['vit_b']
|
17 |
model = build_sam(checkpoint=checkpoint)
|
|
|
11 |
|
12 |
|
13 |
@st.cache_resource
|
14 |
+
def get_model(checkpoint='checkpoint/sam_vit_b_01ec64.pth'):
|
15 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
16 |
build_sam = sam_model_registry['vit_b']
|
17 |
model = build_sam(checkpoint=checkpoint)
|