easy-ml/commom/utils.py

# -*- coding:utf-8 -*-
"""
@author: yq
@time: 2023/12/28
@desc:  各种工具类
"""

import base64
import datetime
import inspect
import os
from json import JSONEncoder
from typing import Union

import numpy as np
import pandas as pd
import pytz

from config import BaseConfig
from .matplotlib_table import TableMaker


def f_format_float(num: float, n=3):
    return f"{num: .{n}f}"


def f_get_date(offset: int = 0, connect: str = "-") -> str:
    current_date = datetime.datetime.now(pytz.timezone("Asia/Shanghai")).date() + datetime.timedelta(days=offset)
    return current_date.strftime(f"%Y{connect}%m{connect}%d")


def f_get_datetime(offset: int = 0, connect: str = "_") -> str:
    current_date = datetime.datetime.now(pytz.timezone("Asia/Shanghai")) + datetime.timedelta(days=offset)
    return current_date.strftime(f"%Y{connect}%m{connect}%d{connect}%H{connect}%M{connect}%S")


def f_get_clazz_in_module(module):
    """
    获取包下的所有类
    """
    classes = []
    for name, member in inspect.getmembers(module):
        if inspect.isclass(member):
            classes.append(member)
    return classes


def f_save_train_df(file_name: str, df: pd.DataFrame):
    file_path = os.path.join(BaseConfig.train_path, file_name)
    df.to_excel(f"{file_path}.xlsx", index=False)


def f_df_to_image(df: pd.DataFrame, filename, fontsize=12):
    converter = TableMaker(fontsize=fontsize, encode_base64=False, for_document=False)
    converter.run(df, filename)

    # if importlib.util.find_spec("dataframe_image"):
    #     import dataframe_image as dfi
    #
    #     dfi.export(obj=df, filename=filename, fontsize=fontsize, table_conversion='matplotlib')
    # elif importlib.util.find_spec("plotly"):
    #     import plotly.graph_objects as go
    #     import plotly.figure_factory as ff
    #     import plotly.io as pio
    #
    #     fig = ff.create_table(df)
    #     fig.update_layout()
    #     fig.write_image(filename)
    #
    #     fig = go.Figure(data=go.Table(
    #         header=dict(
    #             values=df.columns.to_list(),
    #             font=dict(color='black', size=fontsize),
    #             fill_color="white",
    #             line_color='black',
    #             align="center"
    #         ),
    #         cells=dict(
    #             values=[df[k].tolist() for k in df.columns],
    #             font=dict(color='black', size=fontsize),
    #             fill_color="white",
    #             line_color='black',
    #             align="center")
    #     )).update_layout()
    #     pio.write_image(fig, filename)
    # else:
    #     raise GeneralException(ResultCodesEnum.NOT_FOUND, message=f"缺少画图依赖【dataframe_image】或者【plotly】")


def _f_image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        img_str = base64.b64encode(image_file.read())
        return img_str.decode("utf-8")


def f_display_images_by_side(display, image_path_list, title: str = "", width: int = 500,
                             image_path_list2: Union[list, None] = None, title2: str = "", ):
    if isinstance(image_path_list, str):
        image_path_list = [image_path_list]
    # justify-content:space-around; 会导致某些情况下图片越界
    html_str = '<div style="display:flex;">'
    if title != "":
        html_str += '<div>{}</div>'.format(title)
    for image_path in image_path_list:
        html_str += f'<img src="data:image/png;base64,{_f_image_to_base64(image_path)}" style="width:{width}px;"/>'
    html_str += '</div>'
    if not (image_path_list2 is None or len(image_path_list2) == 0):
        html_str += '<div style="display:flex;">'
        if title2 != "":
            html_str += '<div>{}</div>'.format(title2)
        for image_path in image_path_list2:
            html_str += f'<img src="data:image/png;base64,{_f_image_to_base64(image_path)}" style="width:{width}px;"/>'
        html_str += '</div>'
    display.display(display.HTML(html_str))


def f_display_title(display, title):
    html_str = f"<h2>{title}</h2>"
    display.display(display.HTML(html_str))


class f_clazz_to_json(JSONEncoder):
    def default(self, o):
        return o.__dict__


class NumpyEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NumpyEncoder, self).default(obj)