Update app.py

app.py

@@ -7,6 +7,7 @@ import io
 from PIL import Image
 import pickle
 import requests
+import cv2
 
 hf_token = {
     "clipqwentimer": os.environ["HF_CLIPQwenTimer_Token"],
@@ -95,6 +96,72 @@ def load_csv(example_index, file):
     else:
         return gr.skip()
 
+def vision_attention_rollout(attentions, start_layer=0, end_layer=12):
+    seq_len = attentions.shape[-1]
+    result = np.eye(seq_len)
+
+    for attn in attentions[start_layer:end_layer]:
+        attn_heads = attn.mean(axis=0)
+        attn_aug = attn_heads + np.eye(seq_len)
+        attn_aug = attn_aug / attn_aug.sum(axis=-1, keepdims=True)
+        result = attn_aug @ result
+
+    return result[0, -49:]
+
+def plot_vision_heatmap(image, rollout_attention, alpha=0.5, cmap='jet'):
+    num_patches = rollout_attention.shape[0]
+    grid_size = int(np.sqrt(num_patches))
+    
+    attn_grid = rollout_attention.reshape(grid_size, grid_size)
+    
+    H, W = image.shape[:2]
+    attn_map = cv2.resize(attn_grid, (W, H), interpolation=cv2.INTER_CUBIC)
+    attn_map = attn_map / attn_map.max()
+    
+    plt.figure(figsize=(6,6))
+    plt.imshow(image)
+    plt.imshow(attn_map, cmap=cmap, alpha=alpha)
+    plt.axis('off')
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    plot_img = Image.open(buf).convert('RGB')
+    plt.clf()
+
+    return plot_img
+
+def time_series_attention_sum(attentions, context_length, start_layer=0, end_layer=12):
+    import math
+    seq_len = attentions.shape[-1]
+    result = np.zeros(seq_len)
+    for attn in attentions[start_layer:end_layer]:
+        attn_heads = attn.mean(0).squeeze()
+        result += attn_heads
+    att_len = math.ceil(context_length/16)
+    return result[-att_len:]
+
+def plot_time_series_heatmap(context, attention, time_steps):    
+    fig, ax1 = plt.subplots(figsize=(8, 4))
+
+    attention = attention/attention.max()
+    cmap = plt.get_cmap("coolwarm")
+    colors = cmap(attention)
+    colors[:, -1] = attention
+    ax1.bar([16*i for i in range(len(attention))], attention, width=[16 if time_steps-16*(i+1)>0 else time_steps-16*i for i in range(len(attention))], align="edge", color=colors)
+    ax1.set_yticks([])
+    
+    ax2 = ax1.twinx()
+    ax2.plot(context, color="black", linewidth=2)
+    ax2.yaxis.set_ticks_position("left")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    plot_img = Image.open(buf).convert('RGB')
+    plt.clf()
+
+    return plot_img
+
 def predict(dataset, text, example_index, file, vision_encoder, text_encoder, tsfm):
     if (dataset is None or example_index is None) or (example_index == -1 and file is None):
         return (
@@ -138,10 +205,8 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder, ts
     res = requests.post(url, headers=headers, json=payload)
     res_json = res.json()
 
+    # Forecast Plot
     prediction = np.array(res_json['prediction'])
-    vision_attentions = np.array(res_json['vision_attentions'])
-    time_series_attentions = np.array(res_json['time_series_attentions'])
-
     cl = context_length[dataset]
     prediction = prediction[:cl]
     prediction = prediction*std+mean
@@ -151,19 +216,56 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder, ts
     start_time = input_dates_series.iloc[-1] + time_diff
     forecast_dates_series = pd.date_range(start=start_time, periods=len(input_dates_series), freq=time_diff)
 
-    plt.style.use("seaborn-v0_8")
-    fig, ax = plt.subplots()
-    ax.plot(input_dates_series, time_series, color="black", alpha=0.7, linewidth=3, label='Input')
-    ax.plot(forecast_dates_series, prediction, color='C2', alpha=0.7, linewidth=3, label='Forecast')
-    if example_index == -1: # Custom Input
-        true = df["Ground Truth"]
-    else:
-        true = targets[dataset][example_index].iloc[:, -1]
-    if len(true) == context_length[dataset]:
-        ax.plot(forecast_dates_series, true, color='C0', alpha=0.7, linewidth=3, label='Ground Truth')
-    ax.legend()
+    plt.close()
+    with plt.style.context("seaborn-v0_8"):
+        fig, ax = plt.subplots(figsize=(10,4))
+        ax.plot(input_dates_series, time_series, color="black", alpha=0.7, linewidth=3, label='Input')
+        ax.plot(forecast_dates_series, prediction, color='C2', alpha=0.7, linewidth=3, label='Forecast')
+        if example_index == -1: # Custom Input
+            true = df["Ground Truth"]
+        else:
+            true = targets[dataset][example_index].iloc[:, -1]
+        if len(true) == context_length[dataset]:
+            ax.plot(forecast_dates_series, true, color='C0', alpha=0.7, linewidth=3, label='Ground Truth')
+        ax.legend()
+
+    # Vision Heatmap
+    plt.figure(figsize=(384/100, 384/100), dpi=100)
+    plt.plot(time_series_normalized, color="black", linestyle="-", linewidth=1, marker="*", markersize=1)
+    plt.xticks([])
+    plt.yticks([])
+    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+    plt.margins(0,0)
 
-    return gr.Markdown(visible=False), fig
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    context_image = np.array(Image.open(buf).convert('RGB'))
+
+    vision_attentions = np.array(res_json['vision_attentions'])
+    vision_heatmap_gallery_items = []
+    for i in range(0, 12, 3):
+        vis_attn = vision_attention_rollout(vision_attentions, i, i+3)
+        vision_heatmap = plot_vision_heatmap(context_image, vis_attn)
+        vision_heatmap_gallery_items.append((vision_heatmap, f"Heatmap from Layer{i}:{i+3}"))
+
+    # Time Series Heatmap
+    if tsfm == "Chronos":
+        time_series_attentions = np.array(res_json['time_series_attentions'])
+        time_series_heatmap_gallery_items = []
+        for i in range(0, 12, 3):
+            ts_attn = time_series_attention_sum(time_series_attentions, cl, i, i+3)
+            time_series_heatmap = plot_time_series_heatmap(time_series, ts_attn, cl)
+            time_series_heatmap_gallery_items.append((time_series_heatmap, f"Heatmap from Layer{i}:{i+3}"))
+    else:
+        time_series_heatmap_gallery_items = None
+    
+    return (
+        gr.Markdown(visible=False),
+        fig,
+        gr.Gallery(vision_heatmap_gallery_items, interactive=False, height="350px", object_fit="contain", visible=True),
+        gr.Gallery(time_series_heatmap_gallery_items, interactive=False, height="350px", object_fit="contain", visible=True if time_series_heatmap_gallery_items else False)
+    )
 
 def add_example_gallery(dataset, gallery, example_index, file):
     if example_index == -1 and file:
@@ -245,8 +347,10 @@ with gr.Blocks() as demo:
             btn = gr.Button("Run")
         with gr.Column(scale=2):
             forecast_plot = gr.Plot(label="Forecast", format="png")
+            vision_heatmap_gallery = gr.Gallery(visible=False)
+            time_series_heatmap_gallery = gr.Gallery(visible=False)
     
-        btn.click(predict, inputs=[dataset_dropdown, dataset_description_textbox, example_index, time_series_file, vision_encoder_radio, text_encoder_radio, tsfm_radio], outputs=[warning_markdown, forecast_plot])
+        btn.click(predict, inputs=[dataset_dropdown, dataset_description_textbox, example_index, time_series_file, vision_encoder_radio, text_encoder_radio, tsfm_radio], outputs=[warning_markdown, forecast_plot, vision_heatmap_gallery, time_series_heatmap_gallery])
         btn.click(add_example_gallery, inputs=[dataset_dropdown, example_gallery, example_index, time_series_file], outputs=[example_gallery])
 
 if __name__ == "__main__":