Update app.py

app.py

@@ -1,76 +1,99 @@
-import gradio as gr
 import torch
-import timm
-from torch import nn
-from torchvision import transforms
+import gradio as gr
+from transformers import AutoModel
 from PIL import Image
+from torchvision import transforms
 
-# Define the class mapping for the RAF-DB dataset
-class_mapping = {
-    "1": "Surprise",
-    "2": "Fear",
-    "3": "Disgust",
-    "4": "Happiness",
-    "5": "Sadness",
-    "6": "Anger",
-    "7": "Contempt"
-}
+# Load your custom model from Hugging Face (replace with your actual model)
+model_name = 'Sakibrumu/HybridCNNTransformer'  # Replace with your Hugging Face model ID
+model = AutoModel.from_pretrained(model_name)
+
+# If you need to fine-tune or adjust the final layer
+model.fc = torch.nn.Linear(2048, 7)  # Adjust the final layer for 7 emotion categories
 
-# Define the function to map folder names to labels
-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 weights (you might not need this if your model is already fine-tuned in Hugging Face)
+model.load_state_dict(torch.load("transformer_emotion_recognition_model.pth"))
 
-# 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
+# Move to the appropriate device (GPU or CPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# Define the emotion classes
-emotion_classes = list(class_mapping.values())  # Get emotion classes from the class mapping
+# Make sure the model is in evaluation mode
+model.eval()
 
-# Define the transformation pipeline
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),  # Resize image to 224x224 (adjust as needed)
+# Image Preprocessing (e.g., resizing and normalization)
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize to the expected input size
     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
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Standard ImageNet normalization
 ])
 
-# Function for Emotion Prediction
+# Prediction function
 def predict_emotion(image):
-    img = Image.open(image).convert('RGB')  # Convert the uploaded image into RGB
-    img = transform(img).unsqueeze(0)  # Apply transformations and add batch dimension
+    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    image = preprocess(image).unsqueeze(0).to(device)  # Preprocess 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}%"
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)  # Get the class with the highest probability
+    
+    # Assuming you have an emotion label list
+    emotion_labels = ['Anger', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']
+    predicted_label = emotion_labels[predicted.item()]
+    
+    # Confidence is the probability of the predicted class
+    confidence = torch.nn.functional.softmax(outputs, dim=1)
+    predicted_confidence = confidence[0, predicted.item()].item()
+    
+    return predicted_label, round(predicted_confidence * 100, 2)
 
-# Gradio interface with xkcd theme
-with gr.Blocks(theme="gstaff/xkcd") as demo:
-    gr.Markdown("# Emotion Recognition Classifier")
-    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
-        * Fear
-        * Happiness
-        * Sadness
-        * Surprise
-        * Contempt
-    """)
+# Custom CSS for layout styling
+css = """
+body {
+    background-color: #1e1e1e;
+    color: white;
+    font-family: Arial, sans-serif;
+    padding: 20px;
+}
+
+#component-1 {
+    background-color: rgba(255, 255, 255, 0.7);
+    padding: 20px;
+    border-radius: 10px;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+}
 
-    # Upload image widget
-    image_input = gr.Image(type="pil", label="Upload an image of a face")
+#component-2 {
+    color: black;
+    font-weight: bold;
+}
 
-    # Outputs
-    label_output = gr.Textbox(label="Predicted Emotion")
-    confidence_output = gr.Textbox(label="Confidence Percentage")
+#title {
+    color: white;
+    font-size: 36px;
+    font-weight: bold;
+    text-align: center;
+}
+
+#description {
+    color: white;
+    font-size: 16px;
+    text-align: center;
+    margin-bottom: 20px;
+}
+"""
 
-    # Button to predict the emotion
-    image_input.upload(predict_emotion, image_input, [label_output, confidence_output])
+# Gradio Interface
+iface = gr.Interface(
+    fn=predict_emotion,
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Textbox(label="Predicted Emotion"), gr.Textbox(label="Confidence")],
+    live=True,
+    title="Emotion Classification",
+    description="Upload an image to predict the emotion expressed in the image using a fine-tuned SE-ResNet50 model.",
+    css=css
+)
 
-    # Launch the app
-    demo.launch(share=True)
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()