import os
import uuid
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from ultralytics import YOLO

app = Flask(__name__)
CORS(app)  # 允许所有域名进行访问

# 模型字典，可以根据需要更新
models = {
    '追踪': 'models/detect/yolov8n.pt',
    '检测': 'models/detect/弹珠模型.pt',
    '分类': 'models/classify/yolov8n-cls.pt',
    '姿势': 'models/pose/yolov8n-pose.pt',
    '分割': 'models/segment/yolov8n-seg.pt'
  }

model_instances = {}

def load_model(model_path):
    try:
        return YOLO(model_path)  # 直接调用训练好的模型
    except Exception as e:
        print(f"Failed to load model from {model_path}: {e}")
        return None

@app.route('/')
def home():
    return "Welcome to the YOLOv8 Flask App!"

@app.route('/request', methods=['POST'])
def handle_request():
    try:
        selected_model = request.form.get('model')
        if selected_model not in models:
            return jsonify({'error': 'Invalid model selected.'}), 400

        model_path = models[selected_model]
        if selected_model not in model_instances:
            model_instances[selected_model] = load_model(model_path)
        model = model_instances[selected_model]

        if model is None:
            return jsonify({'error': 'Model is not loaded due to connection error.'}), 500

        img = request.files.get('img')
        if img is None:
            return jsonify({'error': 'No image provided.'}), 400

        img_name = str(uuid.uuid4()) + '.jpg'
        img_path = os.path.join('./img', img_name)
        os.makedirs(os.path.dirname(img_path), exist_ok=True)
        img.save(img_path)

        print(f"Image saved at {img_path}")

        save_dir = './runs/detect'
        os.makedirs(save_dir, exist_ok=True)
        results = model.predict(img_path, save=True, project=save_dir, name=img_name.split('.')[0], device='cpu')  # Device should be 'cpu' or 'cuda'

        # 获取预测结果路径
        predicted_img_path = os.path.join(save_dir, img_name.split('.')[0], img_name)

        if not os.path.exists(predicted_img_path):
            print(f"Prediction result not found at {predicted_img_path}")
            return jsonify({'error': 'Prediction result not found.'}), 500

        print(f"Prediction successful. Result saved at {predicted_img_path}")
        return jsonify({'message': 'Prediction successful!', 'img_path': f'/get/{img_name}'})

    except Exception as e:
        print(f"Error during request processing: {e}")
        return jsonify({'error': f'An error occurred during processing: {e}'}), 500

@app.route('/get/<filename>', methods=['GET'])
def get_file(filename):
    try:
        save_dir = './runs/detect'
        predicted_img_path = os.path.join(save_dir, filename.split('.')[0], filename)
        print(f"Trying to send file from path: {predicted_img_path}")

        if not os.path.exists(predicted_img_path):
            return jsonify({'error': 'Prediction result not found.'}), 404

        return send_file(predicted_img_path, mimetype='image/jpeg')

    except Exception as e:
        print(f"Error during file retrieval: {e}")
        return jsonify({'error': f'Failed to retrieve the result image. Error: {e}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)  # 允许从任何网络接口访问