update app.py

app.py

@@ -37,28 +37,48 @@ model = VisionTSpp(
 ).to(DEVICE)
 print(f"Model loaded on {DEVICE}")
 
-# Image normalization constants
+# Image normalization
 imagenet_mean = np.array([0.485, 0.456, 0.406])
 imagenet_std = np.array([0.229, 0.224, 0.225])
 
 
+# ========================
+# 预设数据集
+# ========================
+PRESET_DATASETS = {
+    "ETTm1 (15-min)": "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/ETTm1.csv",
+    "ETTh1 (1-hour)": "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/ETTh1.csv",
+    "Illness": "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/illness.csv",
+    "Weather": "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/weather.csv"
+}
+
+# 本地缓存路径
+PRESET_DIR = "./preset_data"
+os.makedirs(PRESET_DIR, exist_ok=True)
+
+
+def load_preset_data(name):
+    url = PRESET_DATASETS[name]
+    path = os.path.join(PRESET_DIR, f"{name.split(' ')[0]}.csv")
+    if not os.path.exists(path):
+        df = pd.read_csv(url)
+        df.to_csv(path, index=False)
+    else:
+        df = pd.read_csv(path)
+    return df
+
+
 # ========================
 # 可视化函数
 # ========================
 
 def show_image_tensor(image, title='', cur_nvars=1, cur_color_list=None):
-    """
-    image: [H, W, 3] tensor
-    返回 matplotlib figure
-    """
     cur_image = torch.zeros_like(image)
     height_per_var = image.shape[0] // cur_nvars
-
     for i in range(cur_nvars):
         cur_color = cur_color_list[i]
         cur_image[i*height_per_var:(i+1)*height_per_var, :, cur_color] = \
             (image[i*height_per_var:(i+1)*height_per_var, :, cur_color] * imagenet_std[cur_color] + imagenet_mean[cur_color]) * 255
-
     cur_image = torch.clamp(cur_image, 0, 255).cpu().int()
 
     fig, ax = plt.subplots(figsize=(6, 6))
@@ -69,14 +89,18 @@ def show_image_tensor(image, title='', cur_nvars=1, cur_color_list=None):
     return fig
 
 
-def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
+def visual_ts_with_quantiles(true, pred_median, pred_quantiles, lookback_len_visual=300, pred_len=96, quantile_colors=None):
     """
-    绘制时间序列预测图（多变量）
+    可视化中叠加多个 quantile 区间
+    pred_quantiles: list of [pred_len, nvars] tensors
     """
     if isinstance(true, torch.Tensor):
         true = true.cpu().numpy()
-    if isinstance(preds, torch.Tensor):
-        preds = preds.cpu().numpy()
+    if isinstance(pred_median, torch.Tensor):
+        pred_median = pred_median.cpu().numpy()
+    for i, q in enumerate(pred_quantiles):
+        if isinstance(q, torch.Tensor):
+            pred_quantiles[i] = q.cpu().numpy()
 
     nvars = true.shape[1]
     FIG_WIDTH = 12
@@ -92,55 +116,69 @@ def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
 
     lookback_len = true.shape[0] - pred_len
 
+    # Quantile 颜色（从外到内）
+    if quantile_colors is None:
+        quantile_colors = ['lightblue', 'skyblue', 'deepskyblue']
+
     for i, ax in enumerate(axes):
-        ax.plot(true[:, i], label='Ground Truth', color='gray', linewidth=1.8)
-        if preds is not None:
-            ax.plot(np.arange(lookback_len, len(true)), preds[lookback_len:, i],
-                    label='Prediction (Median)', color='blue', linewidth=1.8)
+        ax.plot(true[:, i], label='Ground Truth', color='gray', linewidth=2)
+        ax.plot(np.arange(lookback_len, len(true)), pred_median[lookback_len:, i],
+                label='Prediction (Median)', color='blue', linewidth=2)
+
+        # 绘制 quantile 区间（从外到内）
+        base = pred_median[lookback_len:]
+        quantiles_sorted = sorted(zip(PREDS.quantiles, pred_quantiles), key=lambda x: x[0])
+        for (q, pred_q), color in zip(quantiles_sorted, quantile_colors):
+            upper = pred_q[lookback_len:]
+            lower = 2 * base - upper  # 对称假设
+            ax.fill_between(
+                np.arange(lookback_len, len(true)),
+                lower[:, i], upper[:, i],
+                color=color, alpha=0.5, label=f'Quantile {q:.1f}'
+            )
 
         y_min, y_max = ax.get_ylim()
