Update app.py

app.py

@@ -4,37 +4,42 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
+import torchvision
 from facenet_pytorch import InceptionResnetV1, MTCNN
 import mediapipe as mp
 from fer import FER
 from sklearn.cluster import KMeans
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
-from sklearn.metrics import silhouette_score
-from scipy.spatial.distance import cdist
 import umap
 import pandas as pd
 import matplotlib.pyplot as plt
 from matplotlib.ticker import MaxNLocator
+from moviepy.editor import VideoFileClip
+from PIL import Image
 import gradio as gr
 import tempfile
 import shutil
 import subprocess
 import fractions
 
+
 # Suppress TensorFlow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 import tensorflow as tf
+
 tf.get_logger().setLevel('ERROR')
 
 # 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.999, 0.999, 0.999], min_face_size=100, selection_method='largest')
+mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.999, 0.999, 0.999], min_face_size=100,
+              selection_method='largest')
 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)
 emotion_detector = FER(mtcnn=False)
 
+
 def frame_to_timecode(frame_num, original_fps, desired_fps):
     total_seconds = frame_num / original_fps
     hours = int(total_seconds // 3600)
@@ -43,6 +48,7 @@ def frame_to_timecode(frame_num, original_fps, desired_fps):
     milliseconds = int((total_seconds - int(total_seconds)) * 1000)
     return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}"
 
+
 def get_face_embedding_and_emotion(face_img):
     face_tensor = torch.tensor(face_img).permute(2, 0, 1).unsqueeze(0).float() / 255
     face_tensor = (face_tensor - 0.5) / 0.5
@@ -58,6 +64,7 @@ def get_face_embedding_and_emotion(face_img):
 
     return embedding.cpu().numpy().flatten(), emotion_dict
 
+
 def alignFace(img):
     img_raw = img.copy()
     results = face_mesh.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
@@ -83,53 +90,51 @@ def alignFace(img):
     new_img = cv2.warpAffine(img_raw, rotation_matrix, (width, height))
     return new_img
 
-def extract_frames(video_path, output_folder, fps):
+
+def extract_frames(video_path, output_folder, desired_fps, progress_callback=None):
     os.makedirs(output_folder, exist_ok=True)
-    command = [
-        'ffmpeg',
-        '-i', video_path,
-        '-vf', f'fps={fps}',
-        f'{output_folder}/frame_%04d.jpg'
-    ]
-    try:
-        result = subprocess.run(command, check=True, capture_output=True, text=True)
-        print(f"FFmpeg stdout: {result.stdout}")
-        print(f"FFmpeg stderr: {result.stderr}")
-    except subprocess.CalledProcessError as e:
-        print(f"Error extracting frames: {e}")
-        print(f"FFmpeg stdout: {e.stdout}")
-        print(f"FFmpeg stderr: {e.stderr}")
-        raise
-
-def get_video_info(video_path):
-    ffprobe_command = [
-        'ffprobe',
-        '-v', 'error',
-        '-select_streams', 'v:0',
-        '-count_packets',
-        '-show_entries', 'stream=nb_read_packets,r_frame_rate',
-        '-of', 'csv=p=0',
-        video_path
-    ]
-    ffprobe_output = subprocess.check_output(ffprobe_command, universal_newlines=True).strip().split(',')
-    frame_rate, frame_count = ffprobe_output
-    
-    frac = fractions.Fraction(frame_rate)
-    original_fps = float(frac.numerator) / float(frac.denominator)
-    frame_count = int(frame_count)
-    
+
+    # Load the video clip
+    clip = VideoFileClip(video_path)
+
+    original_fps = clip.fps
+    duration = clip.duration
+    total_frames = int(duration * original_fps)
+    step = max(1, original_fps / desired_fps)
+    total_frames_to_extract = int(total_frames / step)
+
+    frame_count = 0
+    for t in np.arange(0, duration, step / original_fps):
+        # Get the frame at time t
+        frame = clip.get_frame(t)
+
+        # Convert the frame to PIL Image and save it
+        img = Image.fromarray(frame)
+        img.save(os.path.join(output_folder, f"frame_{frame_count:04d}.jpg"))
+
+        frame_count += 1
+
+        # Report progress
+        if progress_callback:
+            progress = frame_count / total_frames_to_extract
+            progress_callback(progress, f"Extracting frame {frame_count} of {total_frames_to_extract}")
+
+        if frame_count >= total_frames_to_extract:
+            break
+
+    clip.close()
     return frame_count, original_fps
 
 def process_frames(frames_folder, aligned_faces_folder, frame_count, progress, batch_size):
     embeddings_by_frame = {}
     emotions_by_frame = {}
     frame_files = sorted([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])
