update app.py

app.py

@@ -18,15 +18,15 @@ REPO_ID = "Lefei/VisionTSpp"
 LOCAL_DIR = "./hf_models/VisionTSpp"
 CKPT_PATH = os.path.join(LOCAL_DIR, "visiontspp_model.ckpt")
 
-ARCH = 'mae_base'  # 可选: 'mae_base', 'mae_large', 'mae_huge'
+ARCH = 'mae_base'
 
-# 下载模型（Space 构建时执行一次）
+# 下载模型
 if not os.path.exists(CKPT_PATH):
     os.makedirs(LOCAL_DIR, exist_ok=True)
     print("Downloading model from Hugging Face Hub...")
     snapshot_download(repo_id=REPO_ID, local_dir=LOCAL_DIR, local_dir_use_symlinks=False)
 
-# 加载模型（全局加载一次）
+# 加载模型
 model = VisionTSpp(
     ARCH,
     ckpt_path=CKPT_PATH,
@@ -37,15 +37,41 @@ model = VisionTSpp(
 ).to(DEVICE)
 print(f"Model loaded on {DEVICE}")
 
+# Image normalization constants
+imagenet_mean = np.array([0.485, 0.456, 0.406])
+imagenet_std = np.array([0.229, 0.224, 0.225])
+
+
 # ========================
-# 核心预测与可视化函数
+# 可视化函数
 # ========================
 
+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))
+    ax.imshow(cur_image.numpy())
+    ax.set_title(title, fontsize=14)
+    ax.axis('off')
+    plt.close(fig)
+    return fig
+
+
 def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
     """
-    可视化真实值 vs 预测值
-    true: [T, nvars]
-    preds: [T, nvars]，与 true 对齐
+    绘制时间序列预测图（多变量）
     """
     if isinstance(true, torch.Tensor):
         true = true.cpu().numpy()
@@ -53,7 +79,6 @@ def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
         preds = preds.cpu().numpy()
 
     nvars = true.shape[1]
-
     FIG_WIDTH = 12
     FIG_HEIGHT_PER_VAR = 1.8
     FONT_S = 10
@@ -73,7 +98,6 @@ def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
             ax.plot(np.arange(lookback_len, len(true)), preds[lookback_len:, i],
                     label='Prediction (Median)', color='blue', linewidth=1.8)
 
-        # 分隔线
         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)
@@ -82,12 +106,10 @@ def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
         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})
 
-    # 计算 MSE/MAE
     if preds is not None:
         true_eval = true[-pred_len:]
         pred_eval = preds[-pred_len:]
@@ -96,77 +118,81 @@ def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
         fig.suptitle(f'MSE: {mse:.4f}, MAE: {mae:.4f}', fontsize=12, y=0.95)
 
     plt.subplots_adjust(hspace=0)
-    return fig  # 返回 matplotlib figure
+    plt.close(fig)
+    return fig
 
 
-def predict_and_visualize(df, context_len=960, pred_len=394, freq="15Min"):
+# ========================
+# 数据预处理与预测
+# ========================
+
+def predict_at_index(df, index, context_len=960, pred_len=394, freq="15Min"):
     """
-    输入: df (pandas.DataFrame)，必须包含 'date' 列和其他数值列
-    输出: matplotlib 图像
+    在指定 index 处预测
+    index: 第 index 个样本（从 0 开始）
+    返回: (ts_fig, input_img_fig, recon_img_fig)
     """
     if 'date' in df.columns:
-        df['date'] = pd.to_datetime(df['date'])
-        df = df.set_index('date')
-    else:
-        # 如果没有 date 列，假设是纯数值序列
-        df = df.copy()
+        df = df.set_index(pd.to_datetime(df['date'])).drop(columns=['date'])
 
     data = df.values  # [T, nvars]
     nvars = data.shape[1]
+    total_samples = len(data) - context_len - pred_len + 1
 
-    if data.shape[0] < context_len + pred_len:
-        raise ValueError(f"数据太短，至少需要 {context_len + pred_len} 行，当前只有 {data.shape[0]} 行。")
+    if total_samples <= 0:
+        raise ValueError(f"数据太短，无法构造任何样本（需要至少 {context_len + pred_len} 行）")
+    if index >= total_samples:
+        raise ValueError(f"索引越界，最大允许索引为 {total_samples - 1}")
 
-    # 归一化（使用训练集前 70% 的统计量）
+    # 归一化
     train_len = int(len(data) * 0.7)
     x_mean = data[:train_len].mean(axis=0, keepdims=True)
     x_std = data[:train_len].std(axis=0, keepdims=True) + 1e-8
     data_norm = (data - x_mean) / x_std
 
-    # 取最后一段作为测试窗口
-    end_idx = len(data_norm)
-    start_idx = end_idx - (context_len + pred_len)
-    x = data_norm[start_idx:start_idx + context_len]  # [context_len, nvars]
-    y_true = data_norm[start_idx + context_len:end_idx]  # [pred_len, nvars]
+    # 提取当前样本
+    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]
 
-    # 设置周期性
+    # 周期性
     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)]  # RGB 循环着色
