import os

from ultralytics import YOLO
from flask import Flask, send_file

app = Flask(__name__)


def print_file_tree(root_dir, indent=""):
    """
    打印文件树
    """
    for root, dirs, files in os.walk(root_dir):
        # 打印当前目录
        print(indent + f"📁 {root}")

        # 打印当前目录下的文件
        for file in files:
            print(indent + f"  📄 {file}")

        # 递归处理子目录
        for subdir in dirs:
            print_file_tree(os.path.join(root, subdir), indent + "  ")

@app.route('/predict', methods=['GET'])
def predict():
    model = YOLO('yolov8s.pt')
    # model.val(data='./data.yaml', imgsz=640)
    results = model.train(data='/app/data.yaml', epochs=200, imgsz=640)
    model.export(format='onnx')


@app.route('/download/<filename>')
def download_file(filename):
    # 在这里进行安全性检查，确保只能访问允许的文件

    # 定义文件路径
    file_path = f'{filename}'

    # 使用send_file函数发送文件
    return send_file(file_path, as_attachment=True)

#print_file_tree("/app")
app.run(host='0.0.0.0', port=7860, debug=True)