Update app.py

app.py

@@ -12,6 +12,9 @@ from moviepy.editor import VideoFileClip
 from sklearn.cluster import DBSCAN
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
+from collections import defaultdict
+
+
 
 # Load models
 @st.cache_resource
@@ -53,48 +56,75 @@ def load_indexes(video_id, output_dir):
     return text_index, image_index
 
 text_index, image_index = load_indexes(video_id, output_dir)
-
-# Comprehensive face summarization
-def create_comprehensive_face_summary(face_index, face_metadata, eps=0.5, min_samples=3):
+def create_comprehensive_face_summary(face_index, face_metadata, eps=0.5, min_samples=3, top_k=5):
+    # Extract face embeddings
     face_embeddings = face_index.reconstruct_n(0, face_index.ntotal)
     
+    # Normalize embeddings
+    face_embeddings = face_embeddings / np.linalg.norm(face_embeddings, axis=1)[:, np.newaxis]
+    
+    # Perform DBSCAN clustering
     clustering = DBSCAN(eps=eps, min_samples=min_samples, metric='cosine').fit(face_embeddings)
     
-    face_clusters = {}
+    # Group faces by cluster
+    face_clusters = defaultdict(list)
     for i, label in enumerate(clustering.labels_):
-        if label not in face_clusters:
-            face_clusters[label] = []
-        face_clusters[label].append(i)
+        face_clusters[label].append(face_metadata[i])
+    
+    # Sort clusters by size
+    sorted_clusters = sorted(face_clusters.items(), key=lambda x: len(x[1]), reverse=True)
+    
+    all_faces_summary = []
+    prominent_faces = []
     
-    summary = []
-    for label, indices in face_clusters.items():
+    for i, (label, cluster) in enumerate(sorted_clusters):
         if label != -1:  # Ignore noise points
