Update api_server.py

api_server.py

@@ -76,71 +76,82 @@ def predict():
 
     file = request.files['image']
 
-    # Get pixels out of file
-    image_data = Image.open(file)
+    # 讀取圖像
+    try:
+        image_data = Image.open(file)
+    except Exception as e:
+        return jsonify({'error': str(e)}), 400
+
+    # 將圖像儲存到一個緩衝區
+    img_io = io.BytesIO()
+    image.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."
 
-    # Preprocess the image
-    processed_image = preprocess_image(image_data)
-
-    # Make a prediction using YOLO
-    results = model(processed_image)
-
-    # Process the YOLO output
-    detections = []
-    for det in results.xyxy[0]:  # Assuming results are in xyxy format (xmin, ymin, xmax, ymax, confidence, class)
-        x_min, y_min, x_max, y_max, confidence, class_idx = det
-        width = x_max - x_min
-        height = y_max - y_min
-        detection = {
-            "label": int(class_idx),
-            "confidence": float(confidence),
-            "bbox": [float(x_min), float(y_min), float(width), float(height)]
-        }
-        detections.append(detection)
-
-    # Calculate latency in milliseconds
-    latency_ms = (time.time() - start_time) * 1000
-
-    # Return the detection results and latency as JSON response
-    response = {
-        'detections': detections,
-        'ml-latency-ms': round(latency_ms, 4)
-    }
-
-    # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary
-    # flask.jsonify vs json.dumps https://sentry.io/answers/difference-between-json-dumps-and-flask-jsonify/
-    # The flask.jsonify() function returns a Response object with Serializable JSON and content_type=application/json.
-    return jsonify(response)
-
-
-# Helper function to preprocess the image
-def preprocess_image(image_data):
-    """Preprocess image for YOLO Model Inference
-
-    :param image_data: Raw image (PIL.Image)
-    :return: image: Preprocessed Image (Tensor)
-    """
-    # Define the YOLO input size (example 640x640, you can modify this based on your model)
-    input_size = (640, 640)
-
-    # Define transformation: Resize the image, convert to Tensor, and normalize pixel values
-    transform = transforms.Compose([
-        transforms.Resize(input_size),       # Resize to YOLO input size
-        transforms.ToTensor(),               # Convert image to PyTorch Tensor (通道數、影像高度和寬度)
-        transforms.Normalize([0.0, 0.0, 0.0], [1.0, 1.0, 1.0])  # Normalization (if needed)
-    ])
-
-    # Apply transformations to the image
-    image = transform(image_data)
-
-    # Add batch dimension (1, C, H, W) since YOLO expects a batch
-    image = image.unsqueeze(0)
-
-    return image
+    # # Preprocess the image
+    # processed_image = preprocess_image(image_data)
+
+    # # Make a prediction using YOLO
+    # results = model(image_data)
+
+    # # Process the YOLO output
+    # detections = []
+    # for det in results.xyxy[0]:  # Assuming results are in xyxy format (xmin, ymin, xmax, ymax, confidence, class)
+    #     x_min, y_min, x_max, y_max, confidence, class_idx = det
+    #     width = x_max - x_min
+    #     height = y_max - y_min
+    #     detection = {
+    #         "label": int(class_idx),
+    #         "confidence": float(confidence),
+    #         "bbox": [float(x_min), float(y_min), float(width), float(height)]
+    #     }
+    #     detections.append(detection)
+
+    # # Calculate latency in milliseconds
+    # latency_ms = (time.time() - start_time) * 1000
+
+    # # Return the detection results and latency as JSON response
+    # response = {
+    #     'detections': detections,
+    #     'ml-latency-ms': round(latency_ms, 4)
+    # }
+
+    # # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary
+    # # flask.jsonify vs json.dumps https://sentry.io/answers/difference-between-json-dumps-and-flask-jsonify/
+    # # The flask.jsonify() function returns a Response object with Serializable JSON and content_type=application/json.
+    # return jsonify(response)
+
+
+# # Helper function to preprocess the image
+# def preprocess_image(image_data):
+#     """Preprocess image for YOLO Model Inference
+
+#     :param image_data: Raw image (PIL.Image)
+#     :return: image: Preprocessed Image (Tensor)
+#     """
+#     # Define the YOLO input size (example 640x640, you can modify this based on your model)
+#     input_size = (640, 640)
+
+#     # Define transformation: Resize the image, convert to Tensor, and normalize pixel values
+#     transform = transforms.Compose([
+#         transforms.Resize(input_size),       # Resize to YOLO input size
+#         transforms.ToTensor(),               # Convert image to PyTorch Tensor (通道數、影像高度和寬度)
+#         transforms.Normalize([0.0, 0.0, 0.0], [1.0, 1.0, 1.0])  # Normalization (if needed)
+#     ])
+
+#     # Apply transformations to the image
+#     image = transform(image_data)
+
+#     # Add batch dimension (1, C, H, W) since YOLO expects a batch
+#     image = image.unsqueeze(0)
+
+#     return image
 
 
 # API route for health check