Aryanshanu's picture
Update app.py
ffab8b7 verified
raw
history blame
1.44 kB
from transformers import pipeline
import gradio as gr
import matplotlib.pyplot as plt
import io
import base64
# Load the emotion classification model
emotion_pipeline = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
# Define the function for emotion analysis and visualization
def analyze_emotions(text):
# Get the emotion predictions
results = emotion_pipeline(text, top_k=None) # Get probabilities for all emotions
emotions = {res['label']: res['score'] for res in results}
# Generate a bar chart for the emotions
plt.figure(figsize=(8, 4))
plt.bar(emotions.keys(), emotions.values(), color='skyblue')
plt.xlabel("Emotions")
plt.ylabel("Probability")
plt.title("Emotion Analysis")
plt.tight_layout()
# Save the plot to a bytes buffer
buf = io.BytesIO()
plt.savefig(buf, format="png")
buf.seek(0)
img_str = base64.b64encode(buf.read()).decode("utf-8")
buf.close()
# Return image data as HTML for Gradio
img_html = f'<img src="data:image/png;base64,{img_str}" alt="Emotion Analysis"/>'
return img_html
# Create Gradio interface
interface = gr.Interface(
fn=analyze_emotions,
inputs="text",
outputs="html",
title="Emotion Detection with Visualization",
description="Enter text to detect and visualize emotions."
)
# Launch the app
if __name__ == "__main__":
interface.launch()