-    
+
     for i in range(0, len(frame_files), batch_size):
-        batch_files = frame_files[i:i+batch_size]
+        batch_files = frame_files[i:i + batch_size]
         batch_frames = []
         batch_nums = []
-        
+
         for frame_file in batch_files:
             frame_num = int(frame_file.split('_')[1].split('.')[0])
             frame_path = os.path.join(frames_folder, frame_file)
@@ -137,12 +142,13 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress, b
             if frame is not None:
                 batch_frames.append(frame)
                 batch_nums.append(frame_num)
-        
+
         if batch_frames:
             # Detect faces in batch
             batch_boxes, batch_probs = mtcnn.detect(batch_frames)
-            
-            for j, (frame, frame_num, boxes, probs) in enumerate(zip(batch_frames, batch_nums, batch_boxes, batch_probs)):
+
+            for j, (frame, frame_num, boxes, probs) in enumerate(
+                    zip(batch_frames, batch_nums, batch_boxes, batch_probs)):
                 if boxes is not None and len(boxes) > 0 and probs[0] >= 0.99:
                     x1, y1, x2, y2 = [int(b) for b in boxes[0]]
                     face = frame[y1:y2, x1:x2]
@@ -155,11 +161,13 @@ def process_frames(frames_folder, aligned_faces_folder, frame_count, progress, b
                             embedding, emotion = get_face_embedding_and_emotion(aligned_face_resized)
                             embeddings_by_frame[frame_num] = embedding
                             emotions_by_frame[frame_num] = emotion
-        
-        progress((i + len(batch_files)) / frame_count, f"Processing frames {i + 1} to {min(i + len(batch_files), frame_count)} of {frame_count}")
+
+        progress((i + len(batch_files)) / frame_count,
+                 f"Processing frames {i + 1} to {min(i + len(batch_files), frame_count)} of {frame_count}")
 
     return embeddings_by_frame, emotions_by_frame
 
+
 def cluster_embeddings(embeddings):
     if len(embeddings) < 2:
         print("Not enough embeddings for clustering. Assigning all to one cluster.")
@@ -171,6 +179,7 @@ def cluster_embeddings(embeddings):
     clusters = kmeans.fit_predict(embeddings_scaled)
     return clusters
 
+
 def organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder):
     for (frame_num, embedding), cluster in zip(embeddings_by_frame.items(), clusters):
         person_folder = os.path.join(organized_faces_folder, f"person_{cluster}")
@@ -179,7 +188,9 @@ def organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder
         dst = os.path.join(person_folder, f"frame_{frame_num}_face.jpg")
         shutil.copy(src, dst)
 
-def save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps, original_fps, output_folder, num_components):
+
+def save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps, original_fps, output_folder,
+                            num_components):
     emotions = ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']
     person_data = {}
 
@@ -224,6 +235,7 @@ def save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, de
 
     return df, largest_cluster
 
+
 class LSTMAutoencoder(nn.Module):
     def __init__(self, input_size, hidden_size=64, num_layers=2):
         super(LSTMAutoencoder, self).__init__()
@@ -239,6 +251,7 @@ class LSTMAutoencoder(nn.Module):
         out = self.fc(outputs)
         return out
 
