Update app.py

app.py

@@ -1,55 +1,62 @@
-import gradio as gr
 import torch
-import timm
-from torch import nn
-from torchvision import transforms
+import torch.nn as nn
+import gradio as gr
+from torchvision import models, transforms
 from PIL import Image
+from transformers import ViTModel
 
+# Define Hybrid CNN + Transformer
+class HybridCNNTransformer(nn.Module):
+    def __init__(self, num_classes=7):
+        super(HybridCNNTransformer, self).__init__()
+        self.cnn = models.resnet50(pretrained=True)
+        self.cnn = nn.Sequential(*list(self.cnn.children())[:-2])
+        self.channel_reduction = nn.Conv2d(2048, 64, kernel_size=1)
+        self.to_rgb = nn.Conv2d(64, 3, kernel_size=1)
+        self.transformer = ViTModel.from_pretrained("google/vit-base-patch16-224")
+        self.fc = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
 
-# Load the pre-trained SE-ResNet50 model from timm
-model = timm.create_model("seresnet50", pretrained=False)
-model.fc = torch.nn.Linear(2048, 7)  # Adjust for the number of emotion categories
-
-# Load the model weights
-model_path = "custom_resnet50_emotion_model.pth"
+    def forward(self, x):
+        x = self.cnn(x)
+        x = self.channel_reduction(x)
+        x = self.to_rgb(x)
+        x = nn.functional.interpolate(x, size=(224, 224), mode="bilinear")
+        x = self.transformer(pixel_values=x).last_hidden_state[:, 0, :]
+        return self.fc(x)
 
-try:
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # Ensure compatibility
-    model.eval()
-    print("✅ Model loaded successfully.")
-except FileNotFoundError:
-    print("❌ Model file not found. Please check the path.")
-except Exception as e:
-    print(f"❌ Error loading model: {e}")
+# Load model
+model = HybridCNNTransformer(num_classes=7)
+model.load_state_dict(torch.load("transformerHybrid_emotation_model.pth", map_location=torch.device('cpu')), strict=False)
+model.eval()
 
-# Define image transforms
+# Transform
 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]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# Emotion classes (adjust based on your dataset)
-emotions = ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']
-
-# Define the prediction function
+# Prediction function
 def predict_emotion(image):
-    img = transform(image).unsqueeze(0)  # Add batch dimension
-
+    image = transform(image).unsqueeze(0)
     with torch.no_grad():
-        output = model(img)
+        output = model(image)
         probs = torch.nn.functional.softmax(output, dim=1)
-        confidence, predicted_class = probs.max(1)
+        conf, pred = torch.max(probs, 1)
 
-    emotion = emotions[predicted_class.item()]
-    percentage = confidence.item() * 100
-    return emotion, f"{percentage:.2f}%"
+    labels = ["Angry", "Disgust", "Fear", "Happy", "Neutral", "Sad", "Surprise"]
+    return labels[pred.item()], f"{conf.item() * 100:.2f}%"
 
-# Custom CSS for layout styling
+# Interface
 css = """
 body {
+    background-color: #1e1e1e;
     color: white;
-    font-family: Arial, sans-serif;
 }
 #component-1 {
     background-color: rgba(255, 255, 255, 0.7);
@@ -58,19 +65,15 @@ body {
 }
 #component-2 {
     color: black;
+    font-weight: bold;
 }
 """
 
-iface = gr.Interface(
+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.",
+    description="Upload an image to predict the emotion expressed using a Hybrid CNN + ViT model.",
     css=css
-)
-
-# Launch the app
-if __name__ == "__main__":
-    iface.launch()
+).launch()