psy/train_api.py

import os
import sys
root_path = os.getcwd()
sys.path.append(root_path)

import time
import datetime
import signal
import uvicorn
import pandas as pd
from fastapi import FastAPI, Request
from pydantic import BaseModel
from typing import List
from utils.common import train_detect
from fastapi.middleware.cors import CORSMiddleware
import logging
import matplotlib.pyplot as plt
import argparse
import numpy as np
import yaml
import threading
import pickle
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from utils.feature_process import create_feature_df, apply_feature_weights, Features

app = FastAPI()

# 定义fastapi返回类
class ClassificationResult(BaseModel):
    precision: list
    recall: list
    f1: list
    wrong_percentage: float

# 允许所有域名的跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["*"],
)

# 定义接口
@app.post("/train/")
async def classify_features(request: Request, features_list: List[Features]):
    # 遍历每个特征对象，并将其添加到 all_features 中
    all_features = create_feature_df(features_list)

    # 读取 YAML 配置文件
    config_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "config/config.yaml"))
    with open(config_path, 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
        feature_names = config['feature_names']
        feature_weights = config['feature_weights']

    # 应用特征权重
    feature_label_weighted = apply_feature_weights(all_features, feature_names, feature_weights)

    start_time = time.time()  # 记录开始时间

    # 创建静态文件存放文件夹
    static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "train_api"))  # 设置模型文件和配置文件的存放目录，和本py同级
    os.makedirs(static_dir, exist_ok=True)

    # 训练前设置
    now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    data_path = os.path.abspath(os.path.join(os.path.dirname(__file__), static_dir, f"all_features_label_{now}.xlsx"))
    config['data_path'] = data_path
    feature_label_weighted.to_excel(data_path, index=False)

    # 添加模型保存路径
    model_path = os.path.abspath(os.path.join(os.path.dirname(__file__), static_dir, f"model_{now}.pth"))
    config['model_path'] = model_path

    # 配置日志
    log_path = os.path.abspath(os.path.join(os.path.dirname(__file__), static_dir, f"train_{now}.log"))
    logging.basicConfig(filename=log_path, level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s', datefmt='%Y-%m-%d %H:%M:%S')

    print("config: ", config)
    logging.info("config: ", config)
    
    # 开始训练
    list_avg_f1 = []
    list_wrong_percentage = []
    list_precision = []
    list_recall = []
    list_f1 = []
    train_times = 1 if config['data_train']==r'all' else config["experiments_count"]
    for i in range(train_times):
        print(config)
        avg_f1, wrong_percentage, precision, recall, f1 = train_detect(config)
        list_avg_f1.append(avg_f1)
        list_wrong_percentage.append(wrong_percentage)
        list_precision.append(precision)
        list_recall.append(recall)
        list_f1.append(f1)

    logging.info(f"Result: Avg F1: {sum(list_avg_f1) / len(list_avg_f1):.4f} Avg Wrong Percentage: {sum(list_wrong_percentage) / len(list_wrong_percentage):.2f}%")
    logging.info(f"Result: Avg Precision: {[sum(p[i] for p in list_precision) / len(list_precision) for i in range(len(list_precision[0]))]} | {np.mean(list_precision)}")
    logging.info(f"Result: Avg Recall: {[sum(r[i] for r in list_recall) / len(list_recall) for i in range(len(list_recall[0]))]} | {np.mean(list_recall)}")
    logging.info(f"Result: Avg F1: {[sum(f1[i] for f1 in list_f1) / len(list_f1) for i in range(len(list_f1[0]))]} | {np.mean(list_f1)}")
    print(f"Result: Avg F1: {sum(list_avg_f1) / len(list_avg_f1):.4f} Avg Wrong Percentage: {sum(list_wrong_percentage) / len(list_wrong_percentage):.2f}%")
    print(f"Result: Avg Precision: {[sum(p[i] for p in list_precision) / len(list_precision) for i in range(len(list_precision[0]))]} | {np.mean(list_precision)}")
    print(f"Result: Avg Recall: {[sum(r[i] for r in list_recall) / len(list_recall) for i in range(len(list_recall[0]))]} | {np.mean(list_recall)}")
    print(f"Result: Avg F1: {[sum(f1[i] for f1 in list_f1) / len(list_f1) for i in range(len(list_f1[0]))]} | {np.mean(list_f1)}")

    end_time = time.time()  # 记录结束时间
    # 训练结束
    print("预测耗时:", end_time - start_time, "秒")  # 打印执行时间

    # 返回分类结果和模型文件下载 URL，static不是程序执行路径，而是app.mount的静态文件夹
    model_file_url = f"{request.base_url}train_api/model_{now}.pth"
    log_file_url = f"{request.base_url}train_api/train_{now}.log"
    data_file_url = f"{request.base_url}train_api/all_features_label_{now}.xlsx"

    # 返回分类结果和模型文件
    return {
        "classification_result": ClassificationResult(
            precision=precision,
            recall=recall,
            f1=f1,
            wrong_percentage=wrong_percentage
        ),
        "data_file": {
            "model_file_url": model_file_url,
            "log_file_url": log_file_url,
            "data_file_url": data_file_url
        }
    }

# 以下是fastapi启动配置
if __name__ == "__main__":
    name_app = os.path.basename(__file__)[0:-3]  # Get the name of the script
    log_config = {
        "version": 1,
        "disable_existing_loggers": True,
        "handlers": {
            "file_handler": {
                "class": "logging.FileHandler",
                "filename": "logfile.log",
            },
        },
        "root": {
            "handlers": ["file_handler"],
            "level": "INFO",
        },
    }
    # 创建静态文件存放文件夹
    static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "train_api"))  # 设置模型文件和配置文件的存放目录，和本py同级
    os.makedirs(static_dir, exist_ok=True)
    # train_api.py同级目录下的static文件夹
    app.mount("/train_api", StaticFiles(directory=static_dir), name="static")
    uvicorn.run(app, host="0.0.0.0", port=3397, reload=False)