from transformers import pipeline, RobertaTokenizer, RobertaForSequenceClassification from sklearn.feature_extraction.text import CountVectorizer from sklearn.naive_bayes import MultinomialNB from sklearn.pipeline import make_pipeline from datasets import load_dataset from PIL import Image, ImageDraw, ImageFont import textwrap import random from diffusers import StableDiffusionPipeline import torch from sklearn.metrics import classification_report, accuracy_score, f1_score from sklearn.model_selection import train_test_split # Import train_test_split import gradio as gr # Load the datasets for emotion detection and quotes emotion_dataset = load_dataset("dair-ai/emotion",trust_remote_code=True) quotes_dataset = load_dataset("Abirate/english_quotes",trust_remote_code=True) # Prepare the Bag-of-Words Model vectorizer = CountVectorizer() naive_bayes_classifier = MultinomialNB() bow_pipeline = make_pipeline(vectorizer, naive_bayes_classifier) texts = [example['text'] for example in emotion_dataset['train']] labels = [example['label'] for example in emotion_dataset['train']] train_texts, test_texts, train_labels, test_labels = train_test_split(texts, labels, test_size=0.2, random_state=42) bow_pipeline.fit(train_texts, train_labels) # Train the Bag-of-Words model predicted_labels = bow_pipeline.predict(test_texts) print("Bag-of-Words Model Evaluation Metrics:") print(classification_report(test_labels, predicted_labels)) print("Accuracy:", accuracy_score(test_labels, predicted_labels)) print("F1 Score:", f1_score(test_labels, predicted_labels, average='weighted')) # Load the emotion classification models distilbert_classifier = pipeline('text-classification', model='bhadresh-savani/distilbert-base-uncased-emotion') roberta_tokenizer = RobertaTokenizer.from_pretrained('roberta-base') roberta_model = RobertaForSequenceClassification.from_pretrained('roberta-base') # Assuming there's a test split for evaluation test_data = emotion_dataset['test'] test_texts = [example['text'] for example in test_data] test_labels = [example['label'] for example in test_data] label_mapping = {0: 'sadness', 1: 'joy', 2: 'love', 3: 'anger', 4: 'fear', 5: 'surprise'} test_labels = [label_mapping[label].lower() for label in test_labels] ''' # Evaluate DistilBERT distilbert_predictions = [distilbert_classifier(text)[0]['label'].lower() for text in test_texts] print("DistilBERT Model Evaluation Metrics:") print("Accuracy:", accuracy_score(test_labels, distilbert_predictions)) print("F1 Score:", f1_score(test_labels, distilbert_predictions, average='weighted')) # Evaluate RoBERTa roberta_predictions = [] for text in test_texts: inputs = roberta_tokenizer(text, return_tensors="pt", padding=True, truncation=True) outputs = roberta_model(**inputs) prediction = torch.argmax(outputs.logits, dim=-1).item() roberta_predictions.append(label_mapping[prediction].lower()) print("RoBERTa Model Evaluation Metrics:") print("Accuracy:", accuracy_score(test_labels, roberta_predictions)) print("F1 Score:", f1_score(test_labels, roberta_predictions, average='weighted')) ''' # Function to generate an image from a prompt using Stable Diffusion def generate_image(prompt): try: pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", variant="fp16", torch_dtype=torch.float32)#.to("cuda") output = pipe(prompt=prompt) image = output.images[0] image.save("output_image.png") return "output_image.png" except Exception as e: print(f"Failed during the image generation process: {e}") return None # Function to create an image with the quote def create_image_with_quote(quote, author, emotion): prompt = f"{emotion} themed image" img_path = generate_image(prompt) if not img_path: print("Image generation failed.") return None img = Image.open(img_path) d = ImageDraw.Draw(img) quote_font = ImageFont.load_default() author_font = ImageFont.load_default() quote_wrapped = textwrap.fill(quote, width=40) author_wrapped = textwrap.fill(f"- {author}", width=40) draw_text_with_background(d, quote_wrapped, author_wrapped, quote_font, author_font,img.width, img.height) return img def draw_text_with_background(d, quote, author, quote_font, author_font, img_width, img_height): # Calculate text bounding boxes quote_bbox = d.textbbox((0, 0), quote, font=quote_font) author_bbox = d.textbbox((0, 0), author, font=author_font) # Calculate the position for the quote quote_x = (img_width - (quote_bbox[2] - quote_bbox[0])) / 2 quote_y = (img_height - (quote_bbox[3] - quote_bbox[1])) / 2 - 20 # Slightly above the center # Calculate the position for the author author_x = (img_width - (author_bbox[2] - author_bbox[0])) / 2 author_y = quote_y + (quote_bbox[3] - quote_bbox[1]) + 10 # Just below the quote # Draw background for quote 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)) # Draw background for author 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)) # Draw text over the boxes d.text((quote_x, quote_y), quote, font=quote_font, fill="black") d.text((author_x, author_y), author, font=author_font, fill="black") def predict_emotion(text, model_choice): if model_choice == 'distilbert': result = distilbert_classifier(text) return result[0]['label'].lower() elif model_choice == 'bow': prediction = bow_pipeline.predict([text])[0] return 'positive' if prediction == 1 else 'negative' elif model_choice == 'roberta': inputs = roberta_tokenizer(text, return_tensors="pt", padding=True, truncation=True) outputs = roberta_model(**inputs) prediction = torch.argmax(outputs.logits, dim=-1) return 'positive' if prediction.item() == 1 else 'negative' def evaluate_quotes_for_emotion(emotion): # Filter quotes based on the predicted emotion suitable_quotes = [q for q in quotes_dataset['train'] if emotion.lower() in q['tags']] if not suitable_quotes: suitable_quotes = quotes_dataset['train'] # fallback to any quote if no tags match selected_quote = random.choice(suitable_quotes) return selected_quote['quote'], selected_quote.get('author', 'Unknown') def predict_emotion_and_generate_quote(feelings, model_choice): emotion = predict_emotion(feelings, model_choice) quote, author = evaluate_quotes_for_emotion(emotion) image = create_image_with_quote(quote, author, emotion) return quote, author, image iface = gr.Interface(fn=predict_emotion_and_generate_quote, inputs=["text", "text"], outputs=["text", "text", "image"], title="Quote Generator: Feeling's Inspired", description="Enter your feelings and choose a model to receive an inspiring quote with an accompanying image. Model Choices include distilbert,roberta,bow", allow_flagging=False, theme="default") iface.launch()