from constants import RESOURCES
from data_preprocessing import RandomizeImageTransform
from utils import beam_search_decode

import streamlit as st
import PIL
import torch
import torchvision.transforms as T

MODEL_PATH = RESOURCES + "/model_2tcuvfsj.pt"

# TODO: make faster
transformer = torch.load(MODEL_PATH)
image_transform = T.Compose((
    T.ToTensor(),
    RandomizeImageTransform(width=transformer.hparams['image_width'],
                            height=transformer.hparams['image_height'],
                            random_magnitude=0)
))

st.markdown("### Image to TeX")
st.image("resources/frontend/latex_example_1.png")
file_png = st.file_uploader("Upload a PNG image", type=([".png"]))
if file_png is not None:
    image = PIL.Image.open(file_png)
    image = image.convert("RGB")
    tex = beam_search_decode(transformer, image, image_transform=image_transform)
    st.latex(tex[0])
    st.text(tex[0])