import gradio as gr
from gradio_client import Client, handle_file
import os

# Initialize the Gradio client
client = Client("sitammeur/PicQ")

# Gradio interface function
def gradio_predict(image):
    # Save the image to a temporary path
    temp_path = os.path.join("temp", "gradio_upload.png")
    image.save(temp_path)

    try:
        # Use the Gradio client to predict the result
        result = client.predict(
            image=handle_file(temp_path),
            question="extract the item_weight from the image only give me the weight only like 100gm or 200 gm",
            api_name="/predict"
        )

        # Return the result as output
        return f"Prediction: {result}"
    
    finally:
        # Clean up the temporary file
        if os.path.exists(temp_path):
            os.remove(temp_path)

# Gradio interface setup
iface = gr.Interface(
    fn=gradio_predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    live=False
)

if __name__ == '__main__':
    if not os.path.exists("temp"):
        os.makedirs("temp")
    
    # Launch Gradio interface
    iface.launch(share=True)