Update app.py

app.py

@@ -4,14 +4,11 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
-import seaborn as sns
 from facenet_pytorch import InceptionResnetV1, MTCNN
 import mediapipe as mp
 from fer import FER
-from scipy import interpolate
-from sklearn.cluster import DBSCAN, KMeans
-from sklearn.preprocessing import StandardScaler, MinMaxScaler
-from sklearn.metrics import silhouette_score
+from sklearn.cluster import DBSCAN
+from sklearn.preprocessing import MinMaxScaler
 from sklearn.decomposition import PCA
 import umap
 import pandas as pd
@@ -22,24 +19,18 @@ from PIL import Image
 import gradio as gr
 import tempfile
 import shutil
-import io
 
-# Suppress TensorFlow warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-import tensorflow as tf
-
-tf.get_logger().setLevel('ERROR')
 
 matplotlib.rcParams['figure.dpi'] = 400
 matplotlib.rcParams['savefig.dpi'] = 400
 
 # Initialize models and other global variables
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda'
 
-mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.975, 0.975, 0.975], min_face_size=100)
+mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.98, 0.98, 0.98], min_face_size=100)
 model = InceptionResnetV1(pretrained='vggface2').eval().to(device)
 mp_face_mesh = mp.solutions.face_mesh
-face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)
+face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.7)
 emotion_detector = FER(mtcnn=False)
 
 
@@ -188,7 +179,7 @@ def organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder
         shutil.copy(src, dst)
 
 
-def find_optimal_components(embeddings, max_components=10):
+def find_optimal_components(embeddings, max_components=20):
     pca = PCA(n_components=max_components)
     pca.fit(embeddings)
 
@@ -269,7 +260,7 @@ def save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, de
 
 
 class LSTMAutoencoder(nn.Module):
-    def __init__(self, input_size, hidden_size=64, num_layers=2):
+    def __init__(self, input_size, hidden_size=128, num_layers=2):
         super(LSTMAutoencoder, self).__init__()
         self.input_size = input_size
         self.hidden_size = hidden_size
@@ -283,8 +274,8 @@ class LSTMAutoencoder(nn.Module):
         return out
 
 
-def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100, batch_size=64):
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100):
+    device = 'cuda'
     X = torch.FloatTensor(X).to(device)
     if X.dim() == 2:
         X = X.unsqueeze(0)
@@ -328,8 +319,8 @@ def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100
 
     return mse_all, mse_comp, mse_raw
 
-def embedding_anomaly_detection(embeddings, epochs=100, batch_size=64):
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+def embedding_anomaly_detection(embeddings, epochs=100):
+    device = 'cpu'
     X = torch.FloatTensor(embeddings).to(device)
     if X.dim() == 2:
         X = X.unsqueeze(0)
@@ -355,14 +346,14 @@ def embedding_anomaly_detection(embeddings, epochs=100, batch_size=64):
     mse = np.mean(np.power(X.squeeze(0).cpu().numpy() - reconstructed, 2), axis=1)
     return mse
 
-def determine_anomalies(mse_values, threshold=3.5):
+def determine_anomalies(mse_values, threshold=4):
     mean = np.mean(mse_values)
     std = np.std(mse_values)
     anomalies = mse_values > (mean + threshold * std)
     return anomalies
 
 
-def plot_mse(df, mse_values, title, color='blue', time_threshold=1, hide_first_n=3):
+def plot_mse(df, mse_values, title, color='blue', time_threshold=1, hide_first_n=5):
     plt.figure(figsize=(16, 8), dpi=300)
     fig, ax = plt.subplots(figsize=(16, 8))
 
@@ -520,13 +511,13 @@ def process_video(video_path, desired_fps, batch_size, progress=gr.Progress()):
                 X, feature_columns, raw_embedding_columns, batch_size=batch_size)
 
             progress(0.95, "Generating plots")
-            mse_plot_all = plot_mse(df, mse_all, "Facial Features + Emotions", color='blue', hide_first_n=3)
-            mse_plot_comp = plot_mse(df, mse_comp, "Facial Features", color='deepskyblue', hide_first_n=3)
-            mse_plot_raw = plot_mse(df, mse_raw, "Facial Embeddings", color='steelblue', hide_first_n=3)
+            mse_plot_all = plot_mse(df, mse_all, "Facial Features + Emotions", color='blue', hide_first_n=5)
+            mse_plot_comp = plot_mse(df, mse_comp, "Facial Features", color='deepskyblue', hide_first_n=5)
+            mse_plot_raw = plot_mse(df, mse_raw, "Facial Embeddings", color='steelblue', hide_first_n=5)
 
             emotion_plots = [
                 plot_mse(df, embedding_anomaly_detection(df[emotion].values.reshape(-1, 1)),
-                         f"MSE: {emotion.capitalize()}", color=color, hide_first_n=3)
+                         f"MSE: {emotion.capitalize()}", color=color, hide_first_n=5)
                 for emotion, color in zip(['fear', 'sad', 'angry', 'happy', 'surprise', 'neutral'],
                                           ['purple', 'green', 'orange', 'darkblue', 'gold', 'grey'])
             ]
@@ -569,8 +560,8 @@ iface = gr.Interface(
     outputs=[
         gr.Textbox(label="Anomaly Detection Results"),
         gr.Plot(label="MSE: Facial Features + Emotions"),
-        gr.Plot(label="MSE: Facial Features (UMAP)"),
-        gr.Plot(label="MSE: Raw Facial Embeddings"),
+        gr.Plot(label="MSE: Facial Features"),
+        gr.Plot(label="MSE: Facial Embeddings"),
         gr.Plot(label="MSE: Fear"),
         gr.Plot(label="MSE: Sad"),
         gr.Plot(label="MSE: Angry"),
@@ -590,7 +581,7 @@ iface = gr.Interface(
 
         Adjust the parameters as needed:
         - Desired FPS: Frames per second to analyze (lower for faster processing)
-        - Batch Size: Affects processing speed and memory usage
+        - Batch Size: Affects processing speed and GPU memory usage
         """,
     allow_flagging="never"
 )