giving a try to crazy imshow thingie

app.py

@@ -9,6 +9,7 @@ from transformers import AutoTokenizer
 import gradio as gr
 import matplotlib.pyplot as plt
 from matplotlib.colors import LinearSegmentedColormap
+import matplotlib.colors as mcolors
 import plotly.express as px
 import seaborn as sns
 
@@ -99,44 +100,31 @@ def prepare_heatmap_data(data):
     heatmap_data.columns = [item["sentence"][:18]+"..." for item in data]
     return heatmap_data
 
-
 def plot_emotion_heatmap(heatmap_data):
-    # heatmap_data = heatmap_data.T  # Transpose to: rows = emotions, cols = sentences
-
-    # Normalize all values to [0, 1] for each emotion
+    # Normalize values to [0, 1] per row (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))
-
-    print(heatmap_data) # debug
-    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_xticks(np.arange(len(heatmap_data.columns)) + 0.5)
-    ax.set_xticklabels(heatmap_data.columns, rotation=0, ha='center')
-    ax.set_yticks(np.arange(len(heatmap_data.index)) + 0.5)
-    ax.set_yticklabels(heatmap_data.index, rotation=0, ha='right')
+        normalized_data.loc[row] = heatmap_data.loc[row] / max_val if max_val > 0 else 0
+
+    # Build custom RGB color matrix
+    color_matrix = np.empty((len(normalized_data.index), len(normalized_data.columns), 3))
+    for i, emotion in enumerate(normalized_data.index):
+        base_rgb = mcolors.to_rgb(emotion_colors[emotion])
+        for j, val in enumerate(normalized_data.loc[emotion]):
+            # Linear interpolation from white to base color
+            color = tuple(1 - val * (1 - c) for c in base_rgb)
+            color_matrix[i, j] = color
+
+    fig, ax = plt.subplots(figsize=(len(normalized_data.columns) * 0.5 + 4, len(normalized_data.index) * 0.5 + 2))
+
+    ax.imshow(color_matrix, aspect='auto')
+
+    # Ticks and labels
+    ax.set_xticks(np.arange(len(normalized_data.columns)))
+    ax.set_xticklabels(normalized_data.columns, rotation=0, ha='center')
+    ax.set_yticks(np.arange(len(normalized_data.index)))
+    ax.set_yticklabels(normalized_data.index, rotation=0, ha='right')
 
     ax.set_xlabel("Sentences")
     ax.set_ylabel("Emotions")