Update app.py

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
-from fastai.vision.all import load_learner, PILImage
-from PIL import Image
 import torch
+import timm
+from torch import nn
+from torchvision import transforms
+from PIL import Image
 
 # Define the class mapping for the RAF-DB dataset
 class_mapping = {
@@ -19,25 +21,37 @@ def get_raf_label(file_path):
     # Use the class mapping to get the label
     return class_mapping[str(file_path.parent.name)]
 
-# Load the model
-model_path = "efficientnet_emotion_model.pkl"  # Replace with your actual model path
-model = load_learner(model_path)
+# Load the SeresNet50 model using timm
+model_path = "custom_seresnet50_emotion_model.pth"  # Replace with your actual .pth model path
+model = timm.create_model('seresnet50', pretrained=False, num_classes=7)  # Load the SeresNet50 model with 7 classes
+model.load_state_dict(torch.load(model_path))  # Load the saved model weights
+model.eval()  # Set the model to evaluation mode
 
 # Define the emotion classes
 emotion_classes = list(class_mapping.values())  # Get emotion classes from the class mapping
 
+# Define the transformation pipeline
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize image to 224x224 (adjust as needed)
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Standard normalization for models trained on ImageNet
+])
+
 # Function for Emotion Prediction
 def predict_emotion(image):
-    img = PILImage.create(image)  # Convert the uploaded image into a PIL image
-    pred_class, pred_idx, outputs = model.predict(img)
-    predicted_emotion = emotion_classes[pred_idx]
-    confidence = outputs[pred_idx] * 100  # Convert to percentage
+    img = Image.open(image).convert('RGB')  # Convert the uploaded image into RGB
+    img = transform(img).unsqueeze(0)  # Apply transformations and add batch dimension
+    with torch.no_grad():
+        outputs = model(img)  # Get model output
+    _, pred_idx = torch.max(outputs, 1)  # Get the predicted class index
+    predicted_emotion = emotion_classes[pred_idx.item()]
+    confidence = torch.nn.functional.softmax(outputs, dim=1)[0][pred_idx].item() * 100  # Get confidence
     return predicted_emotion, f"{confidence:.2f}%"
 
 # Gradio interface with xkcd theme
 with gr.Blocks(theme="gstaff/xkcd") as demo:
     gr.Markdown("# Emotion Recognition Classifier")
-    gr.Markdown("""
+    gr.Markdown(""" 
         This app uses a deep learning model to recognize emotions in facial images. 
         The model has been trained on a dataset to classify images into different emotion categories:
         * Anger
@@ -45,12 +59,12 @@ with gr.Blocks(theme="gstaff/xkcd") as demo:
         * Happiness
         * Sadness
         * Surprise
-        * Neutral
+        * Contempt
     """)
-    
+
     # Upload image widget
     image_input = gr.Image(type="pil", label="Upload an image of a face")
-    
+
     # Outputs
     label_output = gr.Textbox(label="Predicted Emotion")
     confidence_output = gr.Textbox(label="Confidence Percentage")