color coding logic for heatmap

app.py

@@ -100,33 +100,47 @@ def prepare_heatmap_data(data):
     return heatmap_data
 
 
-def plot_emotion_heatmap(heatmap_data):
-    heatmap_data = heatmap_data.T
-
-    fig = plt.figure(figsize=(len(heatmap_data.columns) * 0.5 + 4, len(heatmap_data.index) * 0.5 + 2))
-
-    cmap = LinearSegmentedColormap.from_list("white_to_grey", ["#ffffff", "#aaaaaa"])
-
-    sns.heatmap(
-        heatmap_data,
-        annot=False,
-        cmap=cmap,
-        cbar=True,
-        linewidths=0.5,
-        linecolor='gray',
-        vmin=0,
-        vmax=1
-    )
-
-    plt.xlabel("Emotions")
-    plt.ylabel("Sentences")
-    plt.xticks(rotation=0, ha='center')
-    plt.yticks(rotation=0)
+def plot_emotion_heatmap_colored(heatmap_data, emotion_colors):
+    heatmap_data = heatmap_data.T  # Transpose to: rows = emotions, cols = sentences
+
+    # Normalize all values to [0, 1] for each emotion
+    normalized_data = heatmap_data.copy()
+    for row in heatmap_data.index:
+        max_val = heatmap_data.loc[row].max()
+        if max_val > 0:
+            normalized_data.loc[row] = heatmap_data.loc[row] / max_val
+        else:
+            normalized_data.loc[row] = 0
+
+    fig, ax = plt.subplots(figsize=(len(heatmap_data.columns) * 0.5 + 4, len(heatmap_data.index) * 0.5 + 2))
+
+    for i, emotion in enumerate(heatmap_data.index):
+        # Create custom colormap for each row
+        base_color = emotion_colors[emotion]
+        cmap = LinearSegmentedColormap.from_list(f"{emotion}_map", ["#ffffff", base_color])
+
+        # Create heatmap for one row at a time
+        sns.heatmap(
+            pd.DataFrame([normalized_data.loc[emotion].values], index=[emotion], columns=normalized_data.columns),
+            cmap=cmap,
+            cbar=False,
+            linewidths=0.5,
+            linecolor='gray',
+            vmin=0,
+            vmax=1,
+            ax=ax if i == 0 else ax,  # reuse same axis
+        )
+
+    # Format axis
+    ax.set_xticklabels(heatmap_data.columns, rotation=0, ha='center')
+    ax.set_yticklabels(heatmap_data.index, rotation=0, ha='right')
+
+    ax.set_xlabel("Sentences")
+    ax.set_ylabel("Emotions")
     plt.tight_layout()
 
     return fig
 
-
 def plot_average_emotion_pie(heatmap_data):
     all_emotion_scores = np.array([item['emotions'] for item in heatmap_data])
     mean_scores = all_emotion_scores.mean(axis=0)