flask-docker / api_server.py
ning8429's picture
Update api_server.py
b036e6b verified
raw
history blame
6.15 kB
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
# 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!')
# Initialize the Flask application
app = Flask(__name__)
# API route for prediction(YOLO)
@app.route('/predict', methods=['POST'])
def predict():
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
saved_images = []
# 儲存辨識後的圖片到指定資料夾
for result in results:
# 保存完整圖片
yolo_path = f'{YOLO_DIR}/results_{Path(result.path).name}'
result.save(yolo_path)
saved_images.append(yolo_path)
# 保存裁剪後的圖片(僅當 save_crop 返回有效的路徑時才加入)
cropped_images = result.save_crop(YOLO_DIR) # 有些 YOLO 版本 save_crop 不返回值
if cropped_images: # 確保不會對 None 進行迭代
if isinstance(cropped_images, list): # 如果它返回一個列表
saved_images.extend(cropped_images)
else:
saved_images.append(cropped_images)
end_time = time.time()
inference_time = end_time - start_time
# 返回辨識結果的文件路徑以及推理時間
return jsonify({
'saved_images': saved_images,
'inference_time': inference_time
}), 200
# # 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
@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)