Update api_server.py

api_server.py

@@ -2,7 +2,7 @@ import os
 import time
 import numpy as np
 from PIL import Image
-
+import torchvision.transforms as transforms
 from pathlib import Path
 
 # Disable tensorflow warnings
@@ -10,35 +10,37 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
 from tensorflow import keras
 from flask import Flask, jsonify, request, render_template
+import torch
 
-load_type = 'remote_hub_from_pretrained'
-"""
-local;
-remote_hub_download; 
-remote_hub_from_pretrained;
-remote_hub_pipeline; - needs config.json and this is not easy to grasp how to do it with custom models
-https://discuss.huggingface.co/t/how-to-create-a-config-json-after-saving-a-model/10459/4
-"""
+load_type = 'local'
 
-REPO_ID = "1vash/mnist_demo_model"
+MODEL_NAME = "yolo11_detect_best_241018_1.pt"
 MODEL_DIR = "./artifacts/models"
+#REPO_ID = "1vash/mnist_demo_model"
 
-# Load the saved model into memory
+# Load the saved YOLO model into memory
 if load_type == 'local':
-    model = keras.models.load_model(f'{MODEL_DIR}/mnist_model.h5')
+    # 本地模型路徑
+    model_path = f'{MODEL_DIR}/{MODEL_NAME}'
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found at {model_path}")
+    # 使用 torch 來載入 YOLO 模型
+    model = torch.load(model_path)
+    model.eval()  # 設定模型為推理模式
 elif load_type == 'remote_hub_download':
     from huggingface_hub import hf_hub_download
 
-    model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename="saved_model.pb"))
+    # 從 Hugging Face Hub 下載模型
+    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME)
+    model = torch.load(model_path)
+    model.eval()
 elif load_type == 'remote_hub_from_pretrained':
-    # https://huggingface.co/docs/hub/keras
+    # 使用 Hugging Face Hub 預訓練的模型方式下載
     os.environ['TRANSFORMERS_CACHE'] = str(Path(MODEL_DIR).absolute())
-    from huggingface_hub import from_pretrained_keras
-    model = from_pretrained_keras(REPO_ID, cache_dir=MODEL_DIR)
-elif load_type == 'remote_hub_pipeline':
-    from transformers import pipeline
+    from huggingface_hub import from_pretrained
 
-    model = pipeline("image-classification", model=REPO_ID)
+    model = from_pretrained(REPO_ID, filename=MODEL_NAME, cache_dir=MODEL_DIR)
+    model.eval()
 else:
     raise AssertionError('No load type is specified!')
 
@@ -46,7 +48,7 @@ else:
 app = Flask(__name__)
 
 
-# API route for prediction
+# API route for prediction(YOLO)
 @app.route('/predict', methods=['POST'])
 def predict():
     """
@@ -76,27 +78,35 @@ def predict():
     # Get pixels out of file
     image_data = Image.open(file)
 
-    # Check image shape
-    if image_data.size != (28, 28):
-        return "Invalid image shape. Expected (28, 28), take from 'demo images' folder."
+    # # 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, verbose=0 to disable progress bar in logs
-    prediction = model.predict(processed_image, verbose=0)
-
-    # Get the predicted class label
-    predicted_label = np.argmax(prediction)
-    proba = prediction[0][predicted_label]
+    # 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 prediction result and latency as dictionary response
+    # Return the detection results and latency as JSON response
     response = {
-        'label': int(predicted_label),
-        'pred_proba': float(proba),
+        'detections': detections,
         'ml-latency-ms': round(latency_ms, 4)
     }
 
@@ -108,16 +118,26 @@ def predict():
 
 # Helper function to preprocess the image
 def preprocess_image(image_data):
-    """Preprocess image for Model Inference
+    """Preprocess image for YOLO Model Inference
 
-    :param image_data: Raw image
-    :return: image: Preprocessed Image
+    :param image_data: Raw image (PIL.Image)
+    :return: image: Preprocessed Image (Tensor)
     """
-    # Resize the image to match the input shape of the model
-    image = np.array(image_data).reshape(1, 28, 28)
+    # 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)
 
-    # Normalize the pixel values
-    image = image.astype('float32') / 255.0
+    # Add batch dimension (1, C, H, W) since YOLO expects a batch
+    image = image.unsqueeze(0)
 
     return image