+
 def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, batch_size=64):
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
@@ -287,9 +300,10 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     # Compute anomalies for components only
     component_columns = [col for col in feature_columns if col.startswith('Comp')]
     component_indices = [feature_columns.index(col) for col in component_columns]
-    
+
     if len(component_indices) > 0:
-        mse_comp = np.mean(np.power(X.squeeze(0).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,10 +311,11 @@ def lstm_anomaly_detection(X, feature_columns, num_anomalies=10, epochs=100, bat
     anomalies_comp = np.zeros(len(mse_comp), dtype=bool)
     anomalies_comp[top_indices_comp] = True
 
-    return (anomalies_all, mse_all, top_indices_all, 
-            anomalies_comp, mse_comp, top_indices_comp, 
+    return (anomalies_all, mse_all, top_indices_all,
+            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))
     bars = ax.bar(range(len(df)), anomaly_scores, width=0.8, color='skyblue')
@@ -311,10 +326,12 @@ def plot_anomaly_scores(df, anomaly_scores, top_indices, title):
     ax.set_title(f'Anomaly Scores Over Time ({title})')
     ax.xaxis.set_major_locator(MaxNLocator(nbins=100))
     ticks = ax.get_xticks()
-    ax.set_xticklabels([df['Timecode'].iloc[int(tick)] if tick >= 0 and tick < len(df) else '' for tick in ticks], rotation=90, ha='right')
+    ax.set_xticklabels([df['Timecode'].iloc[int(tick)] if tick >= 0 and tick < len(df) else '' for tick in ticks],
+                       rotation=90, ha='right')
     plt.tight_layout()
     return fig
 
+
 def plot_emotion(df, emotion, num_anomalies, color):
     fig, ax = plt.subplots(figsize=(16, 8))
     values = df[emotion].values
@@ -327,12 +344,15 @@ def plot_emotion(df, emotion, num_anomalies, color):
     ax.set_title(f'{emotion.capitalize()} Anomalies Over Time (Top {num_anomalies} in Red)')
     ax.xaxis.set_major_locator(MaxNLocator(nbins=100))
     ticks = ax.get_xticks()
-    ax.set_xticklabels([df['Timecode'].iloc[int(tick)] if tick >= 0 and tick < len(df) else '' for tick in ticks], rotation=90, ha='right')
+    ax.set_xticklabels([df['Timecode'].iloc[int(tick)] if tick >= 0 and tick < len(df) else '' for tick in ticks],
+                       rotation=90, ha='right')
     plt.tight_layout()
     return fig
 
+
 import base64
 
+
 def get_random_face_sample(organized_faces_folder, largest_cluster, output_folder):
     person_folder = os.path.join(organized_faces_folder, f"person_{largest_cluster}")
     face_files = [f for f in os.listdir(person_folder) if f.endswith('.jpg')]
@@ -340,34 +360,40 @@ def get_random_face_sample(organized_faces_folder, largest_cluster, output_folde
         random_face = np.random.choice(face_files)
         face_path = os.path.join(person_folder, random_face)
         output_path = os.path.join(output_folder, "random_face_sample.jpg")
-        
+
         # Read the image and resize it to be smaller
         face_img = cv2.imread(face_path)
         small_face = cv2.resize(face_img, (100, 100))  # Resize to NxN pixels
         cv2.imwrite(output_path, small_face)
-        
+
         return output_path
     return None
 
+
 def process_video(video_path, num_anomalies, num_components, desired_fps, batch_size, progress=gr.Progress()):
     output_folder = "output"
     os.makedirs(output_folder, exist_ok=True)
-    
+
     with tempfile.TemporaryDirectory() as temp_dir:
         aligned_faces_folder = os.path.join(temp_dir, 'aligned_faces')
         organized_faces_folder = os.path.join(temp_dir, 'organized_faces')
         os.makedirs(aligned_faces_folder, exist_ok=True)
         os.makedirs(organized_faces_folder, exist_ok=True)
 
-        progress(0.1, "Extracting frames")
+        progress(0.05, "Starting frame extraction")
         frames_folder = os.path.join(temp_dir, 'extracted_frames')
-        extract_frames(video_path, frames_folder, desired_fps)
 
-        progress(0.2, "Getting video info")
-        frame_count, original_fps = get_video_info(video_path)
+        def extraction_progress(percent, message):
+            # Adjust the progress to fit within the 5% to 30% range of the overall process
+            overall_progress = 0.05 + (percent * 0.25)
+            progress(overall_progress, message)
+
+        frame_count, original_fps = extract_frames(video_path, frames_folder, desired_fps, extraction_progress)
+
 
         progress(0.3, "Processing frames")
-        embeddings_by_frame, emotions_by_frame = process_frames(frames_folder, aligned_faces_folder, frame_count, progress, batch_size)
+        embeddings_by_frame, emotions_by_frame = process_frames(frames_folder, aligned_faces_folder, frame_count,
+                                                                progress, batch_size)
 
         if not embeddings_by_frame:
             return "No faces were extracted from the video.", None, None, None, None, None, None
@@ -380,15 +406,18 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder)
 
         progress(0.8, "Saving person data")
-        df, largest_cluster = save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps, original_fps, temp_dir, num_components)
+        df, largest_cluster = save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps,
+                                                      original_fps, temp_dir, num_components)
 
         progress(0.9, "Performing anomaly detection")
-        feature_columns = [col for col in df.columns if col not in ['Frame', 'Timecode', 'Time (Minutes)', 'Embedding_Index']]
+        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}")
         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)
