commit

app.py

@@ -1,18 +1,12 @@
 import gradio as gr
 from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForSequenceClassification
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import pipeline
 
 # Load the image captioning model and tokenizer
 caption_model_name = "Salesforce/blip-image-captioning-large"
 caption_processor = BlipProcessor.from_pretrained(caption_model_name)
 caption_model = BlipForConditionalGeneration.from_pretrained(caption_model_name)
 
-# Load the emotion analysis model and tokenizer
-emotion_model_name = "SamLowe/roberta-base-go_emotions"
-emotion_tokenizer = AutoTokenizer.from_pretrained(emotion_model_name)
-emotion_model = AutoModelForSequenceClassification.from_pretrained(emotion_model_name)
-
 def generate_caption_and_analyze_emotions(image):
     # Preprocess the image for caption generation
     caption_inputs = caption_processor(images=image, return_tensors="pt")
@@ -23,23 +17,16 @@ def generate_caption_and_analyze_emotions(image):
     # Decode the output caption
     decoded_caption = caption_processor.decode(caption[0], skip_special_tokens=True)
 
-    # Analyze emotions of the generated caption
-    # Preprocess the caption for emotion analysis
-    emotion_inputs = emotion_tokenizer.encode_plus(
-        decoded_caption,
-        max_length=128,
-        padding="max_length",
-        truncation=True,
-        return_tensors="pt"
-    )
-    emotion_outputs = emotion_model(**emotion_inputs)
-
-    # Get the predicted emotion label
-    emotion_label_id = emotion_outputs.logits.argmax().item()
-    emotion_label = emotion_tokenizer.decode(emotion_label_id)
-
-    # Prepare the final output with sentiment information
-    final_output = f"The sentiment in the provided image shows: {emotion_label}.\n\nGenerated Caption: {decoded_caption}"
+    # Load the emotion analysis model and tokenizer
+    emotion_model_name = "SamLowe/roberta-base-go_emotions"
+    emotion_classifier = pipeline(model=emotion_model_name)
+
+    results = emotion_classifier(decoded_caption)
+    if results[0]['label'] == 'neutral' or results[0]['score'] <= 0.40:
+        final_output = f"Sentiment of image is not clear, image shows {decoded_caption}."
+    else:
+        final_output = f"Sentiment of the image shows {results[0]['label']}."
+
     return final_output
 
 # Define the Gradio interface