import gradio as gr
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from PIL import Image
import io
import base64
import tempfile


css = '''
.gradio-container{max-width: 950px !important}
h1{text-align:center}
'''

DESCRIPTIONz= """## Image to Parquet 📂

"""

def image_to_parquet(files):

    image_data = []

    for file_info in files:
        with open(file_info, "rb") as image_file:
            img = Image.open(image_file)
            buffered = io.BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        image_data.append({"name": file_info, "data": img_str})
    
    df = pd.DataFrame(image_data)
    
    table = pa.Table.from_pandas(df)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".parquet") as tmp_file:
        pq.write_table(table, tmp_file)
        parquet_file_path = tmp_file.name
    
    return parquet_file_path

with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown(DESCRIPTIONz)
    with gr.Row():
        image_input = gr.File(label="Upload Images", type="filepath", file_count="multiple", file_types=["image"])
        download_button = gr.File(label="Download Parquet File", interactive=False)
        
    convert_button = gr.Button("Convert Image to Parquet")
    
    convert_button.click(fn=image_to_parquet, inputs=[image_input], outputs=[download_button])

    gr.Markdown("📌 speed / time of converting images to a .parquet file depends on both the number of images uploaded and the quality and size of the uploaded images invloved for the conversion.")
demo.launch()