+            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}")
@@ -410,38 +439,39 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         except Exception as e:
             return f"Error generating plots: {str(e)}", None, None, None, None, None, None, None, None, None
 
-
         # Get a random face sample
         face_sample = get_random_face_sample(organized_faces_folder, largest_cluster, output_folder)
-        
+
         progress(1.0, "Preparing results")
         results = f"Top {num_anomalies} anomalies (All Features):\n"
-        results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in 
+        results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in
                               zip(anomaly_scores_all[top_indices_all], df['Timecode'].iloc[top_indices_all].values)])
         results += f"\n\nTop {num_anomalies} anomalies (Components Only):\n"
-        results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in 
+        results += "\n".join([f"{score:.4f} at {timecode}" for score, timecode in
                               zip(anomaly_scores_comp[top_indices_comp], df['Timecode'].iloc[top_indices_comp].values)])
 
         for emotion in ['fear', 'sad', 'angry', 'happy', 'surprise', 'neutral']:
             top_indices = np.argsort(df[emotion].values)[-num_anomalies:][::-1]
             results += f"\n\nTop {num_anomalies} {emotion.capitalize()} Scores:\n"
             results += "\n".join([f"{df[emotion].iloc[i]:.4f} at {df['Timecode'].iloc[i]}" for i in top_indices])
-    
+
         return (
             results,  # Text results to a Textbox
             face_sample,  # Random face sample image
             anomaly_plot_all,
             anomaly_plot_comp,
             *emotion_plots
-        )        
-        
-# Gradio interface
+        )
+
+    # Gradio interface
+
+
 iface = gr.Interface(
     fn=process_video,
     inputs=[
         gr.Video(),
         gr.Slider(minimum=1, maximum=20, step=1, value=5, label="Number of Anomalies"),
-        gr.Slider(minimum=2, maximum=5, step=1, value=3, label="Number of Components"),
+        gr.Slider(minimum=1, maximum=20, step=1, value=10, label="Number of Components"),
         gr.Slider(minimum=1, maximum=30, step=1, value=20, label="Desired FPS"),
         gr.Slider(minimum=1, maximum=64, step=1, value=16, label="Batch Size")
     ],
@@ -461,7 +491,7 @@ iface = gr.Interface(
     description="""
     This application detects anomalies in facial expressions and emotions from a video input. 
     It focuses on the most frequently appearing person in the video for analysis.
-    
+
     Adjust the parameters as needed:
     - Number of Anomalies: How many top anomalies or high intensities to highlight
     - Number of Components: Complexity of the facial expression model
@@ -471,4 +501,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()