Cynthiaaaaaaaa's picture
Update app.py
0b3461a verified
raw
history blame
No virus
1.43 kB
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()