update app.py, add choice button for VisionTSpp base and large

app.py

@@ -53,7 +53,7 @@ def cleanup_old_sessions():
         print("Cleanup complete. No old sessions found.")
 
 
-# --- NEW: Function to run the cleanup periodically in the background ---
+# --- Function to run the cleanup periodically in the background ---
 def periodic_cleanup_task():
     """Wrapper function to run cleanup in a loop with a sleep interval."""
     print("Starting background thread for periodic cleanup.")
@@ -96,18 +96,26 @@ for model_size, config in MODEL_CONFIGS.items():
 
 QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
 
-# --- NEW: Global variables to hold the currently loaded model and its size ---
+
+# Image normalization constants
+imagenet_mean = np.array([0.485, 0.456, 0.406])
+imagenet_std = np.array([0.229, 0.224, 0.225])
+
+
+# --- Global variables to hold the currently loaded model and its size ---
 CURRENT_MODEL_SIZE = None
 CURRENT_MODEL = None
 
 def load_model_for_size(model_size: str):
     """Loads the specified VisionTS++ model."""
     global CURRENT_MODEL, CURRENT_MODEL_SIZE
+
     if model_size not in MODEL_CONFIGS:
         raise ValueError(f"Invalid model size: {model_size}. Available: {list(MODEL_CONFIGS.keys())}")
 
     config = MODEL_CONFIGS[model_size]
     print(f"Loading {model_size} model...")
+
     model = VisionTSpp(
         config["arch"],
         ckpt_path=config["ckpt_path"],
@@ -126,6 +134,7 @@ def load_model_for_size(model_size: str):
 
     CURRENT_MODEL = model
     CURRENT_MODEL_SIZE = model_size
+
     return model
 
 # Load the default model (base) on startup
@@ -160,34 +169,44 @@ def load_preset_data(name):
 def show_image_tensor(image_tensor, title='', cur_nvars=1, cur_color_list=None):
     if image_tensor is None:
         return None
+    
     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()
 
     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]
 
     pred_quantiles_list.insert(len(QUANTILES)//2, pred_median)
     sorted_quantiles = sorted(zip(QUANTILES, pred_quantiles_list), key=lambda x: x[0])
@@ -200,10 +219,12 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
         ax.plot(true_data[:, i], label='Ground Truth', color='black', linewidth=1.5)
         pred_range = np.arange(context_len, context_len + pred_len)
         ax.plot(pred_range, pred_median[:, i], label='Prediction (Median)', color='red', linewidth=1.5)
+
         for j in range(num_bands):
             lower_quantile_pred, upper_quantile_pred = quantile_preds[j][:, i], quantile_preds[-(j+1)][:, i]
             q_low, q_high = quantile_vals[j], quantile_vals[-(j+1)]
             ax.fill_between(pred_range, lower_quantile_pred, upper_quantile_pred, color=quantile_colors[j], alpha=0.7, label=f'{int(q_low*100)}-{int(q_high*100)}% Quantile')
+        
         y_min, y_max = ax.get_ylim()
         ax.vlines(x=context_len, ymin=y_min, ymax=y_max, colors='gray', linestyles='--', alpha=0.7)
         ax.set_ylabel(f'Var {i+1}', rotation=0, labelpad=30, ha='right', va='center')
@@ -212,9 +233,11 @@ def visual_ts_with_quantiles(true_data, pred_median, pred_quantiles_list, model_
 
     handles, labels = axes[0].get_legend_handles_labels()
     unique_labels = dict(zip(labels, handles))
+    
     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
 
 
@@ -237,11 +260,13 @@ def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
     try:
         df['date'] = pd.to_datetime(df['date'])
         df = df.sort_values('date').set_index('date')
+
         inferred_freq = pd.infer_freq(df.index)
         if inferred_freq is None:
             time_diff = df.index[1] - df.index[0]
             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).")
@@ -270,7 +295,7 @@ def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
 
     color_list = [i % 3 for i in range(nvars)]
 
-    # --- NEW: Load the requested model if it's not the current one ---
+    # --- Load the requested model if it's not the current one ---
     if CURRENT_MODEL_SIZE != model_size:
         print(f"Switching model from {CURRENT_MODEL_SIZE} to {model_size}")
         load_model_for_size(model_size)
@@ -289,6 +314,7 @@ def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
         all_y_pred_list = copy.deepcopy(y_pred_quantile_list)
         all_y_pred_list.insert(len(QUANTILES)//2, y_pred)
         all_preds = dict(zip(QUANTILES, all_y_pred_list))
+
         pred_median_norm = all_preds.pop(0.5)[0]
         pred_quantiles_norm = [q[0] for q in list(all_preds.values())]
 
@@ -304,15 +330,18 @@ def predict_at_index(df, index, context_len, pred_len, session_dir, model_size):
         pred_quantiles_list=pred_quantiles, model_quantiles=list(all_preds.keys()),
         context_len=context_len, pred_len=pred_len
     )
+
     input_img_fig = show_image_tensor(input_image[0, 0], f'Input Image (Sample {index})', nvars, color_list)
     recon_img_fig = show_image_tensor(reconstructed_image[0, 0], 'Reconstructed Image', nvars, color_list)
 
     csv_path = Path(session_dir) / "prediction_result.csv"
     time_index = df.index[start_idx + context_len : start_idx + context_len + pred_len]
     result_data = {'date': time_index}
+
     for i in range(nvars):
         result_data[f'True_Var{i+1}'] = y_true[:, i]
         result_data[f'Pred_Median_Var{i+1}'] = pred_median[:, i]
+    
     result_df = pd.DataFrame(result_data)
     result_df.to_csv(csv_path, index=False)
 
@@ -329,6 +358,7 @@ def get_session_dir(session_id: gr.State):
         session_dir = Path(SESSION_DIR_ROOT) / session_uuid
         session_dir.mkdir(exist_ok=True, parents=True)
         session_id = str(session_dir)
+    
     return session_id
 
 def run_forecast(data_source, upload_file, index, context_len, pred_len, model_size, session_id: gr.State):
@@ -389,7 +419,7 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
     with gr.Row():
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("### 1. Data & Model Configuration")
-            # --- NEW: Add model selection dropdown ---
+            # --- Add model selection dropdown ---
             model_size = gr.Dropdown(
                 label="Select Model Size",
                 choices=["base", "large"],
@@ -459,7 +489,7 @@ with gr.Blocks(title="VisionTS++ Advanced Forecasting Platform", theme=gr.themes
 
     data_source.change(fn=toggle_upload_visibility, inputs=data_source, outputs=upload_file)
 
-    # --- NEW: Include model_size in the inputs list ---
+    # --- Include model_size in the inputs list ---
     inputs = [data_source, upload_file, sample_index, context_len, pred_len, model_size, session_id_state]
     outputs = [ts_plot, input_img_plot, recon_img_plot, download_csv, sample_index, freq_display, session_id_state]
 
@@ -474,7 +504,7 @@ if __name__ == "__main__":
     # --- Run initial cleanup on startup ---
     cleanup_old_sessions()
 
-    # --- NEW: Start the periodic cleanup in a background daemon thread ---
+    # --- Start the periodic cleanup in a background daemon thread ---
     cleanup_thread = threading.Thread(target=periodic_cleanup_task, daemon=True)
     cleanup_thread.start()