api_server.py

import os
import time
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from pathlib import Path
from ultralytics import YOLO
import io
import base64

# Disable tensorflow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

from tensorflow import keras
from flask import Flask, jsonify, request, render_template, send_file
import torch

load_type = 'local'

MODEL_NAME = "yolo11_detect_best_241018_1.pt"
MODEL_DIR = "./artifacts/models"
YOLO_DIR = "./artifacts/yolo"
#REPO_ID = "1vash/mnist_demo_model"

# Load the saved YOLO model into memory
if load_type == 'local':
    # 本地模型路徑
    model_path = f'{MODEL_DIR}/{MODEL_NAME}'
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found at {model_path}")
        
    model = YOLO(model_path)
    #model.eval()  # 設定模型為推理模式
elif load_type == 'remote_hub_download':
    from huggingface_hub import hf_hub_download

    # 從 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':
    # 使用 Hugging Face Hub 預訓練的模型方式下載
    os.environ['TRANSFORMERS_CACHE'] = str(Path(MODEL_DIR).absolute())
    from huggingface_hub import from_pretrained

    model = from_pretrained(REPO_ID, filename=MODEL_NAME, cache_dir=MODEL_DIR)
    model.eval()
else:
    raise AssertionError('No load type is specified!')


def image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
        return encoded_string


# Initialize the Flask application
app = Flask(__name__)


# API route for prediction(YOLO)
@app.route('/predict', methods=['POST'])
def predict():

    user_id = request.args.get('user_id') 
    
    if 'image' not in request.files:
        # Handle if no file is selected
        return 'No file selected'

    start_time = time.time()

    file = request.files['image']

    # 讀取圖像
    try:
        image_data = Image.open(file)
    except Exception as e:
        return jsonify({'error': str(e)}), 400

    # Make a prediction using YOLO
    results = model(image_data)
    
    # 檢查 YOLO 是否返回了有效的結果
    if results is None or len(results) == 0:
        return jsonify({'error': 'No results from YOLO model'}), 400

    # # 渲染推理結果到圖像
    # img_with_boxes = results[0].plot()  # 使用 results[0]，假設只有一張圖像作推理

    # # 將 numpy array 轉換為 PIL Image
    # img = Image.fromarray(img_with_boxes)

    # # 儲存圖片到內存緩衝區
    # img_io = io.BytesIO()
    # img.save(img_io, 'PNG')
    # img_io.seek(0)

    # # 返回處理後的圖像
    # return send_file(img_io, mimetype='image/png')

    saved_images = []
    
    # 儲存辨識後的圖片到指定資料夾
    for result in results:
        # 保存圖片
        result.save_crop(f"{YOLO_DIR}/{user_id}")
        
        num_detections = len(result.boxes)  # Get the number of detections
        labels = result.boxes.cls  # Get predicted label IDs
        label_names = [model.names[int(label)] for label in labels]  # Convert to names

        encoded_images=[]

        for label_name in label_names:
            output_file=f"{YOLO_DIR}/{user_id}/{label_name}/im.jpg.jpg"
            # 將圖片轉換為 base64 編碼
            encoded_images.append(image_to_base64(output_file))
            
        # 建立回應資料
        response_data = {
            'images': encoded_images,
            'description': label_names
        }
    return jsonify(response_data)


    # # 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
@app.route('/health', methods=['GET'])
def health():
    """
    Health check API to ensure the application is running.
    Returns "OK" if the application is healthy.
    Demo Usage: "curl http://localhost:5000/health" or using alias "curl http://127.0.0.1:5000/health"
    """
    return 'OK'


# API route for version
@app.route('/version', methods=['GET'])
def version():
    """
    Returns the version of the application.
    Demo Usage: "curl http://127.0.0.1:5000/version" or using alias "curl http://127.0.0.1:5000/version"
    """
    return '1.0'


@app.route("/")
def hello_world():
    return render_template("index.html")
    # return "<p>Hello, Team!</p>"


# Start the Flask application
if __name__ == '__main__':
    app.run(debug=True)