REFAC: Improve emotion vector creation and add ECI timeline calculation in Streamlit app

src/streamlit_app.py

@@ -52,13 +52,16 @@ def load_models():
 whisper_model, text_classifier, ser_model, ser_feature_extractor = load_models()
 
 # --- Helper Functions for Analysis ---
-def create_unified_vector(scores_dict):
+def create_unified_vector(scores_dict, mapping_dict):
     vector = np.zeros(len(UNIFIED_EMOTIONS))
-    total_score = sum(scores_dict.values())
+    total_score = 0
+    for label, score in scores_dict.items():
+        unified_label = mapping_dict.get(label)
+        if unified_label in UNIFIED_EMOTIONS:
+            vector[UNIFIED_EMOTIONS.index(unified_label)] += score
+            total_score += score
     if total_score > 0:
-        for label, score in scores_dict.items():
-            if label in UNIFIED_EMOTIONS:
-                vector[UNIFIED_EMOTIONS.index(label)] = score / total_score
+        vector /= total_score
     return vector
 
 def get_consistency_level(cosine_sim):
@@ -138,25 +141,38 @@ if uploaded_file is not None:
             
             st.header("Analysis Results")
             
-            def process_and_get_dominant(timeline, mapping):
-                if not timeline: return "N/A", {}
+            def process_timeline_to_df(timeline, mapping):
+                if not timeline: return pd.DataFrame(columns=UNIFIED_EMOTIONS)
                 df = pd.DataFrame.from_dict(timeline, orient='index')
-                unified_scores = {e: 0.0 for e in UNIFIED_EMOTIONS}
-                for raw_label, scores in df.items():
-                    unified_label = mapping.get(raw_label)
-                    if unified_label:
-                        unified_scores[unified_label] += scores.mean()
-                if sum(unified_scores.values()) == 0: return "N/A", {}
-                dominant_emotion = max(unified_scores, key=unified_scores.get)
-                return dominant_emotion.capitalize(), unified_scores
-
-            dominant_fer, fer_avg_scores = process_and_get_dominant(fer_timeline, FACIAL_TO_UNIFIED)
-            dominant_ser, ser_avg_scores = process_and_get_dominant(ser_timeline, SER_TO_UNIFIED)
-            dominant_text, ter_avg_scores = process_and_get_dominant(ter_timeline, TEXT_TO_UNIFIED)
-
-            fer_vector = create_unified_vector(fer_avg_scores)
-            ser_vector = create_unified_vector(ser_avg_scores)
-            text_vector = create_unified_vector(ter_avg_scores)
+                df_unified = pd.DataFrame(index=df.index, columns=UNIFIED_EMOTIONS).fillna(0.0)
+                for raw_col in df.columns:
+                    unified_col = mapping.get(raw_col)
+                    if unified_col:
+                        df_unified[unified_col] += df[raw_col]
+                return df_unified
+
+            fer_df = process_timeline_to_df(fer_timeline, FACIAL_TO_UNIFIED)
+            ser_df = process_timeline_to_df(ser_timeline, SER_TO_UNIFIED)
+            ter_df = process_timeline_to_df(ter_timeline, TEXT_TO_UNIFIED)
+
+            def get_dominant_emotion_from_df(df):
+                if df.empty or df.sum().sum() == 0: return "N/A"
+                return df.sum().idxmax().capitalize()
+
+            dominant_fer = get_dominant_emotion_from_df(fer_df)
+            dominant_ser = get_dominant_emotion_from_df(ser_df)
+            dominant_text = get_dominant_emotion_from_df(ter_df)
+
+            def get_avg_unified_scores(df):
+                return df.mean().to_dict() if not df.empty else {}
+
+            fer_avg_scores = get_avg_unified_scores(fer_df)
+            ser_avg_scores = get_avg_unified_scores(ser_df)
+            ter_avg_scores = get_avg_unified_scores(ter_df)
+
+            fer_vector = create_unified_vector(fer_avg_scores, {e:e for e in UNIFIED_EMOTIONS})
+            ser_vector = create_unified_vector(ser_avg_scores, {e:e for e in UNIFIED_EMOTIONS})
+            text_vector = create_unified_vector(ter_avg_scores, {e:e for e in UNIFIED_EMOTIONS})
             
             similarities = [cosine_similarity([fer_vector], [text_vector])[0][0], cosine_similarity([fer_vector], [ser_vector])[0][0], cosine_similarity([ser_vector], [text_vector])[0][0]]
             avg_similarity = np.nanmean([s for s in similarities if not np.isnan(s)])
