''' from flask import Flask, request, jsonify, send_from_directory, render_template from flask_cors import CORS from ultralytics import YOLO import gradio as gr from threading import Thread import os import uuid import logging from PIL import Image # 配置日志记录 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s:%(message)s', datefmt='%Y-%m-%d %H:%M:%S') # 创建 Flask 应用 app = Flask(__name__, static_folder='static') CORS(app) # 定义模型路径 models = { '追踪': 'models/yolov8n.pt', '检测': 'models/danzhu.pt', '分类': 'models/yolov8n-cls.pt', '姿势': 'models/yolov8n-pose.pt', '分割': 'models/yolov8n-seg.pt' } model_instances = {} def load_model(model_path): """加载模型""" try: logging.info(f"正在从 {model_path} 加载模型...") model = YOLO(model_path) logging.info(f"模型从 {model_path} 成功加载") return model except Exception as e: logging.error(f"从 {model_path} 加载模型失败: {e}") return None def convert_image_format(img_path, target_format='JPEG'): """转换图像格式""" try: with Image.open(img_path) as img: if img.mode != 'RGB': img = img.convert('RGB') base_name, _ = os.path.splitext(img_path) target_path = f"{base_name}.{target_format.lower()}" img.save(target_path, format=target_format) logging.info(f"图像格式成功转换为 {target_format}，保存到 {target_path}") return target_path except Exception as e: logging.error(f"图像格式转换失败: {e}") raise def predict(model_name, img_path): """进行预测""" try: if model_name not in models: logging.error("选择的模型无效。") return "选择的模型无效。" model_path = models[model_name] if model_name not in model_instances: model_instances[model_name] = load_model(model_path) model = model_instances[model_name] if model is None: logging.error("由于连接错误，模型未加载。") return "由于连接错误，模型未加载。" unique_name = str(uuid.uuid4()) save_dir = './runs/detect' os.makedirs(save_dir, exist_ok=True) logging.info(f"保存目录: {save_dir}") # 转换图像格式 img_path_converted = convert_image_format(img_path, 'JPEG') img_path_converted = os.path.normpath(img_path_converted) logging.info(f"对 {img_path_converted} 进行预测...") results = model.predict(img_path_converted, save=True, project=save_dir, name=unique_name, device='cpu') logging.info(f"预测结果: {results}") result_dir = os.path.join(save_dir, unique_name) result_dir = os.path.normpath(result_dir) logging.info(f"结果目录: {result_dir}") if not os.path.exists(result_dir): logging.error(f"结果目录 {result_dir} 不存在") return "未找到预测结果。" # 查找预测结果文件 predicted_img_path = None for file in os.listdir(result_dir): if file.lower().endswith(('.jpeg', '.jpg')): predicted_img_path = os.path.join(result_dir, file) break if predicted_img_path: logging.info(f"找到预测图像: {predicted_img_path}") return predicted_img_path else: logging.error(f"在 {result_dir} 中未找到预测图像") return "未找到预测结果。" except Exception as e: logging.error(f"预测过程中出错: {e}") return f"预测过程中出错: {e}" # 定义 Gradio 界面 iface = gr.Interface( fn=predict, inputs=[ gr.Dropdown(choices=list(models.keys()), label="选择模型"), gr.Image(type="filepath", label="输入图像") ], outputs=gr.Image(type="filepath", label="输出图像") ) @app.route('/') def home(): """主页""" return render_template('index.html') @app.route('/request', methods=['POST']) def handle_request(): """处理请求""" try: selected_model = request.form.get('model') if selected_model not in models: logging.error("选择的模型无效。") return jsonify({'error': '选择的模型无效。'}), 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: logging.error("由于连接错误，模型未加载。") return jsonify({'error': '由于连接错误，模型未加载。'}), 500 img = request.files.get('img') if img is None: logging.error("未提供图像。") return jsonify({'error': '未提供图像。'}), 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) logging.info(f"图像已保存到: {img_path}") save_dir = './runs/detect' os.makedirs(save_dir, exist_ok=True) unique_name = str(uuid.uuid4()) logging.info(f"对 {img_path} 进行预测...") results = model.predict(img_path, save=True, project=save_dir, name=unique_name, device='cpu') logging.info(f"预测结果: {results}") result_dir = os.path.join(save_dir, unique_name) # 查找预测结果文件 predicted_img_path = None for file in os.listdir(result_dir): if file.endswith('.jpeg') or file.endswith('.jpg'): predicted_img_path = os.path.join(result_dir, file) break if predicted_img_path: img_url = f'/get/{unique_name}/{os.path.basename(predicted_img_path)}' return jsonify({'message': '预测成功！', 'img_path': img_url}) else: saved_files = os.listdir(result_dir) logging.error(f"保存目录中包含文件: {saved_files}") return jsonify({'error': '未找到预测结果。'}), 500 except Exception as e: logging.error(f"处理请求时出错: {e}") return jsonify({'error': f'处理过程中发生错误: {e}'}), 500 @app.route('/get//') def get_image(unique_name, filename): """获取图像""" try: return send_from_directory(os.path.join('runs/detect', unique_name), filename) except Exception as e: logging.error(f"提供文件时出错: {e}") return jsonify({'error': '文件未找到。'}), 404 def run_gradio(): """运行 Gradio 界面""" logging.info("启动 Gradio 界面...") iface.launch(share=True) # 设置 share=True 以便公开访问 def run_flask(): """运行 Flask 应用""" logging.info("启动 Flask 应用...") app.run(host="0.0.0.0", port=5000) if __name__ == '__main__': # 启动 Flask 和 Gradio 线程 gradio_thread = Thread(target=run_gradio) flask_thread = Thread(target=run_flask) gradio_thread.start() flask_thread.start() gradio_thread.join() flask_thread.join() ''' from ultralytics import YOLO # Load a model model = YOLO("yolov8n.yaml") # build a new model from YAML model = YOLO("yolov8n.pt") # load a pretrained model (recommended for training) model = YOLO("yolov8n.yaml").load("yolov8n.pt") # build from YAML and transfer weights # Train the model results = model.train(data="coco8.yaml", epochs=100, imgsz=640)