-        ax.vlines(x=lookback_len, ymin=y_min, ymax=y_max,
-                  colors='gray', linestyles='--', alpha=0.7, linewidth=1)
-
+        ax.vlines(x=lookback_len, ymin=y_min, ymax=y_max, colors='gray', linestyles='--', alpha=0.7)
         ax.set_yticks([])
         ax.set_xticks([])
         ax.text(0.005, 0.8, f'Var {i+1}', transform=ax.transAxes, fontsize=FONT_S, weight='bold')
 
-    if preds is not None:
-        handles, labels = axes[0].get_legend_handles_labels()
-        fig.legend(handles, labels, loc='upper right', bbox_to_anchor=(0.9, 0.9), prop={'size': FONT_S})
-
-    if preds is not None:
-        true_eval = true[-pred_len:]
-        pred_eval = preds[-pred_len:]
-        mse = np.mean((true_eval - pred_eval) ** 2)
-        mae = np.mean(np.abs(true_eval - pred_eval))
-        fig.suptitle(f'MSE: {mse:.4f}, MAE: {mae:.4f}', fontsize=12, y=0.95)
-
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(handles, labels, loc='upper right', bbox_to_anchor=(0.9, 0.9), prop={'size': FONT_S})
     plt.subplots_adjust(hspace=0)
     plt.close(fig)
     return fig
 
 
 # ========================
-# 数据预处理与预测
+# 预测类封装（便于复用）
 # ========================
+class PredictionResult:
+    def __init__(self, ts_fig, input_img_fig, recon_img_fig, csv_path, total_samples):
+        self.ts_fig = ts_fig
+        self.input_img_fig = input_img_fig
+        self.recon_img_fig = recon_img_fig
+        self.csv_path = csv_path
+        self.total_samples = total_samples
+
 
 def predict_at_index(df, index, context_len=960, pred_len=394, freq="15Min"):
-    """
-    在指定 index 处预测
-    index: 第 index 个样本（从 0 开始）
-    返回: (ts_fig, input_img_fig, recon_img_fig)
-    """
-    if 'date' in df.columns:
-        df = df.set_index(pd.to_datetime(df['date'])).drop(columns=['date'])
+    # === 数据校验 ===
+    if 'date' not in df.columns:
+        raise ValueError("❌ 数据集必须包含名为 'date' 的时间列。")
 
-    data = df.values  # [T, nvars]
+    try:
+        df['date'] = pd.to_datetime(df['date'])
+    except Exception:
+        raise ValueError("❌ 'date' 列格式无法解析为时间，请检查日期格式。")
+
+    df = df.sort_values('date').set_index('date')
+    data = df.values
     nvars = data.shape[1]
-    total_samples = len(data) - context_len - pred_len + 1
 
+    total_samples = len(data) - context_len - pred_len + 1
     if total_samples <= 0:
-        raise ValueError(f"数据太短，无法构造任何样本（需要至少 {context_len + pred_len} 行）")
+        raise ValueError(f"数据太短，至少需要 {context_len + pred_len} 行，当前只有 {len(data)} 行。")
     if index >= total_samples:
         raise ValueError(f"索引越界，最大允许索引为 {total_samples - 1}")
 
@@ -150,141 +188,149 @@ def predict_at_index(df, index, context_len=960, pred_len=394, freq="15Min"):
     x_std = data[:train_len].std(axis=0, keepdims=True) + 1e-8
     data_norm = (data - x_mean) / x_std
 
-    # 提取当前样本
     start_idx = index
-    x = data_norm[start_idx:start_idx + context_len]        # [context_len, nvars]
-    y_true = data_norm[start_idx + context_len:start_idx + context_len + pred_len]  # [pred_len, nvars]
+    x = data_norm[start_idx:start_idx + context_len]
+    y_true = data_norm[start_idx + context_len:start_idx + context_len + pred_len]
 
-    # 周期性
     periodicity_list = freq_to_seasonality_list(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)
 
-    x_tensor = torch.FloatTensor(x).unsqueeze(0).to(DEVICE)  # [1, T, N]
+    x_tensor = torch.FloatTensor(x).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
         y_pred, input_image, reconstructed_image, nvars_out, color_list_out = model.forward(
             x_tensor, export_image=True, color_list=color_list
         )
-        y_pred_median = y_pred[0]  # median prediction
+        pred_median, pred_quantiles = y_pred  # list of quantiles
 
     # 反归一化
