Update app.py

app.py

@@ -1,15 +1,44 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from torch.nn.functional import sigmoid
 import torch
+from PIL import Image
+from torchvision import transforms
+import requests
 
 # Load text emotion model
 model_name = "SamLowe/roberta-base-go_emotions"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Load image emotion classification pipeline
-image_emotion_pipeline = pipeline("image-classification", model="nateraw/ferplus-emo-resnet34")
+# Load image emotion model (fine-tuned ResNet-50)
+image_model_name = "Celal11/resnet-50-finetuned-FER2013CKPlus-0.003"
+image_emotion_model = AutoModelForSequenceClassification.from_pretrained(image_model_name)
+image_tokenizer = AutoTokenizer.from_pretrained("microsoft/resnet-50")
+
+# Transform for image preprocessing
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# FER labels
+image_labels = [
+    "Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral", "Contempt"
+]
+
+# Analyze image emotion
+def analyze_image_emotion(image_path):
+    if image_path is None:
+        return "No image provided."
+    image = Image.open(image_path).convert("RGB")
+    img_tensor = image_transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = image_emotion_model(img_tensor)
+    probs = sigmoid(output.logits)[0]
+    top_idx = torch.argmax(probs).item()
+    return f"{image_labels[top_idx]} ({probs[top_idx]:.2f})"
 
 # Emotion label to icon mapping (subset)
 emotion_icons = {
@@ -55,14 +84,6 @@ def get_emotions(text, threshold):
 
     return ", ".join(icons) if icons else "No strong emotion detected."
 
-# Analyze image emotion
-def analyze_image_emotion(image):
-    if image is None:
-        return "No image provided."
-    results = image_emotion_pipeline(image)
-    top = results[0]
-    return f"{top['label']} ({top['score']:.2f})"
-
 # Combined analysis
 def analyze_combined(text, threshold, image):
     text_result = get_emotions(text, threshold)
@@ -117,4 +138,4 @@ demo = gr.Interface(
     css=custom_css
 )
 
-demo.launch()
+demo.launch()