app.py

import base64
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
from io import BytesIO

df = pd.read_csv("herbologist_almanac_checklist_data.csv")
columns_with_tasks = [
    "task1",
    "task2",
    "task3",
    "task4",
    "color1",
    "color2",
    "color3",
    "color4",
    "color5",
    "color6",
]
PLANTS = list(df["plant"].unique())
# 更改字体设置
# rc("font", family="Arial Unicode MS")

# def decimal2base64(decimal_int):
#     # 定义 64 进制字符集
#     base64_chars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

#     if decimal_int == 0:
#         return "0"

#     base64_str = ""
#     while decimal_int > 0:
#         remainder = decimal_int % 64
#         base64_str = base64_chars[remainder] + base64_str
#         decimal_int //= 64

#     return base64_str


def bin_ls2base64(ls):
    # 将二进制列表转换为二进制列表
    binary_str = "".join(str(n) for n in ls)

    # decimal_int = int(bin_str, 2)
    byte_str = int(binary_str, 2).to_bytes((len(binary_str) + 7) // 8, byteorder="big")
    base64_str = base64.b64encode(byte_str).decode("utf-8")
    return base64_str


def base64_to_binary(base64_str):
    if isinstance(base64_str, str):
        # 将base64字符串转换为二进制列表
        byte_str = base64.b64decode(base64_str)
        binary_str = bin(int.from_bytes(byte_str, byteorder="big"))[2:].zfill(
            len(columns_with_tasks) * len(PLANTS)
        )
        # ls = [int(n) for n in binary_str]
        return binary_str


def parse_token(token):
    try:
        if not token:  # 处理空字符串的情况
            token = "\x00"

        if len(token) > 0:
            payload = base64_to_binary(token)
            almanac_data: list = [int(n) for n in payload]
            # print(len(almanac_data))
            parsed_dict = {}
            for _, row in df.iterrows():
                parsed_dict[row["plant"]] = [
                    almanac_data.pop() for _ in range(len(columns_with_tasks))
                ]

        return parsed_dict
    except Exception as e:
        print(e)


# 定义一个简单的函数，模拟接收DataFrame数据
def process_data(*args):
    almanac_dict = dict(zip(PLANTS, args))

    almanac_df = df.filter(items=["plant", "name"] + columns_with_tasks)
    almanac_bin_ls = []
    for pl in PLANTS:
        almanac_inp = almanac_dict[pl]

        almanac_bin = [0 for _ in range(len(columns_with_tasks))]

        for n in almanac_inp:
            almanac_bin[n] = 1

        almanac_bin_ls += almanac_bin

        for _, i in enumerate(almanac_inp):
            almanac_df.loc[almanac_df["plant"] == pl, columns_with_tasks[i]] = "✔"

    almanac_reverse_64 = bin_ls2base64(reversed(almanac_bin_ls))
    # 这里可以添加处理DataFrame的逻辑
    return (
        generate_img(almanac_df.drop(columns=df.columns[0])),
        almanac_reverse_64,
    )


def show_checkbox_groups(token):
    checklist_inputs = []
    parsed_dict = parse_token(token)
    for index, row in df.iterrows():
        tasks = [
            row[col][0].upper() + row[col][1:]
            for col in columns_with_tasks
            if df.notnull().at[index, col]
        ]
        with gr.Row():
            checkbox = gr.CheckboxGroup(
                tasks,
                label=f"{row['name']}",
                value=[
                    tasks[i]
                    for i, v in enumerate(parsed_dict[row["plant"]])
                    if v == 1 and df.notnull().at[index, columns_with_tasks[i]]
                ],
                type="index",
            )
            checklist_inputs.append(checkbox)

    return checklist_inputs


def wrap_text(text, max_width=20):
    """Manually wrap text based on max_width (character count)"""
    wrapped_lines = []
    words = text.split(" ")
    line = ""

    for word in words:
        if len(line) + len(word) + 1 <= max_width:
            line += word + " "
        else:
            wrapped_lines.append(line.strip())
            line = word + " "

    wrapped_lines.append(line.strip())
    return "\n".join(wrapped_lines)


def generate_img(df):
    fig, ax = plt.subplots(
        figsize=(12, 16)
    )  # Adjust the figure size for better readability
    ax.axis("off")  # Turn off the axis

    # Create a table with a light grey background for the header
    table = ax.table(
        cellText=df.values,
        colLabels=df.columns,
        loc="center",
        colColours=["#f2f2f2"] * len(df.columns),
        colWidths=[w / len(df.columns) for w in [1] + [2] * 4 + [1] * 6],
    )

    table.auto_set_font_size(False)
    table.set_fontsize(10)

    # Apply text wrapping and center alignment to non-header cells
    for (row, col), cell in table.get_celld().items():

        if cell.get_text().get_text() == "nan":
            cell.set_text_props(text="", ha="center")
        if row == 0:
            cell.set_text_props(
                weight="bold", ha="center"
            )  # Bold and center-align header text
        if col == 0:
            text = cell.get_text().get_text()
            wrapped_text = wrap_text(text, 9)
            cell.set_text_props(text=wrapped_text, ha="center", linespacing=1)
        else:
            text = cell.get_text().get_text()
            wrapped_text = wrap_text(text)
            cell.set_text_props(
                text=wrapped_text,
                ha="center",
                fontsize=10,
                linespacing=1,
            )  # Enable text wrapping and center-align

    # Manually scale table if needed for better readability with wrapped text
    table.scale(1, 4)  # Adjust row height

    # Ensure the layout is adjusted properly
    plt.tight_layout(pad=0.5)  # Increase padding slightly

    # Save the figure to a bytes buffer
    buffer = BytesIO()
    fig.savefig(
        buffer, format="png", pad_inches=0.2
    )  # Adjust padding around the figure
    plt.close(fig)  # Close the figure to free up memory
    buffer.seek(0)
    image = Image.open(buffer)

    return image


with gr.Blocks() as app:
    gr.Markdown(
        """
        <center><font size=8>👩🏻‍🌾Herbology Almanac Checklist Generator📝</font></center>
        This is a simple web app that generates an almanac checklist for your plants.
        """
    )
    recovery_token = gr.Textbox(
        value="",
        label="Recovery Token",
        info="Save this token or paste your saved one here",
        placeholder="Keep this token to restore your previous input".upper(),
        interactive=True,
    )

    checklist_inputs = show_checkbox_groups(recovery_token.value)

    submit_button = gr.Button("Generate Image and Token")

    # df_out = gr.Dataframe(label="Output Dataframe", interactive=False)

    generated_img = gr.Image(label="Generated Image", format="png", type="pil")

    logs = gr.Markdown(
        """
        # CHANGELOG
        - 204/08/31: Initial release
        - 204/09/03: Add recovery token
    """
    )

    submit_button.click(
        process_data,
        inputs=checklist_inputs,
        outputs=[generated_img, recovery_token],
    )

    recovery_token.change(
        show_checkbox_groups,
        inputs=[recovery_token],
        outputs=checklist_inputs,
    )

    # generate_button.click(
    #     generate_img,
    #     inputs=[df_out],
    #     outputs=[generated_img],
    # )


app.queue()
app.launch()