Spaces:
Runtime error
Runtime error
Upload app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from transformers import pipeline, RobertaTokenizer, RobertaForSequenceClassification
|
| 2 |
+
from sklearn.feature_extraction.text import CountVectorizer
|
| 3 |
+
from sklearn.naive_bayes import MultinomialNB
|
| 4 |
+
from sklearn.pipeline import make_pipeline
|
| 5 |
+
from datasets import load_dataset
|
| 6 |
+
from PIL import Image, ImageDraw, ImageFont
|
| 7 |
+
import textwrap
|
| 8 |
+
import random
|
| 9 |
+
from diffusers import StableDiffusionPipeline
|
| 10 |
+
import torch
|
| 11 |
+
from sklearn.metrics import classification_report, accuracy_score, f1_score
|
| 12 |
+
from sklearn.model_selection import train_test_split # Import train_test_split
|
| 13 |
+
import gradio as gr
|
| 14 |
+
|
| 15 |
+
# Load the datasets for emotion detection and quotes
|
| 16 |
+
emotion_dataset = load_dataset("dair-ai/emotion")
|
| 17 |
+
quotes_dataset = load_dataset("Abirate/english_quotes")
|
| 18 |
+
|
| 19 |
+
# Prepare the Bag-of-Words Model
|
| 20 |
+
vectorizer = CountVectorizer()
|
| 21 |
+
naive_bayes_classifier = MultinomialNB()
|
| 22 |
+
bow_pipeline = make_pipeline(vectorizer, naive_bayes_classifier)
|
| 23 |
+
texts = [example['text'] for example in emotion_dataset['train']]
|
| 24 |
+
labels = [example['label'] for example in emotion_dataset['train']]
|
| 25 |
+
train_texts, test_texts, train_labels, test_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)
|
| 26 |
+
bow_pipeline.fit(train_texts, train_labels) # Train the Bag-of-Words model
|
| 27 |
+
predicted_labels = bow_pipeline.predict(test_texts)
|
| 28 |
+
print("Bag-of-Words Model Evaluation Metrics:")
|
| 29 |
+
print(classification_report(test_labels, predicted_labels))
|
| 30 |
+
print("Accuracy:", accuracy_score(test_labels, predicted_labels))
|
| 31 |
+
print("F1 Score:", f1_score(test_labels, predicted_labels, average='weighted'))
|
| 32 |
+
|
| 33 |
+
# Load the emotion classification models
|
| 34 |
+
distilbert_classifier = pipeline('text-classification', model='bhadresh-savani/distilbert-base-uncased-emotion')
|
| 35 |
+
roberta_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
|
| 36 |
+
roberta_model = RobertaForSequenceClassification.from_pretrained('roberta-base')
|
| 37 |
+
|
| 38 |
+
# Assuming there's a test split for evaluation
|
| 39 |
+
test_data = emotion_dataset['test']
|
| 40 |
+
test_texts = [example['text'] for example in test_data]
|
| 41 |
+
test_labels = [example['label'] for example in test_data]
|
| 42 |
+
label_mapping = {0: 'sadness', 1: 'joy', 2: 'love', 3: 'anger', 4: 'fear', 5: 'surprise'}
|
| 43 |
+
test_labels = [label_mapping[label].lower() for label in test_labels]
|
| 44 |
+
'''
|
| 45 |
+
# Evaluate DistilBERT
|
| 46 |
+
distilbert_predictions = [distilbert_classifier(text)[0]['label'].lower() for text in test_texts]
|
| 47 |
+
print("DistilBERT Model Evaluation Metrics:")
|
| 48 |
+
print("Accuracy:", accuracy_score(test_labels, distilbert_predictions))
|
| 49 |
+
print("F1 Score:", f1_score(test_labels, distilbert_predictions, average='weighted'))
|
| 50 |
+
'''
|
| 51 |
+
|
| 52 |
+
# Function to generate an image from a prompt using Stable Diffusion
|
| 53 |
+
def generate_image(prompt):
|
| 54 |
+
try:
|
| 55 |
+
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", variant="fp16", torch_dtype=torch.float16).to("cuda")
|
| 56 |
+
output = pipe(prompt=prompt)
|
| 57 |
+
image = output.images[0]
|
| 58 |
+
image.save("output_image.png")
|
| 59 |
+
return "output_image.png"
|
| 60 |
+
except Exception as e:
|
| 61 |
+
print(f"Failed during the image generation process: {e}")
|
| 62 |
+
return None
|
| 63 |
+
|
| 64 |
+
# Function to create an image with the quote
|
| 65 |
+
def create_image_with_quote(quote, author, emotion):
|
| 66 |
+
prompt = f"{emotion} themed image"
|
| 67 |
+
img_path = generate_image(prompt)
|
| 68 |
+
if not img_path:
|
| 69 |
+
print("Image generation failed.")
|
| 70 |
+
return None
|
| 71 |
+
img = Image.open(img_path)
|
| 72 |
+
d = ImageDraw.Draw(img)
|
| 73 |
+
quote_font = ImageFont.load_default()
|
| 74 |
+
author_font = ImageFont.load_default()
|
| 75 |
+
quote_wrapped = textwrap.fill(quote, width=40)
|
| 76 |
+
author_wrapped = textwrap.fill(f"- {author}", width=40)
|
| 77 |
+
draw_text_with_background(d, quote_wrapped, author_wrapped, quote_font, author_font,img.width, img.height)
|
| 78 |
+
return img
|
| 79 |
+
|
| 80 |
+
def draw_text_with_background(d, quote, author, quote_font, author_font, img_width, img_height):
|
| 81 |
+
# Calculate text bounding boxes
|
| 82 |
+
quote_bbox = d.textbbox((0, 0), quote, font=quote_font)
|
| 83 |
+
author_bbox = d.textbbox((0, 0), author, font=author_font)
|
| 84 |
+
|
| 85 |
+
# Calculate the position for the quote
|
| 86 |
+
quote_x = (img_width - (quote_bbox[2] - quote_bbox[0])) / 2
|
| 87 |
+
quote_y = (img_height - (quote_bbox[3] - quote_bbox[1])) / 2 - 20 # Slightly above the center
|
| 88 |
+
|
| 89 |
+
# Calculate the position for the author
|
| 90 |
+
author_x = (img_width - (author_bbox[2] - author_bbox[0])) / 2
|
| 91 |
+
author_y = quote_y + (quote_bbox[3] - quote_bbox[1]) + 10 # Just below the quote
|
| 92 |
+
|
| 93 |
+
# Draw background for quote
|
| 94 |
+
d.rectangle([quote_x - 10, quote_y - 5, quote_x + (quote_bbox[2] - quote_bbox[0]) + 10, quote_y + (quote_bbox[3] - quote_bbox[1]) + 5], fill=(255, 255, 255, 128))
|
| 95 |
+
# Draw background for author
|
| 96 |
+
d.rectangle([author_x - 10, author_y - 5, author_x + (author_bbox[2] - author_bbox[0]) + 10, author_y + (author_bbox[3] - author_bbox[1]) + 5], fill=(255, 255, 255, 128))
|
| 97 |
+
|
| 98 |
+
# Draw text over the boxes
|
| 99 |
+
d.text((quote_x, quote_y), quote, font=quote_font, fill="black")
|
| 100 |
+
d.text((author_x, author_y), author, font=author_font, fill="black")
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def predict_emotion(text, model_choice):
|
| 106 |
+
if model_choice == 'distilbert':
|
| 107 |
+
result = distilbert_classifier(text)
|
| 108 |
+
return result[0]['label'].lower()
|
| 109 |
+
elif model_choice == 'bow':
|
| 110 |
+
prediction = bow_pipeline.predict([text])[0]
|
| 111 |
+
return 'positive' if prediction == 1 else 'negative'
|
| 112 |
+
elif model_choice == 'roberta':
|
| 113 |
+
inputs = roberta_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
|
| 114 |
+
outputs = roberta_model(**inputs)
|
| 115 |
+
prediction = torch.argmax(outputs.logits, dim=-1)
|
| 116 |
+
return 'positive' if prediction.item() == 1 else 'negative'
|
| 117 |
+
|
| 118 |
+
def evaluate_quotes_for_emotion(emotion):
|
| 119 |
+
# Filter quotes based on the predicted emotion
|
| 120 |
+
suitable_quotes = [q for q in quotes_dataset['train'] if emotion.lower() in q['tags']]
|
| 121 |
+
if not suitable_quotes:
|
| 122 |
+
suitable_quotes = quotes_dataset['train'] # fallback to any quote if no tags match
|
| 123 |
+
selected_quote = random.choice(suitable_quotes)
|
| 124 |
+
return selected_quote['quote'], selected_quote.get('author', 'Unknown')
|
| 125 |
+
|
| 126 |
+
def predict_emotion_and_generate_quote(feelings, model_choice):
|
| 127 |
+
emotion = predict_emotion(feelings, model_choice)
|
| 128 |
+
quote, author = evaluate_quotes_for_emotion(emotion)
|
| 129 |
+
image = create_image_with_quote(quote, author, emotion)
|
| 130 |
+
return quote, author, image
|
| 131 |
+
|
| 132 |
+
iface = gr.Interface(fn=predict_emotion_and_generate_quote,
|
| 133 |
+
inputs=["text", "text"],
|
| 134 |
+
outputs=["text", "text", "image"],
|
| 135 |
+
title="Quote Generator: Feeling's Inspired",
|
| 136 |
+
description="Enter your feelings and choose a model to receive an inspiring quote with an accompanying image. Model Choices include distilbert,roberta,bow",
|
| 137 |
+
allow_flagging=False,
|
| 138 |
+
theme="default")
|
| 139 |
+
iface.launch()
|