Create app.py

app.py

@@ -0,0 +1,243 @@
+# app.py
+import os
+import gradio as gr
+import torch
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import einops
+
+from huggingface_hub import snapshot_download
+from visionts import VisionTSpp, freq_to_seasonality_list
+
+# ========================
+# 配置
+# ========================
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+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'
+
+# 下载模型（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,
+    quantile=True,
+    clip_input=True,
+    complete_no_clip=False,
+    color=True
+).to(DEVICE)
+print(f"Model loaded on {DEVICE}")
+
+# ========================
+# 核心预测与可视化函数
+# ========================
+
+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()
+    if isinstance(preds, torch.Tensor):
+        preds = preds.cpu().numpy()
+
+    nvars = true.shape[1]
+
+    FIG_WIDTH = 12
+    FIG_HEIGHT_PER_VAR = 1.8
+    FONT_S = 10
+
+    fig, axes = plt.subplots(
+        nrows=nvars, ncols=1, figsize=(FIG_WIDTH, nvars * FIG_HEIGHT_PER_VAR), sharex=True,
+        gridspec_kw={'height_ratios': [1] * nvars}
+    )
+    if nvars == 1:
+        axes = [axes]
+
+    lookback_len = true.shape[0] - pred_len
+
+    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)
+
+        # 分隔线
+        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.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})
+
+    # 计算 MSE/MAE
+    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)
+
+    plt.subplots_adjust(hspace=0)
+    return fig  # 返回 matplotlib figure
+
+
+def predict_and_visualize(df, context_len=960, pred_len=394, freq="15Min"):
+    """
+    输入: df (pandas.DataFrame)，必须包含 'date' 列和其他数值列
+    输出: matplotlib 图像
+    """
+    if 'date' in df.columns:
+        df['date'] = pd.to_datetime(df['date'])
+        df = df.set_index('date')
+    else:
+        # 如果没有 date 列，假设是纯数值序列
+        df = df.copy()
+
+    data = df.values  # [T, nvars]
+    nvars = data.shape[1]
+
+    if data.shape[0] < context_len + pred_len:
+        raise ValueError(f"数据太短，至少需要 {context_len + pred_len} 行，当前只有 {data.shape[0]} 行。")
+
+    # 归一化（使用训练集前 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]
+
+    # 设置周期性
+    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'
+    )
+
+    # 转为 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_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
+
+
+# ========================
+# 默认数据加载
+# ========================
+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
+
+
+# ========================
+# Gradio 界面
+# ========================
+def run_forecast(file_input, context_len, pred_len, freq):
+    if file_input is not None:
+        df = pd.read_csv(file_input.name)
+        title = "Uploaded Data Prediction"
+    else:
+        df = load_default_data()
+        title = "Default ETTm1 Dataset Prediction"
+
+    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
+    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
+
+
+# Gradio UI
+with gr.Blocks(title="VisionTS++ 时间序列预测") as demo:
+    gr.Markdown("# 🕰️ VisionTS++ 时间序列预测平台")
+    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="预测结果")
+
+    btn.click(
+        fn=run_forecast,
+        inputs=[file_input, context_len, pred_len, freq],
+        outputs=output_plot
+    )
+
+    # 示例：使用默认数据
+    gr.Examples(
+        examples=[
+            [None, 960, 394, "15Min"]
+        ],
+        inputs=[file_input, context_len, pred_len, freq],
+        outputs=output_plot,
+        fn=run_forecast,
+        label="点击运行默认示例"
+    )
+
+# 启动
+if __name__ == "__main__":
+    demo.launch()