Update app.py

app.py

@@ -235,8 +235,8 @@ class LSTMAutoencoder(nn.Module):
         self.fc = nn.Linear(hidden_size, input_size)
 
     def forward(self, x):
-        _, (hidden, _) = self.lstm(x)
-        out = self.fc(hidden[-1])
+        outputs, (hidden, _) = self.lstm(x)
+        out = self.fc(outputs)
         return out
 
 def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, batch_size=64):
@@ -244,38 +244,42 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
 
     X = torch.FloatTensor(X).to(device)
 
-    # Ensure X is 2D
-    if X.dim() == 1:
-        X = X.unsqueeze(1)  # Add a feature dimension
-    elif X.dim() > 2:
-        raise ValueError(f"Input X should be 1D or 2D, but got {X.dim()} dimensions")
+    # Ensure X is 3D (batch, sequence, features)
+    if X.dim() == 2:
+        X = X.unsqueeze(0)
+    elif X.dim() == 1:
+        X = X.unsqueeze(0).unsqueeze(2)
+    elif X.dim() > 3:
+        raise ValueError(f"Input X should be 1D, 2D or 3D, but got {X.dim()} dimensions")
 
-    train_size = int(0.85 * len(X))
-    X_train, X_val = X[:train_size], X[train_size:]
+    print(f"X shape after reshaping: {X.shape}")
 
-    model = LSTMAutoencoder(input_size=X.shape[1]).to(device)
+    train_size = int(0.85 * X.shape[1])
+    X_train, X_val = X[:, :train_size, :], X[:, train_size:, :]
+
+    model = LSTMAutoencoder(input_size=X.shape[2]).to(device)
     criterion = nn.MSELoss()
     optimizer = optim.Adam(model.parameters())
 
     for epoch in range(epochs):
         model.train()
         optimizer.zero_grad()
-        output_train = model(X_train.unsqueeze(0))
-        loss_train = criterion(output_train, X_train)
+        output_train = model(X_train)
+        loss_train = criterion(output_train, X_train.squeeze(0))
         loss_train.backward()
         optimizer.step()
 
         model.eval()
         with torch.no_grad():
-            output_val = model(X_val.unsqueeze(0))
-            loss_val = criterion(output_val, X_val)
+            output_val = model(X_val)
+            loss_val = criterion(output_val, X_val.squeeze(0))
 
     model.eval()
     with torch.no_grad():
-        reconstructed = model(X.unsqueeze(0)).squeeze(0).cpu().numpy()
+        reconstructed = model(X).squeeze(0).cpu().numpy()
 
     # Compute anomalies for all features
-    mse_all = np.mean(np.power(X.cpu().numpy() - reconstructed, 2), axis=1)
+    mse_all = np.mean(np.power(X.squeeze(0).cpu().numpy() - reconstructed, 2), axis=1)
     top_indices_all = mse_all.argsort()[-num_anomalies:][::-1]
     anomalies_all = np.zeros(len(mse_all), dtype=bool)
     anomalies_all[top_indices_all] = True
@@ -285,10 +289,7 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     component_indices = [feature_columns.index(col) for col in component_columns]
     
     if len(component_indices) > 0:
-        if X.dim() == 1:
-            mse_comp = mse_all  # If X is 1D, we can't select specific components
-        else:
-            mse_comp = np.mean(np.power(X.cpu().numpy()[:, component_indices] - reconstructed[:, component_indices], 2), axis=1)
+        mse_comp = np.mean(np.power(X.squeeze(0).cpu().numpy()[:, component_indices] - reconstructed[:, component_indices], 2), axis=1)
     else:
         mse_comp = mse_all  # If no components, use all features
 
@@ -297,8 +298,8 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     anomalies_comp[top_indices_comp] = True
 
     return (anomalies_all, mse_all, top_indices_all, 
-                anomalies_comp, mse_comp, top_indices_comp, 
-                model)
+            anomalies_comp, mse_comp, top_indices_comp, 
+            model)
 
 def plot_anomaly_scores(df, anomaly_scores, top_indices, title):
     fig, ax = plt.subplots(figsize=(16, 8))
@@ -375,10 +376,14 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         progress(0.9, "Performing anomaly detection")
         feature_columns = [col for col in df.columns if col not in ['Frame', 'Timecode', 'Time (Minutes)', 'Embedding_Index']]
         X = df[feature_columns].values
-        print(f"Shape of input data: {X.shape}")  # Debug print
+        print(f"Shape of input data: {X.shape}")
+        print(f"Feature columns: {feature_columns}")
         try:
             anomalies_all, anomaly_scores_all, top_indices_all, anomalies_comp, anomaly_scores_comp, top_indices_comp, _ = lstm_anomaly_detection(X, feature_columns, num_anomalies=num_anomalies, batch_size=batch_size)
         except Exception as e:
+            print(f"Error details: {str(e)}")
+            print(f"X shape: {X.shape}")
+            print(f"X dtype: {X.dtype}")
             return f"Error in anomaly detection: {str(e)}", None, None, None, None, None, None
         
         progress(0.95, "Generating plots")