Update video_processing.py

video_processing.py

@@ -130,6 +130,9 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
         clusters = cluster_faces(embeddings)
         num_clusters = len(set(clusters))
 
+        # Adding the 'Cluster' column to the DataFrame
+        cluster_by_frame = {frame_num: cluster for frame_num, cluster in zip(embeddings_by_frame.keys(), clusters)}
+
         progress(0.65, "Organizing faces")
         organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder)
 
@@ -139,6 +142,7 @@ def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
 
         df['Seconds'] = df['Timecode'].apply(
             lambda x: sum(float(t) * 60 ** i for i, t in enumerate(reversed(x.split(':')))))
+        df['Cluster'] = df['Frame'].map(cluster_by_frame)
 
         progress(0.75, "Getting face samples")
         face_samples = get_all_face_samples(organized_faces_folder, output_folder, largest_cluster)