Upload app.py

app.py

@@ -0,0 +1,93 @@
+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)  # 允许从任何网络接口访问