Nandoguer commited on
Commit
9c26d5d
·
verified ·
1 Parent(s): d3cb9da

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +139 -0
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()