Update handler.py
Browse files- handler.py +3 -3
handler.py
CHANGED
@@ -7,15 +7,15 @@ from ultralytics import YOLO
|
|
7 |
|
8 |
|
9 |
class EndpointHandler:
|
10 |
-
def __init__(self,
|
11 |
-
|
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 =
|
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")
|