update app.py

app.py

@@ -1,11 +1,14 @@
 # app.py
 import os
+os.environ["GRADIO_TEMP_DIR"] = "/home/mouxiangchen/VisionTSpp/gradio_tmp"
+
 import gradio as gr
 import torch
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
 import einops
+import copy
 
 from huggingface_hub import snapshot_download
 from visionts import VisionTSpp, freq_to_seasonality_list
@@ -31,6 +34,7 @@ model = VisionTSpp(
     ARCH,
     ckpt_path=CKPT_PATH,
     # quantiles=QUANTILES,
+    quantile=True,
     clip_input=True,
     complete_no_clip=False,
     color=True
@@ -48,8 +52,8 @@ imagenet_std = np.array([0.229, 0.224, 0.225])
 # This dictionary maps user-friendly names to local file paths
 # ASSUMPTION: These files exist in a 'datasets' subfolder
 
-# data_dir = "./datasets/"
-data_dir = "./"
+data_dir = "./datasets/"
+# data_dir = "./"
 PRESET_DATASETS = {
     "ETTm1": data_dir + "ETTm1.csv",
     "ETTm2": data_dir + "ETTm2.csv",
@@ -71,27 +75,40 @@ def load_preset_data(name):
 # 3. Visualization Functions (No changes needed)
 # ========================
 def show_image_tensor(image_tensor, title='', cur_nvars=1, cur_color_list=None):
-    if image_tensor is None: return None
-    image = image_tensor.permute(1, 2, 0).cpu()
+    if image_tensor is None: 
+        return None
+    
+    # no need for permute?
+    # image = image_tensor.permute(1, 2, 0).cpu()
+    image = image_tensor.cpu()
+    
     cur_image = torch.zeros_like(image)
+
     height_per_var = image.shape[0] // cur_nvars
     for i in range(cur_nvars):
         cur_color_idx = cur_color_list[i]
         var_slice = image[i*height_per_var:(i+1)*height_per_var, :, :]
         unnormalized_channel = var_slice[:, :, cur_color_idx] * imagenet_std[cur_color_idx] + imagenet_mean[cur_color_idx]
         cur_image[i*height_per_var:(i+1)*height_per_var, :, cur_color_idx] = unnormalized_channel * 255
+    
     cur_image = torch.clamp(cur_image, 0, 255).int().numpy()
+
     fig, ax = plt.subplots(figsize=(6, 6))
     ax.imshow(cur_image)
     ax.set_title(title, fontsize=14)
     ax.axis('off')
+    
     plt.tight_layout()
     plt.close(fig)
+    
     return fig
 
 def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_quantiles, context_len, pred_len):
-    if isinstance(true_data, torch.Tensor): true_data = true_data.cpu().numpy()
-    if isinstance(pred_median, torch.Tensor): pred_median = pred_median.cpu().numpy()
+    if isinstance(true_data, torch.Tensor): 
+        true_data = true_data.cpu().numpy()
+    if isinstance(pred_median, torch.Tensor): 
+        pred_median = pred_median.cpu().numpy()
+    
     for i, q in enumerate(pred_quantiles_list):
         if isinstance(q, torch.Tensor):
             pred_quantiles_list[i] = q.cpu().numpy()
@@ -99,11 +116,17 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
     nvars = true_data.shape[1]
     FIG_WIDTH, FIG_HEIGHT_PER_VAR = 15, 2.0
     fig, axes = plt.subplots(nvars, 1, figsize=(FIG_WIDTH, nvars * FIG_HEIGHT_PER_VAR), sharex=True)
-    if nvars == 1: axes = [axes]
+    if nvars == 1: 
+        axes = [axes]
+
+    print(f"{len(pred_quantiles_list) = }")
+    print(f"{len(model_quantiles) = }")
+    print(f"{pred_quantiles_list[0].shape = }")
 
     sorted_quantiles = sorted(zip(model_quantiles, pred_quantiles_list + [pred_median]), key=lambda x: x[0])
     quantile_preds = [item[1] for item in sorted_quantiles if item[0] != 0.5]
     quantile_vals = [item[0] for item in sorted_quantiles if item[0] != 0.5]
+    
     num_bands = len(quantile_preds) // 2
     quantile_colors = plt.cm.Blues(np.linspace(0.3, 0.8, num_bands))[::-1]
 
@@ -128,6 +151,7 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
     fig.legend(unique_labels.values(), unique_labels.keys(), loc='upper center', bbox_to_anchor=(0.5, 1.05), ncol=num_bands + 2)
     plt.tight_layout(rect=[0, 0, 1, 0.95])
     plt.close(fig)
+    
     return fig
 
 
@@ -159,6 +183,7 @@ def predict_at_index(df, index, context_len, pred_len):
             inferred_freq = pd.tseries.frequencies.to_offset(time_diff).freqstr
             gr.Warning(f"Could not reliably infer frequency. Using fallback based on first two timestamps: {inferred_freq}")
         print(f"Inferred frequency: {inferred_freq}")
+    
     except Exception as e:
         raise gr.Error(f"❌ Date processing failed: {e}. Please check the date format (e.g., YYYY-MM-DD HH:MM:SS).")
 
@@ -167,7 +192,7 @@ def predict_at_index(df, index, context_len, pred_len):
 
     total_samples = len(data) - context_len - pred_len + 1
     if total_samples <= 0:
-        raise gr.Error(f"Data is too short. It needs at least {context_len + pred_len} rows, but has {len(data)}.")
+        raise gr.Error(f"Data is too short. It needs at least context_len + pred_len = {context_len + pred_len} rows, but has {len(data)}.")
     
     index = max(0, min(index, total_samples - 1))
 
@@ -184,6 +209,7 @@ def predict_at_index(df, index, context_len, pred_len):
     # *** Use inferred frequency ***
     periodicity_list = freq_to_seasonality_list(inferred_freq)
     periodicity = periodicity_list[0] if periodicity_list else 1
+    
     color_list = [i % 3 for i in range(nvars)]
     model.update_config(context_len=context_len, pred_len=pred_len, periodicity=periodicity)
 
@@ -191,10 +217,39 @@ def predict_at_index(df, index, context_len, pred_len):
         y_pred, input_image, reconstructed_image, _, _ = model.forward(
             x_tensor, export_image=True, color_list=color_list
         )
-        all_preds = dict(zip(model.quantiles, y_pred))
+        y_pred, y_pred_quantile_list = y_pred
+
+        print(f"{x_tensor.shape = }")
+        print(f"{y_pred.shape = }")
+        print(f"{input_image.shape = }")
+        print(f"{reconstructed_image.shape = }")
+        print(f"{len(y_pred_quantile_list) = }")
+
+        print(f"{input_image = }")
+        print(f"{input_image[0,0,0, :, 0] = }")
+        print(f"{input_image[0,0,0, 50:70, 0] = }")
+        print(f"{input_image[0,0,0, 100:120, 0] = }")
+        # print(f"{input_image[0] = }")
+        # print(f"{reconstructed_image = }")
+
+        all_y_pred_list = copy.deepcopy(y_pred_quantile_list)
+        
+        # insert in the place of 0.5 quantile, ie:len(QUANTILES)//2
+        all_y_pred_list.insert(len(QUANTILES)//2, y_pred)
+
+        print(f"{len(all_y_pred_list) = }")
+        print(f"{all_y_pred_list[0].shape = }")
+
+        all_preds = dict(zip(QUANTILES, all_y_pred_list))
+
+        print(f"{all_preds.keys() = }")
+
         pred_median_norm = all_preds.pop(0.5)[0]
         pred_quantiles_norm = [q[0] for q in list(all_preds.values())]
 
+        print(f"{pred_median_norm.shape = }")
+        print(f"{len(pred_quantiles_norm) = }")
+
     y_true = y_true_norm * x_std + x_mean
     pred_median = pred_median_norm.cpu().numpy() * x_std + x_mean
     pred_quantiles = [q.cpu().numpy() * x_std + x_mean for q in pred_quantiles_norm]
@@ -251,6 +306,7 @@ def run_forecast(data_source, upload_file, index, context_len, pred_len):
             gr.update(maximum=result.total_samples - 1, value=final_index),
             gr.update(value=result.inferred_freq) # *** Update frequency textbox ***
         )
+    
     except Exception as e:
         error_fig = plt.figure(figsize=(10, 5))
         plt.text(0.5, 0.5, f"An error occurred:\n{str(e)}", ha='center', va='center', wrap=True, color='red', fontsize=12)