From 71028b5bcd2acd35dca42c619cf46d5cbf01f724 Mon Sep 17 00:00:00 2001
From: Francesco Mattioli <Francesco.mttl@gmail.com>
Date: Wed, 9 Oct 2024 18:25:23 +0200
Subject: [PATCH] fixed server and cli runner

---
 ultralytics/solutions/serve.py | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/ultralytics/solutions/serve.py b/ultralytics/solutions/serve.py
index 1856170e25..a8e31a81c3 100644
--- a/ultralytics/solutions/serve.py
+++ b/ultralytics/solutions/serve.py
@@ -14,7 +14,6 @@ from ultralytics.engine.results import Results
 class UltralyticsRequest(BaseModel):
     image: str
 
-
 class UltralyticsResponse(BaseModel):
     status: str = "success"
     results: List[List[Dict]]
@@ -22,11 +21,22 @@ class UltralyticsResponse(BaseModel):
 
 class YOLOServe(ls.LitAPI):
     def __init__(self, model) -> None:
-        self.model = model
+        """
+        Litserve API for YOLO model:
+        Call order is:
+            setup -> batch -> decode_request -> predict -> unbatch -> encode_response
+
+        Args:
+            model (str): Model name to use
+        """
+        self.model_name = model
         super().__init__()
 
     def setup(self, device):
-        self.model = YOLO(model=self.model)
+        self.model = YOLO(model=self.model_name)
+        
+    def batch(self, inputs):
+        return list(inputs)
 
     def decode_request(self, request: UltralyticsRequest):
         base64_str = request.image
@@ -34,22 +44,24 @@ class YOLOServe(ls.LitAPI):
         image = Image.open(io.BytesIO(img_data))
         image_array = np.array(image)
         return image_array
-
-    def encode_response(self, results: Results) -> UltralyticsResponse:
-        image_results = []
-        for result in results:
-            image_results.append(result.to_json())
-        return UltralyticsResponse(results=[r.to_dict() for r in results])
-
+    
     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):
     if "model" in args:
         model = args["model"]
     else:
         model = "yolo11n.pt"
-    server = ls.LitServer(YOLOServe(model), accelerator="auto", max_batch_size=1, batch_timeout=0.05)
+    server = ls.LitServer(YOLOServe(model), accelerator="auto", max_batch_size=15, batch_timeout=0.05)
     server.run(port=8000)