Update app.py
Browse files
app.py
CHANGED
@@ -8,7 +8,7 @@ from itertools import chain
|
|
8 |
from transformers import SamModel, SamProcessor
|
9 |
|
10 |
|
11 |
-
model = SamModel.from_pretrained('facebook/sam-vit-huge')
|
12 |
processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
|
13 |
|
14 |
|
@@ -16,7 +16,7 @@ def set_predictor(image):
|
|
16 |
"""
|
17 |
Creates a Sam predictor object based on a given image and model.
|
18 |
"""
|
19 |
-
device = '
|
20 |
inputs = processor(image, return_tensors='pt').to(device)
|
21 |
image_embedding = model.get_image_embeddings(inputs['pixel_values'])
|
22 |
|
@@ -28,11 +28,9 @@ def get_polygon(points, image, image_embedding):
|
|
28 |
Returns the points of the polygon given a bounding box and a prediction
|
29 |
made by Sam.
|
30 |
"""
|
31 |
-
#points = [int(w) for w in points.split(',')]
|
32 |
points = list(chain.from_iterable(points))
|
33 |
-
print(points)
|
34 |
|
35 |
-
device = '
|
36 |
inputs = processor(image, input_boxes=[points], return_tensors="pt").to(device)
|
37 |
|
38 |
# pop the pixel_values as they are not neded
|
|
|
8 |
from transformers import SamModel, SamProcessor
|
9 |
|
10 |
|
11 |
+
model = SamModel.from_pretrained('facebook/sam-vit-huge')
|
12 |
processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
|
13 |
|
14 |
|
|
|
16 |
"""
|
17 |
Creates a Sam predictor object based on a given image and model.
|
18 |
"""
|
19 |
+
device = 'cpu'
|
20 |
inputs = processor(image, return_tensors='pt').to(device)
|
21 |
image_embedding = model.get_image_embeddings(inputs['pixel_values'])
|
22 |
|
|
|
28 |
Returns the points of the polygon given a bounding box and a prediction
|
29 |
made by Sam.
|
30 |
"""
|
|
|
31 |
points = list(chain.from_iterable(points))
|
|
|
32 |
|
33 |
+
device = 'cpu'
|
34 |
inputs = processor(image, input_boxes=[points], return_tensors="pt").to(device)
|
35 |
|
36 |
# pop the pixel_values as they are not neded
|