{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6603a8fc-d9da-4037-b845-d9c38bae4ce4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, Dataset, random_split\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c28cc123-71be-47ff-b78f-3a4d5592df39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum pixel value in the dataset: 107.49169921875\n"
     ]
    }
   ],
   "source": [
    "max_pixel_value = 107.49169921875\n",
    "\n",
    "print(f\"Maximum pixel value in the dataset: {max_pixel_value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dbfe80ce-4394-449c-a9a4-22ed15b2b8f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NO2Dataset(Dataset):\n",
    "    \n",
    "    def __init__(self, image_dir, mask_dir):\n",
    "        \n",
    "        self.image_dir = image_dir\n",
    "        self.mask_dir = mask_dir\n",
    "        self.image_filenames = [f for f in os.listdir(image_dir) if f.endswith('.npy')]  # 仅加载 .npy 文件\n",
    "        self.mask_filenames = [f for f in os.listdir(mask_dir) if f.endswith('.jpg')]  # 仅加载 .jpg 文件\n",
    "        \n",
    "    def __len__(self):\n",
    "        \n",
    "        return len(self.image_filenames)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        image_path = os.path.join(self.image_dir, self.image_filenames[idx])\n",
    "        mask_idx = np.random.choice(self.mask_filenames)\n",
    "        mask_path = os.path.join(self.mask_dir, mask_idx)\n",
    "\n",
    "        # 加载图像数据 (.npy 文件)\n",
    "        image = np.expand_dims(np.load(image_path).astype(np.float32), axis=2) / max_pixel_value  # 形状为 (96, 96, 1)\n",
    "\n",
    "        # 加载掩码数据 (.jpg 文件)\n",
    "        mask = np.array(Image.open(mask_path).convert('L')).astype(np.float32)\n",
    "\n",
    "        # 将掩码数据中非0值设为1，0值保持不变\n",
    "        mask = np.where(mask != 0, 1.0, 0.0)\n",
    "\n",
    "        # 保持掩码数据形状为 (96, 96, 1)\n",
    "        mask = mask[:, :, np.newaxis]  # 将形状调整为 (96, 96, 1)\n",
    "\n",
    "        # 应用掩码\n",
    "        masked_image = image.copy()\n",
    "        masked_image[:, :, 0] = image[:, :, 0] * mask.squeeze()  # 遮盖NO2数据\n",
    "\n",
    "        # cGAN的输入和目标\n",
    "        X = masked_image[:, :, :1]  # 形状为 (96, 96, 8)\n",
    "        y = image[:, :, 0:1]  # 目标输出为NO2数据，形状为 (96, 96, 1)\n",
    "\n",
    "        # 转换形状为 (channels, height, width)\n",
    "        X = np.transpose(X, (2, 0, 1))  # 转换为 (1, 96, 96)\n",
    "        y = np.transpose(y, (2, 0, 1))  # 转换为 (1, 96, 96)\n",
    "        mask = np.transpose(mask, (2, 0, 1))  # 转换为 (1, 96, 96)\n",
    "\n",
    "        return torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "70797703-1619-4be7-b965-5506b3d1e775",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 可视化特定特征的函数\n",
    "def visualize_feature(input_feature,masked_feature, output_feature, title):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.imshow(input_feature[0].cpu().numpy(), cmap='RdYlGn_r')\n",
    "    plt.title(title + \" Input\")\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.imshow(masked_feature[0].cpu().numpy(), cmap='RdYlGn_r')\n",
    "    plt.title(title + \" Masked\")\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.imshow(output_feature[0].detach().cpu().numpy(), cmap='RdYlGn_r')\n",
    "    plt.title(title + \" Recovery\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "645114e8-65a4-4867-b3fe-23395288e855",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Conv(nn.Sequential):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, stride=1, bias=False):\n",
    "        super(Conv, self).__init__(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, bias=bias,\n",
    "                      dilation=dilation, stride=stride, padding=((stride - 1) + dilation * (kernel_size - 1)) // 2)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2af52d0e-b785-4a84-838c-6fcfe2568722",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvBNReLU(nn.Sequential):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, stride=1, norm_layer=nn.BatchNorm2d,\n",
    "                 bias=False):\n",
    "        super(ConvBNReLU, self).__init__(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, bias=bias,\n",
    "                      dilation=dilation, stride=stride, padding=((stride - 1) + dilation * (kernel_size - 1)) // 2),\n",
    "            norm_layer(out_channels),\n",
    "            nn.ReLU()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "31ecf247-e98b-4977-a145-782914a042bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SeparableBNReLU(nn.Sequential):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, norm_layer=nn.BatchNorm2d):\n",
    "        super(SeparableBNReLU, self).__init__(\n",
    "            nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, dilation=dilation,\n",
    "                      padding=((stride - 1) + dilation * (kernel_size - 1)) // 2, groups=in_channels, bias=False),\n",
    "            # 分离卷积，仅调整空间信息\n",
    "            norm_layer(in_channels),  # 对输入通道进行归一化\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),  # 这里进行升维操作\n",
    "            nn.ReLU6()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7827bee2-74f7-4e47-b8c6-e41d5670e8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1, downsample=None):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "\n",
    "        # 如果输入和输出通道不一致，进行降采样操作\n",
    "        self.downsample = downsample\n",
    "        if in_channels != out_channels or stride != 1:\n",
    "            self.downsample = nn.Sequential(\n",
    "                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(out_channels)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "\n",
    "        out += identity\n",
    "        out = self.relu(out)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7853bf62-02f5-4917-b950-6fdfe467df4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mlp(nn.Module):\n",
    "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU6, drop=0.):\n",
    "        super().__init__()\n",
    "        out_features = out_features or in_features\n",
    "        hidden_features = hidden_features or in_features\n",
    "        self.fc1 = nn.Conv2d(in_features, hidden_features, 1, 1, 0, bias=True)\n",
    "\n",
    "        self.act = act_layer()\n",
    "        self.fc2 = nn.Conv2d(hidden_features, out_features, 1, 1, 0, bias=True)\n",
    "        self.drop = nn.Dropout(drop, inplace=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.act(x)\n",
    "        x = self.drop(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.drop(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e2375881-a11b-47a7-8f56-2eadb25010b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionBlock(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads, dropout=0.1):\n",
    "        super(MultiHeadAttentionBlock, self).__init__()\n",
    "        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)\n",
    "        self.norm = nn.LayerNorm(embed_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # (B, C, H, W) -> (HW, B, C) for MultiheadAttention compatibility\n",
    "        B, C, H, W = x.shape\n",
    "        x = x.view(B, C, H * W).permute(2, 0, 1)  # (B, C, H, W) -> (HW, B, C)\n",
    "\n",
    "        # Apply multihead attention\n",
    "        attn_output, _ = self.attention(x, x, x)\n",
    "\n",
    "        # Apply normalization and dropout\n",
    "        attn_output = self.norm(attn_output)\n",
    "        attn_output = self.dropout(attn_output)\n",
    "\n",
    "        # Reshape back to (B, C, H, W)\n",
    "        attn_output = attn_output.permute(1, 2, 0).view(B, C, H, W)\n",
    "\n",
    "        return attn_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "82a15d3d-2f8d-42ec-9146-87c8a4abe384",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpatialAttentionBlock(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SpatialAttentionBlock, self).__init__()\n",
    "        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False)\n",
    "\n",
    "    def forward(self, x): #(B, 64, H, W)\n",
    "        avg_out = torch.mean(x, dim=1, keepdim=True) #(B, 1, H, W)\n",
    "        max_out, _ = torch.max(x, dim=1, keepdim=True)#(B, 1, H, W)\n",
    "        out = torch.cat([avg_out, max_out], dim=1)#(B, 2, H, W)\n",
    "        out = torch.sigmoid(self.conv(out))#(B, 1, H, W)\n",
    "        return x * out #(B, C, H, W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "497bb9f1-1ac5-4d7f-a930-0ea222b9d1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderAttentionBlock(nn.Module):\n",
    "    def __init__(self, in_channels):\n",
    "        super(DecoderAttentionBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1)\n",
    "        self.conv2 = nn.Conv2d(in_channels // 2, in_channels, kernel_size=1)\n",
    "        self.spatial_attention = SpatialAttentionBlock()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 通道注意力\n",
    "        b, c, h, w = x.size()\n",
    "        avg_pool = F.adaptive_avg_pool2d(x, 1)\n",
    "        max_pool = F.adaptive_max_pool2d(x, 1)\n",
    "\n",
    "        avg_out = self.conv1(avg_pool)\n",
    "        max_out = self.conv1(max_pool)\n",
    "\n",
    "        out = avg_out + max_out\n",
    "        out = torch.sigmoid(self.conv2(out))\n",
    "\n",
    "        # 添加空间注意力\n",
    "        out = x * out\n",
    "        out = self.spatial_attention(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "15b9d453-d8d9-43b8-aca2-904735fb3a99",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SEBlock(nn.Module):\n",
    "    def __init__(self, in_channels, reduced_dim):\n",
    "        super(SEBlock, self).__init__()\n",
    "        self.se = nn.Sequential(\n",
    "            nn.AdaptiveAvgPool2d(1),  # 全局平均池化\n",
    "            nn.Conv2d(in_channels, reduced_dim, kernel_size=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(reduced_dim, in_channels, kernel_size=1),\n",
    "            nn.Sigmoid()  # 使用Sigmoid是因为我们要对通道进行权重归一化\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * self.se(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6379adb7-8a87-4dd8-a695-4013a7b37830",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 定义Masked Autoencoder模型\n",
    "class MaskedAutoencoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MaskedAutoencoder, self).__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            Conv(1, 32, kernel_size=3, stride=2),\n",
    "            \n",
    "            nn.ReLU(),\n",
    "            \n",
    "            SEBlock(32,32),\n",
    "            \n",
    "            ConvBNReLU(32, 64, kernel_size=3, stride=2),\n",
    "            \n",
    "            ResidualBlock(64,64),\n",
    "            \n",
    "            SeparableBNReLU(64, 128, kernel_size=3, stride=2),\n",
    "            \n",
    "            MultiHeadAttentionBlock(embed_dim=128, num_heads=4),\n",
    "            \n",
    "            SEBlock(128, 128)\n",
    "        )\n",
    "        self.mlp = Mlp(in_features=128, hidden_features=256, out_features=128, act_layer=nn.ReLU6, drop=0.1)\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose2d(128, 32, kernel_size=3, stride=2, padding=1, output_padding=1),\n",
    "            nn.ReLU(),\n",
    "            \n",
    "            DecoderAttentionBlock(32),\n",
    "            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1),\n",
    "            nn.ReLU(),\n",
    "            \n",
    "            DecoderAttentionBlock(16),\n",
    "            nn.ReLU(),\n",
    "            \n",
    "            nn.ConvTranspose2d(16, 1, kernel_size=3, stride=2, padding=1, output_padding=1),  # 修改为 output_padding=1\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        encoded = self.encoder(x)\n",
    "        decoded = self.decoder(encoded)\n",
    "        return decoded\n",
    "\n",
    "# 实例化模型、损失函数和优化器\n",
    "model = MaskedAutoencoder()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e9c804e0-6f5c-40a7-aba7-a03a496cf427",
   "metadata": {},
   "outputs": [],
   "source": [
    "def masked_mse_loss(preds, target, mask):\n",
    "    loss = (preds - target) ** 2\n",
    "    loss = loss.mean(dim=-1)  # 对每个像素点求平均\n",
    "    loss = (loss * mask).sum() / mask.sum()  # 只计算被mask的像素点的损失\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "94457c6b-4c6e-4aff-946d-fe4c670bfe16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 评估函数\n",
    "def evaluate(model, device, data_loader, criterion):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (X, y, mask) in enumerate(data_loader):\n",
    "            X, y, mask = X.to(device), y.to(device), mask.to(device)\n",
    "            reconstructed = model(X)\n",
    "            if batch_idx == 8:\n",
    "                rand_ind = np.random.randint(0, len(y))\n",
    "                # visualize_feature(y[rand_ind], X[rand_ind], reconstructed[rand_ind], title='NO_2')\n",
    "            loss = masked_mse_loss(reconstructed, y, mask)\n",
    "            running_loss += loss.item()\n",
    "    return running_loss / (batch_idx + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "296ba6bd-2239-4948-b278-7edcb29bfd14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "# 数据准备\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "775bb9b8-1d6a-40f0-82e5-e1d6bc369e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model10 = torch.load('./models/MAE/final_10.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "44b36101-69a1-4b12-823b-8653110863c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model20 = torch.load('./models/MAE/final_20.pt')\n",
    "model30 = torch.load('./models/MAE/final_30.pt')\n",
    "model40 = torch.load('./models/MAE/final_40.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a8467686-0655-4056-8e01-56299eb89d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_percentage_error, mean_absolute_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9d13fb84-65e2-4e67-91a2-d6a4b36a0842",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 实例化数据集和数据加载器\n",
    "image_dir = './2022data/selected_data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5bb0c2d4-e05d-4611-b247-4b8b000e6fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_ioa(y_true, y_pred):\n",
    "    # 计算平均值\n",
    "    mean_observed = np.mean(y_true)\n",
    "    mean_predicted = np.mean(y_pred)\n",
    "\n",
    "    # 计算IoA\n",
    "    numerator = np.sum((y_true - y_pred) ** 2)\n",
    "    denominator = np.sum((np.abs(y_true - mean_observed) + np.abs(y_pred - mean_predicted)) ** 2)\n",
    "    IoA = 1 - (numerator / denominator)\n",
    "\n",
    "    return IoA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cab43bce-4d37-4f13-9153-9b9ced72fdaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_frame(model, mask_dir):\n",
    "    test_set = NO2Dataset(image_dir, mask_dir)\n",
    "    test_loader = DataLoader(test_set, batch_size=32, shuffle=False, num_workers=4)\n",
    "    eva_list_frame = list()\n",
    "    device = 'cpu'\n",
    "    model = model.to(device)\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (X, y, mask) in enumerate(test_loader):\n",
    "            X, y, mask = X.to(device), y.to(device), mask.to(device)\n",
    "            mask_rev = (torch.squeeze(mask, dim=1)==0) * 1 # mask取反获得修复区域\n",
    "            reconstructed = model(X)\n",
    "            rev_data = y * max_pixel_value\n",
    "            rev_recon = reconstructed * max_pixel_value\n",
    "            # todo: 这里需要只评估修补出来的模块\n",
    "            for i, sample in enumerate(rev_data):\n",
    "                used_mask = mask_rev[i]\n",
    "                data_label = sample[0] * used_mask\n",
    "                recon_no2 = rev_recon[i][0] * used_mask\n",
    "                data_label = data_label[used_mask==1]\n",
    "                recon_no2 = recon_no2[used_mask==1]\n",
    "                mae = mean_absolute_error(data_label, recon_no2)\n",
    "                rmse = np.sqrt(mean_squared_error(data_label, recon_no2))\n",
    "                mape = mean_absolute_percentage_error(data_label, recon_no2)\n",
    "                r2 = r2_score(data_label, recon_no2)\n",
    "                ioa = cal_ioa(data_label.detach().numpy(), recon_no2.detach().numpy())\n",
    "                r = np.corrcoef(data_label, recon_no2)[0, 1]\n",
    "                eva_list_frame.append([mae, rmse, mape, r2, ioa, r])\n",
    "    return eva_list_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7d791903-c6eb-4170-b816-07c127471aa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch(model, mask_dir):\n",
    "    test_set = NO2Dataset(image_dir, mask_dir)\n",
    "    test_loader = DataLoader(test_set, batch_size=32, shuffle=False, num_workers=4)\n",
    "    eva_list = list()\n",
    "    device = 'cpu'\n",
    "    model = model.to(device)\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (X, y, mask) in enumerate(test_loader):\n",
    "            X, y, mask = X.to(device), y.to(device), mask.to(device)\n",
    "            mask_rev = (torch.squeeze(mask, dim=1)==0) * 1 # mask取反获得修复区域\n",
    "            reconstructed = model(X)\n",
    "            rev_data = y * max_pixel_value\n",
    "            rev_recon = reconstructed * max_pixel_value\n",
    "            # todo: 这里需要只评估修补出来的模块\n",
    "            data_label = torch.squeeze(rev_data, dim=1) * mask_rev\n",
    "            data_label = data_label[mask_rev==1]\n",
    "            recon_no2 = torch.squeeze(rev_recon, dim=1) * mask_rev\n",
    "            recon_no2 = recon_no2[mask_rev==1]\n",
    "            mae = mean_absolute_error(data_label, recon_no2)\n",
    "            rmse = np.sqrt(mean_squared_error(data_label, recon_no2))\n",
    "            mape = mean_absolute_percentage_error(data_label, recon_no2)\n",
    "            r2 = r2_score(data_label, recon_no2)\n",
    "            ioa = cal_ioa(data_label.detach().numpy(), recon_no2.detach().numpy())\n",
    "            r = np.corrcoef(data_label, recon_no2)[0, 1]\n",
    "            eva_list.append([mae, rmse, mape, r2, ioa, r])\n",
    "    return eva_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0174210e-a771-47ed-a719-65c18e0185fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "eva_10 = predict_batch(model10, './out_mat/96/mask/10/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "1e9a5196-b8a6-42d8-b1fc-8f37decead81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>mae</th>\n",
       "      <th>rmse</th>\n",
       "      <th>mape</th>\n",
       "      <th>r2</th>\n",
       "      <th>ioa</th>\n",
       "      <th>r</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>1.091435</td>\n",
       "      <td>1.800776</td>\n",
       "      <td>0.144901</td>\n",
       "      <td>0.921502</td>\n",
       "      <td>0.978528</td>\n",
       "      <td>0.961399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.121988</td>\n",
       "      <td>0.277084</td>\n",
       "      <td>0.013961</td>\n",
       "      <td>0.021990</td>\n",
       "      <td>0.006471</td>\n",
       "      <td>0.011300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.780621</td>\n",
       "      <td>1.203425</td>\n",
       "      <td>0.112769</td>\n",
       "      <td>0.814169</td>\n",
       "      <td>0.938012</td>\n",
       "      <td>0.912011</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>1.010365</td>\n",
       "      <td>1.604889</td>\n",
       "      <td>0.134638</td>\n",
       "      <td>0.908699</td>\n",
       "      <td>0.974835</td>\n",
       "      <td>0.954552</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>1.084495</td>\n",
       "      <td>1.778022</td>\n",
       "      <td>0.143943</td>\n",
       "      <td>0.924748</td>\n",
       "      <td>0.979503</td>\n",
       "      <td>0.962814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>1.167333</td>\n",
       "      <td>1.943834</td>\n",
       "      <td>0.153004</td>\n",
       "      <td>0.936659</td>\n",
       "      <td>0.983079</td>\n",
       "      <td>0.969209</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>1.663302</td>\n",
       "      <td>3.638290</td>\n",
       "      <td>0.195296</td>\n",
       "      <td>0.966477</td>\n",
       "      <td>0.991180</td>\n",
       "      <td>0.984394</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              mae        rmse        mape          r2         ioa           r\n",
       "count  944.000000  944.000000  944.000000  944.000000  944.000000  944.000000\n",
       "mean     1.091435    1.800776    0.144901    0.921502    0.978528    0.961399\n",
       "std      0.121988    0.277084    0.013961    0.021990    0.006471    0.011300\n",
       "min      0.780621    1.203425    0.112769    0.814169    0.938012    0.912011\n",
       "25%      1.010365    1.604889    0.134638    0.908699    0.974835    0.954552\n",
       "50%      1.084495    1.778022    0.143943    0.924748    0.979503    0.962814\n",
       "75%      1.167333    1.943834    0.153004    0.936659    0.983079    0.969209\n",
       "max      1.663302    3.638290    0.195296    0.966477    0.991180    0.984394"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame.from_records(eva_10, columns=['mae', 'rmse', 'mape', 'r2', 'ioa', 'r']).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a34173c5-b193-4f5f-9a2a-8f577c013156",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>1.355741</td>\n",
       "      <td>2.298790</td>\n",
       "      <td>0.187385</td>\n",
       "      <td>0.874066</td>\n",
       "      <td>0.964384</td>\n",
       "      <td>0.936755</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.175765</td>\n",
       "      <td>0.420152</td>\n",
       "      <td>0.021570</td>\n",
       "      <td>0.036664</td>\n",
       "      <td>0.011227</td>\n",
       "      <td>0.019697</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.937516</td>\n",
       "      <td>1.491321</td>\n",
       "      <td>0.135232</td>\n",
       "      <td>0.670371</td>\n",
       "      <td>0.899252</td>\n",
       "      <td>0.821530</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>1.227770</td>\n",
       "      <td>2.003057</td>\n",
       "      <td>0.172465</td>\n",
       "      <td>0.852918</td>\n",
       "      <td>0.958502</td>\n",
       "      <td>0.925009</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>1.338980</td>\n",
       "      <td>2.233278</td>\n",
       "      <td>0.184734</td>\n",
       "      <td>0.879677</td>\n",
       "      <td>0.966322</td>\n",
       "      <td>0.939587</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>1.461651</td>\n",
       "      <td>2.517045</td>\n",
       "      <td>0.200911</td>\n",
       "      <td>0.900681</td>\n",
       "      <td>0.972358</td>\n",
       "      <td>0.950725</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>2.233936</td>\n",
       "      <td>4.265592</td>\n",
       "      <td>0.289657</td>\n",
       "      <td>0.949471</td>\n",
       "      <td>0.986560</td>\n",
       "      <td>0.976791</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                0           1           2           3           4           5\n",
       "count  944.000000  944.000000  944.000000  944.000000  944.000000  944.000000\n",
       "mean     1.355741    2.298790    0.187385    0.874066    0.964384    0.936755\n",
       "std      0.175765    0.420152    0.021570    0.036664    0.011227    0.019697\n",
       "min      0.937516    1.491321    0.135232    0.670371    0.899252    0.821530\n",
       "25%      1.227770    2.003057    0.172465    0.852918    0.958502    0.925009\n",
       "50%      1.338980    2.233278    0.184734    0.879677    0.966322    0.939587\n",
       "75%      1.461651    2.517045    0.200911    0.900681    0.972358    0.950725\n",
       "max      2.233936    4.265592    0.289657    0.949471    0.986560    0.976791"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eva_20 = predict_batch(model20, './out_mat/96/mask/20/')\n",
    "pd.DataFrame.from_records(eva_20).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c8a2ecdc-ef09-4fe2-a95e-9ede1a6a5e32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "      <td>944.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>1.539012</td>\n",
       "      <td>2.592209</td>\n",
       "      <td>0.198195</td>\n",
       "      <td>0.849743</td>\n",
       "      <td>0.956817</td>\n",
       "      <td>0.924245</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.199099</td>\n",
       "      <td>0.457944</td>\n",
       "      <td>0.021195</td>\n",
       "      <td>0.037078</td>\n",
       "      <td>0.011842</td>\n",
       "      <td>0.020082</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>1.072083</td>\n",
       "      <td>1.713604</td>\n",
       "      <td>0.153092</td>\n",
       "      <td>0.674728</td>\n",
       "      <td>0.878825</td>\n",
       "      <td>0.837543</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>1.404373</td>\n",
       "      <td>2.280456</td>\n",
       "      <td>0.183094</td>\n",
       "      <td>0.829249</td>\n",
       "      <td>0.950952</td>\n",
       "      <td>0.912680</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>1.509811</td>\n",
       "      <td>2.494275</td>\n",
       "      <td>0.195810</td>\n",
       "      <td>0.853937</td>\n",
       "      <td>0.958379</td>\n",
       "      <td>0.926343</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>1.649892</td>\n",
       "      <td>2.814400</td>\n",
       "      <td>0.211518</td>\n",
       "      <td>0.875720</td>\n",
       "      <td>0.964866</td>\n",
       "      <td>0.938508</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>2.427394</td>\n",
       "      <td>5.086926</td>\n",
       "      <td>0.281582</td>\n",
       "      <td>0.936612</td>\n",
       "      <td>0.982576</td>\n",
       "      <td>0.969190</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                0           1           2           3           4           5\n",
       "count  944.000000  944.000000  944.000000  944.000000  944.000000  944.000000\n",
       "mean     1.539012    2.592209    0.198195    0.849743    0.956817    0.924245\n",
       "std      0.199099    0.457944    0.021195    0.037078    0.011842    0.020082\n",
       "min      1.072083    1.713604    0.153092    0.674728    0.878825    0.837543\n",
       "25%      1.404373    2.280456    0.183094    0.829249    0.950952    0.912680\n",
       "50%      1.509811    2.494275    0.195810    0.853937    0.958379    0.926343\n",
       "75%      1.649892    2.814400    0.211518    0.875720    0.964866    0.938508\n",
       "max      2.427394    5.086926    0.281582    0.936612    0.982576    0.969190"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eva_30 = predict_batch(model30, './out_mat/96/mask/30/')\n",
    "pd.DataFrame.from_records(eva_30).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "478ad241-6774-4fb2-ae72-9abea0ca2a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "eva_40 = predict_batch(model40, './out_mat/96/mask/40/')\n",
    "pd.DataFrame.from_records(eva_40).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946d8ee3-608b-4327-b576-88bf723449d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14be731a-0334-4912-9d7b-5d040bcffa33",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
