vmoras commited on
Commit
91feba4
1 Parent(s): b808c95

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -5
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').to('cuda')
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 = 'cuda'
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 = 'cuda'
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