import streamlit as st
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import torch

# Carregar o modelo e tokenizador
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, 
                                  device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval().cuda()

# Título do app no Streamlit
st.title('OCR com Modelo Pré-treinado')

# Carregando a imagem
uploaded_file = st.file_uploader("Faça o upload de uma imagem", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Exibir a imagem
    image = Image.open(uploaded_file)
    st.image(image, caption='Imagem carregada', use_column_width=True)
    
    # Salvar a imagem para processamento (opcional)
    image.save("uploaded_image.jpg")

    # Executando OCR na imagem carregada
    with st.spinner('Processando a imagem...'):
        # Realizando o OCR
        res = model.chat(tokenizer, "uploaded_image.jpg", ocr_type='ocr')

    # Exibindo o resultado do OCR
    st.subheader("Texto Extraído:")
    st.text(res)
else:
    st.info("Por favor, faça o upload de uma imagem para continuar.")