Update app.py

app.py

@@ -1,52 +1,49 @@
 import torch
+import timm
+import torch.nn as nn
 import gradio as gr
-from transformers import AutoModel
 from PIL import Image
 from torchvision import transforms
 
-# 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)
+# Define your custom model architecture (HybridCNNTransformer in this case)
+class HybridCNNTransformer(nn.Module):
+    def __init__(self, num_classes=7):
+        super(HybridCNNTransformer, self).__init__()
+        
+        # Example: Using ResNet50 from timm as a CNN feature extractor
+        self.backbone = timm.create_model('resnet50', pretrained=True)
+        
+        # Example Transformer part (modify according to your model)
+        self.transformer = nn.Transformer(d_model=2048, nhead=8, num_encoder_layers=6)
+        
+        # Final fully connected layer (7 classes for emotion recognition)
+        self.fc = nn.Linear(2048, num_classes)
 
-# 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
+    def forward(self, x):
+        # CNN feature extraction
+        cnn_features = self.backbone(x)
+        
+        # Transformer encoding (if applicable, you might not need this part)
+        transformer_features = self.transformer(cnn_features, cnn_features)
+        
+        # Final classification layer
+        output = self.fc(transformer_features)
+        return output
 
-# 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 model
+model = HybridCNNTransformer(num_classes=7)
 
-# Move to the appropriate device (GPU or CPU)
+# Load the weights from the .pth file
+model_path = "transformer_emotion_recognition_model.pth"  # Replace with the path to your .pth file
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # For CPU; change 'cpu' to 'cuda' for GPU
+
+# Move the model to the appropriate device (CUDA or CPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Make sure the model is in evaluation mode
+# Set the model to evaluation mode
 model.eval()
 
-# 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 ImageNet normalization
-])
-
-# Prediction function
-def predict_emotion(image):
-    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(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)
-
 # Custom CSS for layout styling
 css = """
 body {
@@ -55,26 +52,22 @@ body {
     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);
 }
-
 #component-2 {
     color: black;
     font-weight: bold;
 }
-
 #title {
     color: white;
     font-size: 36px;
     font-weight: bold;
     text-align: center;
 }
-
 #description {
     color: white;
     font-size: 16px;
@@ -83,6 +76,29 @@ body {
 }
 """
 
+# Image Preprocessing for the model (assuming the model was trained with resized and normalized images)
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),  # Adjust according to your model's input size
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Standard ImageNet normalization
+])
+
+# Prediction function
+def predict_emotion(image):
+    # Preprocess the image
+    image_tensor = preprocess(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
+
+    # Make prediction
+    with torch.no_grad():
+        output = model(image_tensor)
+        _, predicted = torch.max(output, 1)  # Get the predicted class
+        confidence = torch.nn.functional.softmax(output, dim=1).max().item()  # Confidence score
+
+    # Return the predicted emotion label and confidence score
+    emotions = ["Anger", "Disgust", "Fear", "Happiness", "Sadness", "Surprise", "Neutral"]  # Modify labels as per your model
+    predicted_emotion = emotions[predicted.item()]
+    return predicted_emotion, confidence
+
 # Gradio Interface
 iface = gr.Interface(
     fn=predict_emotion,