Update api_server.py

api_server.py

@@ -70,31 +70,28 @@ def predict():
 
     # Make a prediction using YOLO
     results = model(image_data)
-    
-    # 準備返回多張圖像
-    images_io = []
-    
-   # 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()
-        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))  # Store each image with its filename
-
-    # 打包多張圖像為 ZIP 文件進行返回
-    zip_io = io.BytesIO()
-    with zipfile.ZipFile(zip_io, 'w') as zip_file:
-        for filename, image in images_io:
-            zip_file.writestr(filename, image.getvalue())
-    
-    zip_io.seek(0)
 
-    # 返回壓縮包
-    return send_file(zip_io, mimetype='application/zip', as_attachment=True, download_name='predictions.zip')
-
-    # # Preprocess the image
-    # processed_image = preprocess_image(image_data)
+    saved_images = []
+    
+    # 儲存辨識後的圖片到指定資料夾
+    for result in results:
+        # 保存完整圖片
+        yolo_path = YOLO_DIR + f"results_{Path(result.path).name}"
+        result.save(yolo_path)
+        saved_images.append(yolo_path)
+
+        # 保存裁剪後的圖片
+        cropped_images = result.save_crop(YOLO_DIR)  # save_crop usually returns list of cropped image paths
+        saved_images.extend(cropped_images)  # 將裁剪後的圖片加入到返回列表
+
+    end_time = time.time()
+    inference_time = end_time - start_time
+
+    # 返回辨識結果的文件路徑以及推理時間
+    return jsonify({
+        'saved_images': saved_images,
+        'inference_time': inference_time
+    }), 200