import gradio as gr
import subprocess
from PIL import Image
import tempfile
import os

def predict(input_image: Image.Image):
    # Save the input image to a temporary file
    temp_dir = tempfile.mkdtemp()
    temp_image_path = os.path.join(temp_dir, 'temp_image.jpg')
    input_image.save(temp_image_path)

    # Update img_list to include the path of the temporary image
    with open('test_img_list.txt', 'w') as f:
        f.write(temp_image_path)

    # Run the pylaia-htr-decode-ctc command
    subprocess.run(['pylaia-htr-decode-ctc', '--config', 'my_decode_config.yaml'])

    # Read the output from predict.txt
    with open('predict.txt', 'r') as f:
        output = f.read().strip().split('\n')[-1]  # Assuming the last line of predict.txt contains the desired output

    # Cleanup temporary files
    os.remove(temp_image_path)
    os.rmdir(temp_dir)

    return output

# Gradio interface
title = "PyLaia HTR"
description = "Inference using PyLaia models."
examples =[["examples/example01.jpg"], ["examples/example02.jpg"]]

iface = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Textbox(),
    title=title,
    description=description,
    examples=examples
)

iface.launch()