import streamlit as st
from transformers import pipeline

# Summarization
def summarization(text):
    image_to_text_model = pipeline("text-generation", model="ainize/bart-base-cnn")
    summary = image_to_text_model(text, max_length=100, do_sample=False)[0]["generated_text"]
    return summary

# Sentiment Classification
def sentiment_classification(summary):
    sentiment_model = pipeline("text-classification", model="wxrrrrrrr/finetuned_sentiment_analysis")
    result = sentiment_model(summary, max_length=100, do_sample=False)[0]['label']
    return result

def main():
    st.set_page_config(page_title="Your Image to Text Analysis", page_icon="🦜")
    st.header("Tell me your comments!")
    uploaded_file = st.file_uploader("Select an Image...")

    if uploaded_file is not None:
        with open(uploaded_file.name, "wb") as file:
            file.write(uploaded_file.getbuffer())
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

        # Stage 1: Summarization
        st.text('Processing image to text...')
        summary = summarization(uploaded_file.name)
        st.write(summary)

        # Stage 2: Sentiment Classification
        st.text('Analyzing sentiment...')
        sentiment = sentiment_classification(summary)
        st.write(sentiment)

        # Display the classification result
        st.write("Sentiment:", sentiment)

if __name__ == '__main__':
    main()