Create visualization.py

visualization.py

@@ -0,0 +1,76 @@
+import gradio as gr
+import plotly.graph_objs as go
+import re
+from collections import Counter
+
+def extract_speaker_data(text):
+    speakers = {}
+    current_speaker = None
+    for line in text.split('\n'):
+        if line.lower().startswith("speaker"):
+            current_speaker = line.split(":")[1].strip() if ":" in line else line.split()[1]
+            speakers[current_speaker] = {}
+        elif current_speaker and ":" in line:
+            key, value = line.split(":", 1)
+            try:
+                speakers[current_speaker][key.strip()] = float(value.strip())
+            except ValueError:
+                match = re.search(r"[-+]?\d*\.\d+|\d+", value)
+                if match:
+                    speakers[current_speaker][key.strip()] = float(match.group())
+    return speakers
+
+def create_bar_chart(data, title, speaker):
+    fig = go.Figure(data=[go.Bar(
+        x=list(data.keys()), 
+        y=list(data.values()),
+        marker_color=['red', 'green', 'blue', 'yellow', 'purple', 'orange', 'pink', 'cyan', 'magenta', 'brown'][:len(data)]
+    )])
+    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):
+    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))
+        
+        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))
+
+    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