Spaces:
Runtime error
Runtime error
File size: 3,369 Bytes
3e2a4df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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) # 允许从任何网络接口访问
|