Update anomaly_detection.py

anomaly_detection.py

@@ -58,7 +58,7 @@ def vae_loss(recon_x, x, mu, logvar):
     KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
     return BCE + KLD
 
-def anomaly_detection(X_embeddings, X_posture, epochs=200, patience=5):
+def anomaly_detection(X_embeddings, X_posture, X_voice, epochs=200, patience=5):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     # Normalize posture
@@ -75,16 +75,24 @@ def anomaly_detection(X_embeddings, X_posture, epochs=200, patience=5):
     if X_posture_scaled.dim() == 2:
         X_posture_scaled = X_posture_scaled.unsqueeze(0)
     
+    # Process voice embeddings
+    X_voice = torch.FloatTensor(X_voice).to(device)
+    if X_voice.dim() == 2:
+        X_voice = X_voice.unsqueeze(0)
+    
     model_embeddings = VAE(input_size=X_embeddings.shape[2]).to(device)
     model_posture = VAE(input_size=X_posture_scaled.shape[2]).to(device)
+    model_voice = VAE(input_size=X_voice.shape[2]).to(device)
     
     optimizer_embeddings = optim.Adam(model_embeddings.parameters())
     optimizer_posture = optim.Adam(model_posture.parameters())
+    optimizer_voice = optim.Adam(model_voice.parameters())
     
     # Train models
-    for epoch in range(epochs):
+    for epoch in range(int(epochs)):  # Ensure epochs is an integer
         for model, optimizer, X in [(model_embeddings, optimizer_embeddings, X_embeddings),
-                                    (model_posture, optimizer_posture, X_posture_scaled)]:
+                                    (model_posture, optimizer_posture, X_posture_scaled),
+                                    (model_voice, optimizer_voice, X_voice)]:
             model.train()
             optimizer.zero_grad()
             recon_batch, mu, logvar = model(X)
@@ -92,16 +100,19 @@ def anomaly_detection(X_embeddings, X_posture, epochs=200, patience=5):
             loss.backward()
             optimizer.step()
     
-    # Compute reconstruction error for embeddings and posture
+    # Compute reconstruction error for embeddings, posture, and voice
     model_embeddings.eval()
     model_posture.eval()
+    model_voice.eval()
     with torch.no_grad():
         recon_embeddings, _, _ = model_embeddings(X_embeddings)
         recon_posture, _, _ = model_posture(X_posture_scaled)
+        recon_voice, _, _ = model_voice(X_voice)
         mse_embeddings = F.mse_loss(recon_embeddings, X_embeddings, reduction='none').mean(dim=2).cpu().numpy().squeeze()
         mse_posture = F.mse_loss(recon_posture, X_posture_scaled, reduction='none').mean(dim=2).cpu().numpy().squeeze()
+        mse_voice = F.mse_loss(recon_voice, X_voice, reduction='none').mean(dim=2).cpu().numpy().squeeze()
     
-    return mse_embeddings, mse_posture
+    return mse_embeddings, mse_posture, mse_voice
 
 def determine_anomalies(mse_values, threshold):
     mean = np.mean(mse_values)