Update app.py

app.py

@@ -1,44 +1,31 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoImageProcessor, AutoModelForImageClassification
 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 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")
+# Load image emotion model
+image_model_name = "Celal11/resnet-50-finetuned-FER2013-0.001"
+image_processor = AutoImageProcessor.from_pretrained(image_model_name)
+image_model = AutoModelForImageClassification.from_pretrained(image_model_name)
 
-# 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:
+# Analyze image emotion using processor and model
+def analyze_image_emotion(image):
+    if image is None:
         return "No image provided."
-    image = Image.open(image_path).convert("RGB")
-    img_tensor = image_transform(image).unsqueeze(0)
+    inputs = image_processor(images=image, return_tensors="pt")
     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})"
+        logits = image_model(**inputs).logits
+    probs = torch.nn.functional.softmax(logits, dim=1)[0]
+    pred_idx = torch.argmax(probs).item()
+    label = image_model.config.id2label[pred_idx]
+    score = probs[pred_idx].item()
+    return f"{label} ({score:.2f})"
 
 # Emotion label to icon mapping (subset)
 emotion_icons = {
@@ -126,7 +113,7 @@ demo = gr.Interface(
     inputs=[
         gr.Textbox(lines=5, placeholder="Write a sentence or a full paragraph...", label="Your Text"),
         gr.Slider(minimum=0.1, maximum=0.9, value=0.3, step=0.05, label="Threshold"),
-        gr.Image(type="filepath", label="Upload Face Photo")
+        gr.Image(type="pil", label="Upload Face Photo")
     ],
     outputs=[
         gr.Textbox(label="Detected Text Emotions", elem_classes=["output-textbox"]),
@@ -138,4 +125,4 @@ demo = gr.Interface(
     css=custom_css
 )
 
-demo.launch()
+demo.launch()