UltralyticsAssistant 2 months ago
parent 71028b5bcd
commit 6b3dc62f32
  1. 11
      ultralytics/solutions/serve.py

@ -14,6 +14,7 @@ from ultralytics.engine.results import Results
class UltralyticsRequest(BaseModel):
image: str
class UltralyticsResponse(BaseModel):
status: str = "success"
results: List[List[Dict]]
@ -24,7 +25,7 @@ class YOLOServe(ls.LitAPI):
"""
Litserve API for YOLO model:
Call order is:
setup -> batch -> decode_request -> predict -> unbatch -> encode_response
setup -> batch -> decode_request -> predict -> unbatch -> encode_response.
Args:
model (str): Model name to use
@ -34,7 +35,7 @@ class YOLOServe(ls.LitAPI):
def setup(self, device):
self.model = YOLO(model=self.model_name)
def batch(self, inputs):
return list(inputs)
@ -44,18 +45,16 @@ class YOLOServe(ls.LitAPI):
image = Image.open(io.BytesIO(img_data))
image_array = np.array(image)
return image_array
def predict(self, x):
result = self.model(x)
return result
def unbatch(self, output):
return list(output)
def encode_response(self, result: List[Results]) -> UltralyticsResponse:
return UltralyticsResponse(results=[r.to_dict() for r in result])
def run(args):

Loading…
Cancel
Save