-
-    # 更新模型配置
-    model.update_config(
-        context_len=context_len,
-        pred_len=pred_len,
-        periodicity=periodicity,
-        num_patch_input=7,
-        padding_mode='constant'
-    )
+    color_list = [i % 3 for i in range(nvars)]
+
+    model.update_config(context_len=context_len, pred_len=pred_len, periodicity=periodicity)
 
-    # 转为 tensor
     x_tensor = torch.FloatTensor(x).unsqueeze(0).to(DEVICE)  # [1, T, N]
-    y_true_tensor = torch.FloatTensor(y_true).unsqueeze(0).to(DEVICE)
 
-    # 预测
     with torch.no_grad():
-        y_pred, _, _, _, _ = model.forward(x_tensor, export_image=True, color_list=color_list)
+        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
 
     # 反归一化
     y_true_original = y_true * x_std + x_mean
     y_pred_original = y_pred_median[0].cpu().numpy() * x_std + x_mean
 
-    # 构造完整序列用于可视化
+    # 完整序列（用于可视化）
     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)
 
-    # 可视化
-    fig = visual_ts(true=full_true, preds=full_pred, lookback_len_visual=context_len, pred_len=pred_len)
-    return fig
+    # === 可视化 ===
+    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
+    )
+
+    return ts_fig, input_img_fig, recon_img_fig, total_samples
 
 
 # ========================
-# 默认数据加载
+# 默认数据
 # ========================
 def load_default_data():
     data_path = "./datasets/ETTm1.csv"
@@ -181,63 +207,84 @@ def load_default_data():
 
 
 # ========================
-# Gradio 界面
+# Gradio 接口
 # ========================
-def run_forecast(file_input, context_len, pred_len, freq):
+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 = "Uploaded Data Prediction"
+        title_prefix = "Uploaded Data"
     else:
         df = load_default_data()
-        title = "Default ETTm1 Dataset Prediction"
+        title_prefix = "ETTm1 Dataset"
 
     try:
-        fig = predict_and_visualize(df, context_len=int(context_len), pred_len=int(pred_len), freq=freq)
-        fig.suptitle(title, fontsize=14, y=0.98)
-        plt.close(fig)  # 防止重复显示
-        return fig
+        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
+        )
+
+        # 修改标题
+        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:
-        # 返回错误信息图像
-        fig, ax = plt.subplots()
-        ax.text(0.5, 0.5, f"Error: {str(e)}", ha='center', va='center', wrap=True)
-        ax.axis('off')
-        plt.close(fig)
-        return fig
+        # 错误图
+        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()
 
+
+# ========================
 # Gradio UI
-with gr.Blocks(title="VisionTS++ 时间序列预测") as demo:
+# ========================
+with gr.Blocks(title="VisionTS++ 多变量预测") as demo:
     gr.Markdown("# 🕰️ VisionTS++ 时间序列预测平台")
-    gr.Markdown("上传你的多变量时间序列 CSV 文件，或使用默认 ETTm1 数据进行预���。")
+    gr.Markdown("上传 CSV 或使用默认 ETTm1 数据。滑动选择不同样本进行预测，并查看原始图像表示。")
 
     with gr.Row():
-        file_input = gr.File(label="上传 CSV 文件（含 date 列或纯数值）", file_types=['.csv'])
-        with gr.Column():
-            context_len = gr.Number(label="历史长度 (context_len)", value=960)
-            pred_len = gr.Number(label="预测长度 (pred_len)", value=394)
-            freq = gr.Textbox(label="时间频率 (如 15Min, H)", value="15Min")
-
-    btn = gr.Button("🚀 开始预测")
-
-    output_plot = gr.Plot(label="预测结果")
-
+        with gr.Column(scale=2):
+            file_input = gr.File(label="上传 CSV 文件", file_types=['.csv'])
+            context_len = gr.Number(label="历史长度", value=960)
+            pred_len = gr.Number(label="预测长度", value=394)
+            freq = gr.Textbox(label="频率 (如 15Min)", 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="时间序列预测")
+            with gr.Row():
+                input_img_plot = gr.Plot(label="Input Image")
+                recon_img_plot = gr.Plot(label="Reconstructed Image")
+
+    btn = gr.Button("🚀 更新预测")
+
+    # 点击按钮或滑动条变化时更新
     btn.click(
         fn=run_forecast,
-        inputs=[file_input, context_len, pred_len, freq],
-        outputs=output_plot
+        inputs=[file_input, sample_index, context_len, pred_len, freq],
+        outputs=[ts_plot, input_img_plot, recon_img_plot, sample_index]
+    )
+
+    # 滑动条变化时也触发（但只在点击后才允许滑动）
+    # 我们用 sample_index.change 依赖于前一次运行的结果
+    demo.load(
+        fn=lambda: gr.update(maximum=100, value=0),
+        outputs=sample_index
     )
 
-    # 示例：使用默认数据
+    # 示例
     gr.Examples(
         examples=[
             [None, 960, 394, "15Min"]
         ],
         inputs=[file_input, context_len, pred_len, freq],
-        outputs=output_plot,
-        fn=run_forecast,
-        label="点击运行默认示例"
+        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="运行默认示例"
     )
 
 # 启动
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()