@@ -179,23 +195,27 @@ if uploaded_file is not None:
             with col2:
                 st.subheader("Unified Emotion Timeline")
                 
-                def create_timeline_df(timeline, mapping):
-                    if not timeline: return pd.DataFrame(columns=UNIFIED_EMOTIONS)
-                    df = pd.DataFrame.from_dict(timeline, orient='index')
-                    df_unified = pd.DataFrame(index=df.index, columns=UNIFIED_EMOTIONS).fillna(0.0)
-                    for raw_col in df.columns:
-                        unified_col = mapping.get(raw_col)
-                        if unified_col:
-                            df_unified[unified_col] += df[raw_col]
-                    return df_unified
-
-                fer_df = create_timeline_df(fer_timeline, FACIAL_TO_UNIFIED)
-                ser_df = create_timeline_df(ser_timeline, SER_TO_UNIFIED)
-                ter_df = create_timeline_df(ter_timeline, TEXT_TO_UNIFIED)
-
                 full_index = np.arange(0, duration, 0.5)
                 combined_df = pd.DataFrame(index=full_index)
 
+                # --- NEW: ECI Timeline Calculation ---
+                eci_timeline = {}
+                for t_stamp in full_index:
+                    vectors = []
+                    # Get interpolated facial vector if available
+                    if not fer_df.empty and t_stamp in fer_df.index:
+                        vectors.append(create_unified_vector(fer_df.loc[t_stamp].to_dict(), {e:e for e in UNIFIED_EMOTIONS}))
+                    # Get speech vector if available
+                    if not ser_df.empty and int(t_stamp) in ser_df.index:
+                        vectors.append(create_unified_vector(ser_df.loc[int(t_stamp)].to_dict(), {e:e for e in UNIFIED_EMOTIONS}))
+                    # Get text vector if available
+                    if not ter_df.empty and int(t_stamp) in ter_df.index:
+                        vectors.append(create_unified_vector(ter_df.loc[int(t_stamp)].to_dict(), {e:e for e in UNIFIED_EMOTIONS}))
+                    
+                    if len(vectors) >= 2:
+                        sims = [cosine_similarity([v1], [v2])[0][0] for i, v1 in enumerate(vectors) for v2 in vectors[i+1:]]
+                        eci_timeline[t_stamp] = np.mean(sims)
+
                 if not fer_df.empty:
                     fer_df_resampled = fer_df.reindex(fer_df.index.union(full_index)).interpolate(method='linear').reindex(full_index)
                     for e in UNIFIED_EMOTIONS: combined_df[f'Facial_{e}'] = fer_df_resampled.get(e, 0.0)
@@ -207,6 +227,11 @@ if uploaded_file is not None:
                 if not ter_df.empty:
                     ter_df_resampled = ter_df.reindex(ter_df.index.union(full_index)).interpolate(method='linear').reindex(full_index)
                     for e in UNIFIED_EMOTIONS: combined_df[f'Text_{e}'] = ter_df_resampled.get(e, 0.0)
+                
+                # Add ECI timeline to the DataFrame
+                if eci_timeline:
+                    eci_df = pd.Series(eci_timeline, name="ECI")
+                    combined_df['ECI'] = eci_df
 
                 combined_df.fillna(0, inplace=True)
                 
@@ -216,9 +241,14 @@ if uploaded_file is not None:
                     styles = {'Facial': '-', 'Speech': '--', 'Text': ':'}
 
                     for col in combined_df.columns:
+                        if col == 'ECI': continue
                         modality, emotion = col.split('_')
                         if emotion in colors:
-                            ax.plot(combined_df.index, combined_df[col], label=f'{modality} {emotion.capitalize()}', color=colors[emotion], linestyle=styles[modality], alpha=0.8)
+                            ax.plot(combined_df.index, combined_df[col], label=f'{modality} {emotion.capitalize()}', color=colors[emotion], linestyle=styles[modality], alpha=0.7)
+                    
+                    # Plot ECI on the same axis
+                    if 'ECI' in combined_df.columns:
+                        ax.plot(combined_df.index, combined_df['ECI'], label='Emotion Consistency', color='black', linewidth=2.5, alpha=0.9)
 
                     ax.set_title("Emotion Confidence Over Time (Normalized)")
                     ax.set_xlabel("Time (seconds)")