Aryanshanu's picture
Update app.py
7713707 verified
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")
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}
# Extract emotion labels and probabilities
emotion_labels = list(emotions.keys())
emotion_probabilities = list(emotions.values())
# Generate a pie chart for the emotions
plt.figure(figsize=(8, 6))
plt.pie(emotion_probabilities, labels=emotion_labels, autopct='%1.1f%%', startangle=140)
plt.title("Emotion Distribution")
plt.axis('equal')
# 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()