easy-ml/commom/utils.py

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

import base64
import datetime
import inspect
import os
import sys
from contextlib import contextmanager
from json import JSONEncoder
from typing import Union

import numpy as np
import pandas as pd
import pytz
from PIL import Image

from config import BaseConfig
from .matplotlib_table import TableMaker


def f_is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False


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_image_crop_white_borders(image_path, output_path):
    # 打开图片
    image = Image.open(image_path)
    # 将图片转换为灰度图
    gray_image = image.convert('L')
    # 获取图片的宽度和高度
    width, height = gray_image.size
    # 初始化边界
    left, top, right, bottom = width, height, 0, 0

    # 遍历图片的每一行和每一列
    for y in range(height):
        for x in range(width):
            # 获取当前像素的灰度值
            pixel = gray_image.getpixel((x, y))
            # 如果像素不是白色（灰度值小于 255）
            if pixel < 255:
                # 更新边界
                if x < left:
                    left = x
                if x > right:
                    right = x
                if y < top:
                    top = y
                if y > bottom:
                    bottom = y

    # 裁剪图片
    cropped_image = image.crop((left, top, right + 1, bottom + 1))
    # 保存裁剪后的图片
    cropped_image.save(output_path)


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)


@contextmanager
def silent_print():
    original_stdout = sys.stdout

    class NullWriter:
        def write(self, text):
            pass

    null_writer = NullWriter()
    sys.stdout = null_writer
    try:
        yield
    finally:
        sys.stdout = original_stdout

@contextmanager
def df_print_nolimit():
    max_columns = pd.get_option('display.max_columns')
    max_rows = pd.get_option('display.max_rows')
    pd.set_option('display.max_columns', None)
    pd.set_option('display.max_rows', None)
    try:
        yield
    finally:
        pd.set_option('display.max_columns', max_columns)
        pd.set_option('display.max_rows', max_rows)