Update app.py

app.py

@@ -14,13 +14,14 @@ import soundfile as sf
 language_mapping = {'malayalam': 0, 'english': 1, 'tamil': 2,'hindi':3}
 
 class ANNModel(nn.Module):
-    def __init__(self):
+    def __init__(self, input_size=300, hidden_size1=256, hidden_size2=64, num_classes=4):
         super(ANNModel, self).__init__()
-        self.fc1 = nn.Linear(300, 256)
+        self.fc1 = nn.Linear(input_size, hidden_size1)
         self.relu1 = nn.ReLU()
-        self.fc2 = nn.Linear(256, 64)
+        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
         self.relu2 = nn.ReLU()
-        self.fc3 = nn.Linear(64, 4)
+        self.fc3 = nn.Linear(hidden_size2, num_classes)
+        self.softmax = nn.Softmax(dim=1)
 
     def forward(self, x):
         x = self.fc1(x)
@@ -28,16 +29,16 @@ class ANNModel(nn.Module):
         x = self.fc2(x)
         x = self.relu2(x)
         x = self.fc3(x)
-        return x
+        x = self.softmax(x)
 
 # Create an instance of your model
 ann_model = ANNModel()
 
 # Load the trained model
-ann_model.load_state_dict(torch.load('ann_model_256_01_94.pth'))
+ann_model.load_state_dict(torch.load('ann_model_256_97400.pth'))
 
 # Load the PCA instance
-pca = load('pca_256_01_94.pkl')
+pca = load('pca_400.pkl')
 
 vgg16 = models.vgg16(pretrained=True).features
 # Function to load and preprocess a single audio file