diff --git a/ultralytics/solutions/serve.py b/ultralytics/solutions/serve.py index a8e31a81c3..ddc265f48f 100644 --- a/ultralytics/solutions/serve.py +++ b/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):