oscarfu0501 commited on
Commit
a0403c5
1 Parent(s): 452921c

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -7,15 +7,15 @@ from ultralytics import YOLO
7
 
8
 
9
  class EndpointHandler:
10
- def __init__(self, key): # pass api key to model
11
- pass
12
 
13
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
14
  inputs = data.get("inputs")
15
  isurl = inputs.get("isurl")
16
  path = inputs.get("path")
17
 
18
- model = YOLO("MakeNTU-v2024/Image-preprocessor/best.pt")
19
  ########################### Load Image #################################
20
  if(isurl): # for url set isurl = 1
21
  print("checkpoint 2-1")
 
7
 
8
 
9
  class EndpointHandler:
10
+ def __init__(self, path='.'): # pass api key to model
11
+ self.model = YOLO("./best.pt")
12
 
13
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
14
  inputs = data.get("inputs")
15
  isurl = inputs.get("isurl")
16
  path = inputs.get("path")
17
 
18
+ model = self.model
19
  ########################### Load Image #################################
20
  if(isurl): # for url set isurl = 1
21
  print("checkpoint 2-1")