Update api_server.py

api_server.py

@@ -17,7 +17,8 @@ import torch
 load_type = 'local'
 
 MODEL_NAME = "yolo11_detect_best_241018_1.pt"
-MODEL_DIR = "./artifacts/models"
+MODEL_DIR = "./artifacts/models".
+YOLO_DIR = "./artifacts/yolo"
 #REPO_ID = "1vash/mnist_demo_model"
 
 # Load the saved YOLO model into memory
@@ -73,12 +74,13 @@ def predict():
     # 準備返回多張圖像
     images_io = []
     
-    for i, result_img in enumerate(results.render()):  # 假設 results.render() 返回的是 PIL Image 格式的圖像
+   # YOLOv5/YOLOv8 results have a .ims attribute that contains the rendered images
+    for i, img in enumerate(results.ims):  # Assuming `ims` contains the rendered images
         img_io = io.BytesIO()
-        result_img.save(img_io, 'PNG')  # 儲存 YOLO 處理過的圖像到緩衝區
+        img = Image.fromarray(img)  # Convert numpy array back to PIL image if necessary
+        img.save(img_io, 'PNG')  # Save YOLO-processed image to buffer
         img_io.seek(0)
-        images_io.append((f'image_{i}.png', img_io))  # 使用名稱區分不同的圖像
-        send_file(img_io, mimetype='image/png')
+        images_io.append((f'image_{i}.png', img_io))  # Store each image with its filename
 
     # 打包多張圖像為 ZIP 文件進行返回
     zip_io = io.BytesIO()