Update visualization.py

visualization.py

@@ -29,48 +29,34 @@ def create_bar_chart(data, title, speaker):
     fig.update_layout(title=f"{title} - Speaker {speaker}", xaxis_title="Traits", yaxis_title="Score")
     return fig
 
-def create_heatmap(data, title, speaker):
-    fig = go.Figure(data=go.Heatmap(
-        z=[[data[k] for k in data]],
-        x=list(data.keys()),
-        y=[speaker],
-        colorscale='Viridis'
-    ))
-    fig.update_layout(title=f"{title} - Speaker {speaker}")
-    return fig
-
-def update_visibility_and_charts(status, exec_time, lang, info, attachments, bigfive, personalities):
+def update_visibility_and_charts(status, exec_time, lang, attachments, bigfive, personalities):
     charts = []
     
-    speakers_data = extract_speaker_data(attachments + "\n" + bigfive + "\n" + personalities)
-    main_speakers = list(speakers_data.keys())[:2]  # Limit to 2 main speakers
-    
-    for speaker in main_speakers:
-        data = speakers_data[speaker]
-        
-        attachment_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"]}
-        if attachment_data:
-            charts.append(create_bar_chart(attachment_data, "Attachment Styles", speaker))
-        
-        bigfive_data = {k: v for k, v in data.items() if k in ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]}
-        if bigfive_data:
-            charts.append(create_bar_chart(bigfive_data, "Big Five Traits", speaker))
-        
-        personality_data = {k: v for k, v in data.items() if k in ["Depressed", "Paranoid", "Schizoid-Schizotypal", "Antisocial-Psychopathic", "Borderline-Dysregulated", "Hysteric-Histrionic", "Narcissistic", "Anxious-Avoidant", "Dependent-Victimized", "Obsessional"]}
-        if personality_data:
-            charts.append(create_bar_chart(personality_data, "Personality Traits", speaker))
-        
-        anxiety_avoidance_data = {k: v for k, v in data.items() if k in ["Anxiety", "Avoidance"]}
-        if anxiety_avoidance_data:
-            charts.append(create_heatmap(anxiety_avoidance_data, "Anxiety-Avoidance", speaker))
+    for analysis_text in [attachments, bigfive, personalities]:
+        speakers_data = extract_speaker_data(analysis_text)
+        if not speakers_data:
+            print(f"No speaker data extracted from: {analysis_text}")
+
+        # Determine the two main speakers
+        speaker_counts = Counter(speakers_data.keys())
+        main_speakers = [speaker for speaker, count in speaker_counts.most_common(2)]
         
-        self_others_data = {k: v for k, v in data.items() if k in ["Self", "Others"]}
-        if self_others_data:
-            charts.append(create_heatmap(self_others_data, "Self-Others", speaker))
+        for speaker in main_speakers:
+            data = speakers_data.get(speaker, {})
+            attachment_data = {k: v for k, v in data.items() if k in ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"]}
+            if attachment_data:
+                charts.append(create_bar_chart(attachment_data, "Attachment Styles", speaker))
+            
+            bigfive_data = {k: v for k, v in data.items() if k in ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"]}
+            if bigfive_data:
+                charts.append(create_bar_chart(bigfive_data, "Big Five Traits", speaker))
+            
+            personality_data = {k: v for k, v in data.items() if k in ["Depressed", "Paranoid", "Schizoid-Schizotypal", "Antisocial-Psychopathic", "Borderline-Dysregulated", "Hysteric-Histrionic", "Narcissistic", "Anxious-Avoidant", "Dependent-Victimized", "Obsessional"]}
+            if personality_data:
+                charts.append(create_bar_chart(personality_data, "Personality Traits", speaker))
 
     return [
         gr.update(value=status, visible=True),
         gr.update(value=exec_time, visible=True),
         gr.update(value=lang, visible=True),
-        gr.update(value=info, visible=True),
     ] + charts