Update app.py

app.py

@@ -39,13 +39,6 @@ class IR50(nn.Module):
         self.layer2 = resnet.layer2
         self.downsample = nn.Conv2d(512, 256, 1, stride=2)
         self.bn_downsample = nn.BatchNorm2d(256, eps=1e-5)
-        # Fine-tuned layers (as in training)
-        for param in self.conv1.parameters():
-            param.requires_grad = True
-        for param in self.bn1.parameters():
-            param.requires_grad = True
-        for param in self.layer1.parameters():
-            param.requires_grad = True
 
     def forward(self, x):
         x = self.conv1(x)
@@ -89,7 +82,7 @@ class HLA(nn.Module):
 
 # ViT Stream
 class ViT(nn.Module):
-    def __init__(self, in_channels=256, patch_size=1, embed_dim=768, num_layers=12, num_heads=12):
+    def __init__(self, in_channels=256, patch_size=1, embed_dim=768, num_layers=8, num_heads=12):  # 8 layers as in the 82.93% version
         super(ViT, self).__init__()
         self.patch_embed = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
@@ -196,10 +189,13 @@ class TripleStreamHLAViT(nn.Module):
 
 # Load the model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 model = TripleStreamHLAViT(num_classes=7).to(device)
 model_path = "triple_stream_model_rafdb.pth"  # Ensure this file is in the Hugging Face Space repository
 try:
-    model.load_state_dict(torch.load(model_path, weights_only=True))
+    # Map the weights to the appropriate device
+    map_location = torch.device('cpu') if not torch.cuda.is_available() else None
+    model.load_state_dict(torch.load(model_path, map_location=map_location, weights_only=True))
     model.eval()
     print("Model loaded successfully")
 except Exception as e:
@@ -260,10 +256,9 @@ iface = gr.Interface(
         gr.Image(label="Input Image and HLA Heatmap")
     ],
     title="Facial Emotion Recognition with TripleStreamHLAViT",
-    description="Upload an image to predict the facial emotion (Surprise, Fear, Disgust, Happiness, Sadness, Anger, Neutral). The model also visualizes the HLA heatmap showing where it focuses.",
+    description="Upload an image to predict the facial emotion (Surprise, Fear, Disgust, Happiness, Sadness, Anger, Neutral). This model achieves 82.93% test accuracy on the RAF-DB dataset. The HLA heatmap shows where the model focuses.",
     examples=[
-        ["examples/Surprise.jpg"],
-        ["examples/happy.JPEG"],
+        ["examples/surprise.jpg"],
         ["examples/sadness.jpg"]
     ]
 )