Update api_server.py

api_server.py

@@ -53,22 +53,6 @@ app = Flask(__name__)
 # API route for prediction(YOLO)
 @app.route('/predict', methods=['POST'])
 def predict():
-    """
-    Predicts the class label of an input image.
-
-    Request format:
-    {
-        "image": [[pixel_values_gray]]
-    }
-
-    Response format:
-    {
-        "label": predicted_label,
-        "pred_proba" prediction class probability
-        "ml-latency-ms": latency_in_milliseconds
-            (Measures time only for ML operations preprocessing with predict)
-    }
-    """
     if 'image' not in request.files:
         # Handle if no file is selected
         return 'No file selected'
@@ -83,23 +67,34 @@ def predict():
     except Exception as e:
         return jsonify({'error': str(e)}), 400
 
-    # 將圖像儲存到一個緩衝區
-    img_io = io.BytesIO()
-    image_data.save(img_io, 'PNG')
-    img_io.seek(0)
-
-    # 返回圖像
-    return send_file(img_io, mimetype='image/png')
-
-    # # Check image shape
-    # if image_data.size != (28, 28):
-    #     return "Invalid image shape. Expected (28, 28), take from 'demo images' folder."
+    # Make a prediction using YOLO
+    results = model(image_data)
+    
+    # 準備返回多張圖像
+    images_io = []
+    
+    for i, result_img in enumerate(results.render()):  # 假設 results.render() 返回的是 PIL Image 格式的圖像
+        img_io = io.BytesIO()
+        result_img.save(img_io, 'PNG')  # 儲存 YOLO 處理過的圖像到緩衝區
+        img_io.seek(0)
+        images_io.append((f'image_{i}.png', img_io))  # 使用名稱區分不同的圖像
+        send_file(img_io, mimetype='image/png')
+
+    # 打包多張圖像為 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)
 
-    # # Make a prediction using YOLO
-    # results = model(image_data)
+
 
     # # Process the YOLO output
     # detections = []