import streamlit as st
import easyocr
import asyncio
import re
from transformers import pipeline

asyncio.set_event_loop(asyncio.new_event_loop())

# Load OCR model
reader = easyocr.Reader(['lt'])

# Load summarization model
summarizer_model_name = "LukasStankevicius/t5-base-lithuanian-news-summaries-175"
try:
    summarizer = pipeline("summarization", model=summarizer_model_name)
except Exception as e:
    st.error(f"Klaida su santraukos modeliu: {e}")
    st.stop()

# Streamlit UI Setup
st.title("Lietuviško teksto iš nuotraukos santraukos sukūrimas naudojant DI (naudojant lietuvišką modelį)")
st.write("Įkelkite nuotrauką su tekstu:")

# Upload image file
uploaded_file = st.file_uploader("Įkelkite nuotrauką...", type=["png", "jpg", "jpeg"])

def preprocess_text(text):
    text = text.replace("-\n", "").replace("- ", "")  
    text = re.sub(r"[^a-zA-ZąčęėįšųūžĄČĘĖĮŠŲŪŽ0-9\s\.,;:]", "", text)  
    return text

if uploaded_file:
    st.image(uploaded_file, caption="Įkelta nuotrauka", use_container_width=True)

    with st.spinner("Gaunamas tekstas..."):
        extracted_text = reader.readtext(uploaded_file.read(), detail=0)
        extracted_text = " ".join(extracted_text)

    if extracted_text:
        st.subheader("Gautas tekstas:")
        st.write(extracted_text)

        # Preprocess the extracted text
        processed_text = preprocess_text(extracted_text)

        st.subheader("Sutvarkytas tekstas:")
        st.write(processed_text)

        # Generate Summary
        with st.spinner("Gaunama santrauka..."):
            try:
                summary_output = summarizer(processed_text, max_length=100, min_length=30, do_sample=False)
                summary_lithuanian = summary_output[0]['summary_text']

                st.subheader("Santrauka Lietuviškai:")
                st.write(summary_lithuanian)

            except Exception as e:
                st.error(f"Klaida gaunant santrauką: {e}")
    else:
        st.warning("Nerasta jokio teksto, pabandykite iš naujo.")