easy-ml/webui/utils.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/12/5
@desc: 
"""
import os
import shutil
from typing import List

import gradio as gr
import pandas as pd

from config import BaseConfig
from data import DataLoaderExcel, DataExplore
from .manager import engine

DATA_DIR = "data"
UPLOAD_DATA_PREFIX = "prefix_upload_data_"
data_loader = DataLoaderExcel()


def _check_save_dir(data):
    project_name = engine.get(data, "project_name")
    if project_name is None or len(project_name) == 0:
        gr.Warning(message='项目名称不能为空', duration=5)
        return False
    return True


def _get_prefix_file(save_path, prefix):
    file_name_list: List[str] = os.listdir(save_path)
    for file_name in file_name_list:
        if prefix in file_name:
            return os.path.join(save_path, file_name)


def _get_base_dir(data):
    project_name = engine.get(data, "project_name")
    base_dir = os.path.join(BaseConfig.train_path, project_name)
    return base_dir


def _get_upload_data(data) -> pd.DataFrame:
    base_dir = _get_base_dir(data)
    save_path = os.path.join(base_dir, DATA_DIR)
    file_path = _get_prefix_file(save_path, UPLOAD_DATA_PREFIX)
    df = data_loader.get_data(file_path)
    return df


def f_project_is_exist(data):
    project_name = engine.get(data, "project_name")
    if project_name is None or len(project_name) == 0:
        gr.Warning(message='项目名称不能为空', duration=5)
    elif os.path.exists(_get_base_dir(data)):
        gr.Warning(message='项目名称已被使用', duration=5)


def f_get_save_path(data, file_name: str, sub_dir="", name_prefix=""):
    base_dir = _get_base_dir(data)
    save_path = os.path.join(base_dir, sub_dir)
    os.makedirs(save_path, exist_ok=True)
    # 有前缀标示的先删除
    if name_prefix:
        file = _get_prefix_file(save_path, name_prefix)
        if file:
            os.remove(file)
    save_path = os.path.join(save_path, name_prefix + os.path.basename(file_name))
    return save_path


def f_data_upload(data):
    if not _check_save_dir(data):
        return
    file_data = engine.get(data, "file_data")
    data_path = f_get_save_path(data, file_data.name, DATA_DIR, UPLOAD_DATA_PREFIX)
    shutil.copy(file_data.name, data_path)
    df = _get_upload_data(data)
    distribution = DataExplore.distribution(df)

    return gr.update(value=df, visible=True), gr.update(value=distribution, visible=True),