Gilad Avidan commited on
Commit
78cd05e
1 Parent(s): cffdace

move to smaller model

Browse files
__pycache__/handler.cpython-311.pyc ADDED
Binary file (3.1 kB). View file
 
handler.py CHANGED
@@ -9,8 +9,8 @@ class EndpointHandler():
9
  def __init__(self, path=""):
10
  # Preload all the elements you are going to need at inference.
11
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
- self.model = SamModel.from_pretrained("facebook/sam-vit-huge").to(self.device)
13
- self.processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
14
 
15
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
16
  """
 
9
  def __init__(self, path=""):
10
  # Preload all the elements you are going to need at inference.
11
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ self.model = SamModel.from_pretrained("facebook/sam-vit-base").to(self.device)
13
+ self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
14
 
15
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
16
  """
test_handler.py CHANGED
@@ -2,7 +2,7 @@ import base64
2
  from handler import EndpointHandler
3
 
4
  my_handler = EndpointHandler(path=".")
5
- image = base64.b64encode(open("./car.png", "rb").read())
6
 
7
  input = {"image": image, "points": [[450, 600]]}
8
  result = my_handler(input)
 
2
  from handler import EndpointHandler
3
 
4
  my_handler = EndpointHandler(path=".")
5
+ image = base64.b64encode(open("../sam-vit-huge/car.png", "rb").read())
6
 
7
  input = {"image": image, "points": [[450, 600]]}
8
  result = my_handler(input)