-    y_true_original = y_true * x_std + x_mean
-    y_pred_original = y_pred_median[0].cpu().numpy() * x_std + x_mean
+    y_true_orig = y_true * x_std + x_mean
+    pred_med_orig = pred_median[0].cpu().numpy() * x_std + x_mean
+    pred_quants_orig = [q[0].cpu().numpy() * x_std + x_mean for q in pred_quantiles]
 
-    # 完整序列（用于可视化）
-    full_true = np.concatenate([x * x_std + x_mean, y_true_original], axis=0)
-    full_pred = np.concatenate([x * x_std + x_mean, y_pred_original], axis=0)
+    # 完整序列
+    full_true = np.concatenate([x * x_std + x_mean, y_true_orig], axis=0)
+    full_pred_med = np.concatenate([x * x_std + x_mean, pred_med_orig], axis=0)
 
     # === 可视化 ===
-    ts_fig = visual_ts(true=full_true, preds=full_pred, lookback_len_visual=context_len, pred_len=pred_len)
-
-    input_img_fig = show_image_tensor(
-        input_image[0, 0], title=f'Input Image (Sample {index})', cur_nvars=nvars, cur_color_list=color_list
-    )
-    recon_img_fig = show_image_tensor(
-        reconstructed_image[0, 0], title=f'Reconstructed Image', cur_nvars=nvars, cur_color_list=color_list
+    ts_fig = visual_ts_with_quantiles(
+        true=full_true,
+        pred_median=full_pred_med,
+        pred_quantiles=pred_quants_orig,
+        lookback_len_visual=context_len,
+        pred_len=pred_len
     )
 
-    return ts_fig, input_img_fig, recon_img_fig, total_samples
+    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 ===
+    os.makedirs("outputs", exist_ok=True)
+    csv_path = "outputs/prediction_result.csv"
+    time_index = df.index[start_idx:start_idx + context_len + pred_len]
+    combined = np.concatenate([full_true, full_pred_med], axis=1)  # [T, 2*nvars]
+    col_names = [f"True_Var{i+1}" for i in range(nvars)] + [f"Pred_Var{i+1}" for i in range(nvars)]
+    result_df = pd.DataFrame(combined, index=time_index, columns=col_names)
+    result_df.to_csv(csv_path)
 
-# ========================
-# 默认数据
-# ========================
-def load_default_data():
-    data_path = "./datasets/ETTm1.csv"
-    if not os.path.exists(data_path):
-        os.makedirs("./datasets", exist_ok=True)
-        url = "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/ETTm1.csv"
-        df = pd.read_csv(url)
-        df.to_csv(data_path, index=False)
-    else:
-        df = pd.read_csv(data_path)
-    return df
+    return PredictionResult(ts_fig, input_img_fig, recon_img_fig, csv_path, total_samples)
 
 
 # ========================
-# Gradio 接口
+# Gradio 接口函数
 # ========================
-def run_forecast(file_input, sample_index, context_len, pred_len, freq):
-    if file_input is not None:
-        df = pd.read_csv(file_input.name)
-        title_prefix = "Uploaded Data"
+def run_forecast(data_source, upload_file, index, context_len, pred_len, freq):
+    if data_source == "Upload CSV":
+        if upload_file is None:
+            raise ValueError("请上传一个 CSV 文件")
+        df = pd.read_csv(upload_file.name)
     else:
-        df = load_default_data()
-        title_prefix = "ETTm1 Dataset"
+        df = load_preset_data(data_source)
 
     try:
-        ts_fig, input_img_fig, recon_img_fig, total_samples = predict_at_index(
-            df, int(sample_index), context_len=int(context_len), pred_len=int(pred_len), freq=freq
+        result = predict_at_index(df, int(index), context_len=int(context_len), pred_len=int(pred_len), freq=freq)
+        return (
+            result.ts_fig,
+            result.input_img_fig,
+            result.recon_img_fig,
+            result.csv_path,
+            gr.update(maximum=result.total_samples - 1, value=min(index, result.total_samples - 1))
         )
-
-        # 修改标题
-        ts_fig.suptitle(f"{title_prefix} - Sample {int(sample_index)}", fontsize=14, y=0.98)
-
-        return ts_fig, input_img_fig, recon_img_fig, gr.update(maximum=total_samples - 1, value=total_samples - 1)
     except Exception as e:
-        # 错误图
-        def error_fig(msg):
-            fig, ax = plt.subplots()
-            ax.text(0.5, 0.5, msg, ha='center', va='center', wrap=True)
-            ax.axis('off')
-            plt.close(fig)
-            return fig
-
-        return error_fig("Error"), error_fig("Error"), error_fig("Error"), gr.Number()
+        fig_err = plt.figure(figsize=(6, 4))
+        plt.text(0.5, 0.5, f"Error: {str(e)}", ha='center', va='center', wrap=True)
+        plt.axis('off')
+        plt.close(fig_err)
+        return fig_err, fig_err, fig_err, None, gr.update()
 
 
 # ========================
 # Gradio UI
 # ========================
-with gr.Blocks(title="VisionTS++ 多变量预测") as demo:
-    gr.Markdown("# 🕰️ VisionTS++ 时间序列预测平台")
-    gr.Markdown("上传 CSV 或使用默认 ETTm1 数据。滑动选择不同样本进行预测，并查看原始图像表示。")
+with gr.Blocks(title="VisionTS++ 高级预测平台") as demo:
+    gr.Markdown("# 🕰️ VisionTS++ 多变量时间序列预测平台")
+    gr.Markdown("""
+    - ✅ 支持预设数据集或本地上传
+    - ✅ 上传规则：必须是 `.csv`，且包含 `date` 列
+    - ✅ 显示多分位数预测区间
+    - ✅ 支持下载预测结果
+    - ✅ 滑动样本实时预测
+    """)
 
     with gr.Row():
         with gr.Column(scale=2):
-            file_input = gr.File(label="上传 CSV 文件", file_types=['.csv'])
+            data_source = gr.Dropdown(
+                label="选择数据源",
+                choices=["ETTm1 (15-min)", "ETTh1 (1-hour)", "Illness", "Weather", "Upload CSV"],
+                value="ETTm1 (15-min)"
+            )
+            upload_file = gr.File(label="上传 CSV 文件", file_types=['.csv'], visible=False)
             context_len = gr.Number(label="历史长度", value=960)
             pred_len = gr.Number(label="预测长度", value=394)
-            freq = gr.Textbox(label="频率 (如 15Min)", value="15Min")
+            freq = gr.Textbox(label="频率 (如 15Min, H)", value="15Min")
             sample_index = gr.Slider(label="样本索引", minimum=0, maximum=100, step=1, value=0)
 
         with gr.Column(scale=3):
-            ts_plot = gr.Plot(label="时间序列预测")
+            ts_plot = gr.Plot(label="时间序列预测（含分位数区间）")
             with gr.Row():
                 input_img_plot = gr.Plot(label="Input Image")
                 recon_img_plot = gr.Plot(label="Reconstructed Image")
+            download_csv = gr.File(label="下载预测结果")
+
+    btn = gr.Button("🚀 初始运行")
 
-    btn = gr.Button("🚀 更新预测")
+    # 上传切换
+    def toggle_upload(choice):
+        return gr.update(visible=choice == "Upload CSV")
 
-    # 点击按钮或滑动条变化时更新
+    data_source.change(fn=toggle_upload, inputs=data_source, outputs=upload_file)
+
+    # 初始运行 + 滑动条变化都触发
     btn.click(
         fn=run_forecast,
-        inputs=[file_input, sample_index, context_len, pred_len, freq],
-        outputs=[ts_plot, input_img_plot, recon_img_plot, sample_index]
+        inputs=[data_source, upload_file, sample_index, context_len, pred_len, freq],
+        outputs=[ts_plot, input_img_plot, recon_img_plot, download_csv, sample_index]
     )
 
-    # 滑动条变化时也触发（但只在点击后才允许滑动）
-    # 我们用 sample_index.change 依赖于前一次运行的结果
-    demo.load(
-        fn=lambda: gr.update(maximum=100, value=0),
-        outputs=sample_index
+    # 【关键】滑动条变化时重新预测
+    sample_index.change(
+        fn=run_forecast,
+        inputs=[data_source, upload_file, sample_index, context_len, pred_len, freq],
+        outputs=[ts_plot, input_img_plot, recon_img_plot, download_csv, sample_index]
     )
 
     # 示例
     gr.Examples(
         examples=[
-            [None, 960, 394, "15Min"]
+            ["ETTm1 (15-min)", None, 960, 394, "15Min"],
+            ["Illness", None, 36, 24, "D"]
         ],
-        inputs=[file_input, context_len, pred_len, freq],
-        outputs=[ts_plot, input_img_plot, recon_img_plot, sample_index],
-        fn=lambda f, i, c, p, fr: run_forecast(f, 0, c, p, fr),  # 默认 index=0
-        label="运行默认示例"
+        inputs=[data_source, upload_file, context_len, pred_len, freq],
+        fn=lambda a,b,c,d,e: run_forecast(a,b,0,c,d,e),
+        label="点击运行示例"
     )
 
-# 启动
 demo.launch()