-            cluster_appearances = [face_metadata[i] for i in indices]
-            cluster_summary = {
+            # Collect all appearances
+            appearances = [
+                {
+                    'start': face['start'],
+                    'end': face['end'],
+                    'size_ratio': face.get('size_ratio', 1.0)  # Use 1.0 as default if size_ratio is not present
+                }
+                for face in cluster
+            ]
+            
+            # Sort appearances by start time
+            appearances.sort(key=lambda x: x['start'])
+            
+            # Select representative face (e.g., largest face in the cluster)
+            representative_face = max(cluster, key=lambda f: f.get('size_ratio', 1.0))
+            
+            face_summary = {
+                "id": f"face_{i}",
                 "cluster_id": f"cluster_{label}",
-                "face_count": len(indices),
-                "appearances": cluster_appearances,
-                "timeline": [
-                    {"start": app['start'], "end": app['end']} 
-                    for app in cluster_appearances
-                ],
-                "total_screen_time": sum(app['end'] - app['start'] for app in cluster_appearances),
-                "first_appearance": min(app['start'] for app in cluster_appearances),
-                "last_appearance": max(app['end'] for app in cluster_appearances)
+                "bbox": representative_face.get('bbox', []),
+                "embedding": representative_face.get('embedding', []),
+                "appearances": appearances,
+                "total_appearances": len(appearances),
+                "total_screen_time": sum(app['end'] - app['start'] for app in appearances),
+                "first_appearance": appearances[0]['start'],
+                "last_appearance": appearances[-1]['end'],
+                "thumbnail": representative_face.get('thumbnail', '')
             }
-            summary.append(cluster_summary)
+            
+            all_faces_summary.append(face_summary)
+            
+            if i < top_k:
+                prominent_faces.append(face_summary)
     
-    return summary, face_embeddings, clustering.labels_
+    return all_faces_summary, prominent_faces, face_embeddings, clustering.labels_
 
-# Create comprehensive face summary
-face_summary, face_embeddings, face_labels = create_comprehensive_face_summary(face_index, face_metadata)
+# Usage in the main Streamlit app:
+all_faces_summary, prominent_faces, face_embeddings, face_labels = create_comprehensive_face_summary(face_index, face_metadata)
 
 # Face cluster visualization
-def plot_face_clusters_interactive(face_embeddings, labels, face_summary):
+# Update the face cluster visualization function
+def plot_face_clusters_interactive(face_embeddings, face_labels, all_faces_summary, prominent_faces):
     pca = PCA(n_components=3)
     embeddings_3d = pca.fit_transform(face_embeddings)
     
-    unique_labels = set(labels)
+    unique_labels = set(face_labels)
     colors = [f'rgb({int(r*255)},{int(g*255)},{int(b*255)})' 
               for r, g, b, _ in plt.cm.rainbow(np.linspace(0, 1, len(unique_labels)))]
     
@@ -102,20 +132,16 @@ def plot_face_clusters_interactive(face_embeddings, labels, face_summary):
     for label, color in zip(unique_labels, colors):
         if label == -1:
             continue  # Skip noise points
-        cluster_points = embeddings_3d[labels == label]
-        cluster_info = next((c for c in face_summary if c['cluster_id'] == f'cluster_{label}'), None)
         
-        if cluster_info:
-            hover_text = [
-                f"Cluster {label}<br>"
-                f"Face count: {cluster_info['face_count']}<br>"
-                f"Total screen time: {cluster_info['total_screen_time']:.2f}s<br>"
-                f"First appearance: {cluster_info['first_appearance']:.2f}s<br>"
-                f"Last appearance: {cluster_info['last_appearance']:.2f}s"
-                for _ in cluster_points
-            ]
-        else:
-            hover_text = [f"Cluster {label}" for _ in cluster_points]
+        cluster_points = embeddings_3d[face_labels == label]
+        cluster_faces = [face for face in all_faces_summary if face['cluster_id'] == f'cluster_{label}']
+        
+        hover_text = [
+            f"Cluster {label}<br>"
+            f"Time: {face['appearances'][0]['start']:.2f}s - {face['appearances'][-1]['end']:.2f}s<br>"
+            f"Appearances: {face['total_appearances']}"
+            for face in cluster_faces
+        ]
         
         trace = go.Scatter3d(
             x=cluster_points[:, 0],
@@ -129,6 +155,20 @@ def plot_face_clusters_interactive(face_embeddings, labels, face_summary):
         )
         traces.append(trace)
     
+    # Add markers for prominent faces
+    prominent_points = [embeddings_3d[face_labels == int(face['cluster_id'].split('_')[1])][0] for face in prominent_faces]
+    prominent_trace = go.Scatter3d(
+        x=[p[0] for p in prominent_points],
+        y=[p[1] for p in prominent_points],
+        z=[p[2] for p in prominent_points],
+        mode='markers',
+        name='Prominent Faces',
+        marker=dict(size=10, color='red', symbol='star'),
+        text=[f"Prominent Face<br>Cluster {face['cluster_id']}" for face in prominent_faces],
+        hoverinfo='text'
+    )
+    traces.append(prominent_trace)
+    
     layout = go.Layout(
         title='Face Clusters Visualization',
         scene=dict(xaxis_title='PCA 1', yaxis_title='PCA 2', zaxis_title='PCA 3'),
@@ -190,22 +230,25 @@ st.sidebar.text_area("Full Transcript", transcript_text, height=300)
 st.header("Video Summary")
 
 # Face Clusters
-st.subheader("Face Clusters")
-for cluster in face_summary[:5]:  # Display first 5 clusters
-    st.write(f"Cluster {cluster['cluster_id']}:")
-    st.write(f"  Face count: {cluster['face_count']}")
-    st.write(f"  Total screen time: {cluster['total_screen_time']:.2f} seconds")
-    st.write(f"  First appearance: {cluster['first_appearance']:.2f} seconds")
-    st.write(f"  Last appearance: {cluster['last_appearance']:.2f} seconds")
-    st.write(f"  Timeline: {len(cluster['timeline'])} appearances")
+st.subheader("Prominent Face Clusters")
+for face in prominent_faces:  # Use prominent_faces instead of face_summary
+    st.write(f"Face Cluster {face['cluster_id']}:")
+    st.write(f"  Total appearances: {face['total_appearances']}")
+    st.write(f"  Total screen time: {face['total_screen_time']:.2f} seconds")
+    st.write(f"  First appearance: {face['first_appearance']:.2f} seconds")
+    st.write(f"  Last appearance: {face['last_appearance']:.2f} seconds")
+    st.write(f"  Timeline: {len(face['appearances'])} appearances")
     st.write("  First 5 appearances:")
-    for app in cluster['timeline'][:5]:
+    for app in face['appearances'][:5]:
         st.write(f"    {app['start']:.2f}s - {app['end']:.2f}s")
+    if face['thumbnail']:
+        image = Image.open(io.BytesIO(base64.b64decode(face['thumbnail'])))
+        st.image(image, caption=f"Representative face for {face['cluster_id']}", width=100)
     st.write("---")
 
 # Face Cluster Visualization
 st.subheader("Face Cluster Visualization")
-fig = plot_face_clusters_interactive(face_embeddings, face_labels, face_summary)
+fig = plot_face_clusters_interactive(face_embeddings, face_labels, all_faces_summary, prominent_faces)
 st.plotly_chart(fig)
 
 # Themes