from dotenv import find_dotenv, load_dotenv
import requests
import os
import streamlit as st

# Set up
load_dotenv(find_dotenv())
HUGGING_FACE_TOKEN = os.getenv('HUGGING_FACE_TOKEN')

API_URL = "https://router.huggingface.co/hf-inference/models/microsoft/trocr-base-handwritten"
headers = {"Authorization": f'Bearer {HUGGING_FACE_TOKEN}'}

# Query function
def query(data):
    response = requests.post(API_URL, headers=headers, data=data)

    # Print JSON response for debugging
    print(response.json())

    return response.json()[0]['generated_text']


# UI
def main():
    st.set_page_config(page_icon='random')
    st.title('READ-me : Weird Handwriting Translator')
    st.write('Upload an image of a 1 line of handwriting!')

    uploaded_file = st.file_uploader("Choose a file")

    if uploaded_file is not None:
        bytes_data = uploaded_file.getvalue()

        # Show the image
        st.image(bytes_data)

        st.header('This is what that says:')

        # Query the model
        st.write(query(bytes_data))

if __name__ == '__main__':
    main()