vmoras commited on
Commit
a86e443
1 Parent(s): 0ad7527

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -6,7 +6,7 @@ import numpy as np
6
  from transformers import SamModel, SamProcessor
7
 
8
 
9
- model = SamModel.from_pretrained('facebook/sam-vit-huge').to('cuda')
10
  processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
11
 
12
 
@@ -14,7 +14,7 @@ def set_predictor(image):
14
  """
15
  Creates a Sam predictor object based on a given image and model.
16
  """
17
- inputs = processor(image, return_tensors='pt').to('cuda')
18
  image_embedding = model.get_image_embeddings(inputs['pixel_values'])
19
 
20
  return [image, image_embedding, 'Done']
@@ -27,7 +27,7 @@ def get_polygon(points, image, image_embedding):
27
  """
28
  points = [int(w) for w in points.split(',')]
29
 
30
- inputs = processor(image, input_boxes=[points], return_tensors="pt").to('cuda')
31
 
32
  # pop the pixel_values as they are not neded
33
  inputs.pop("pixel_values", None)
 
6
  from transformers import SamModel, SamProcessor
7
 
8
 
9
+ model = SamModel.from_pretrained('facebook/sam-vit-huge').to('cpu')
10
  processor = SamProcessor.from_pretrained('facebook/sam-vit-huge')
11
 
12
 
 
14
  """
15
  Creates a Sam predictor object based on a given image and model.
16
  """
17
+ inputs = processor(image, return_tensors='pt').to('cpu')
18
  image_embedding = model.get_image_embeddings(inputs['pixel_values'])
19
 
20
  return [image, image_embedding, 'Done']
 
27
  """
28
  points = [int(w) for w in points.split(',')]
29
 
30
+ inputs = processor(image, input_boxes=[points], return_tensors="pt").to('cpu')
31
 
32
  # pop the pixel_values as they are not neded
33
  inputs.pop("pixel_values", None)