autoencoder_gridpatch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torchvision import transforms as tf\n",
    "from models.mynn import functional as myf\n",
    "from models.mynn.loss import ssim\n",
    "from train import train\n",
    "import typing\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from skimage.data import shepp_logan_phantom\n",
    "from skimage.transform import resize\n",
    "import models as mymodels\n",
    "\n",
    "from datasets import Fastmri_320p\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torch import nn\n",
    "# import torch\n",
    "# import math\n",
    "# from typing import List\n",
    "# from einops import rearrange, repeat, pack, unpack\n",
    "\n",
    "# class SinusoidalPositionEncoding1d(nn.Module):\n",
    "#     \"\"\"Sinusoidal Position Encoding for 1-dimensions\"\"\"\n",
    "#     def __init__(self, d_model, pos_scale=1e2):\n",
    "#         super(SinusoidalPositionEncoding1d, self).__init__()\n",
    "#         self.d_model = d_model\n",
    "#         self.pos_scale = pos_scale\n",
    "#         self.conv1d = nn.Conv1d(1, d_model, 1)\n",
    "\n",
    "#         self.register_buffer('_div_term', torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)))\n",
    "\n",
    "#     def forward(self, tensor = None, position = None):\n",
    "#         \"\"\"tensor:[batch_size, seq_len, d_model]  \n",
    "#         position:[batch_size, seq_len, 1]  \n",
    "#         tensor is None: return positional encoding  \n",
    "#         position is None: return [0-1]*scale positional encoding for tensor  \n",
    "#         \"\"\"\n",
    "#         assert tensor is not None or position is not None, \"Either tensor or position must be provided\"\n",
    "\n",
    "#         if position is None:\n",
    "#             position = torch.arange(tensor.size(1), device=tensor.device).unsqueeze(0).unsqueeze(-1).float()\n",
    "#         position = (position - position.min()) / (position.max() - position.min())\n",
    "#         position = position * self.pos_scale\n",
    "\n",
    "#         assert len(position.size()) == 3, \"position must be [batch_size, seq_len, 1]\"\n",
    "\n",
    "#         pe = torch.zeros(position.size(0), position.size(1), self.d_model, device=position.device)\n",
    "#         pe[:, :, 0::2] = torch.sin(position * self._div_term)\n",
    "#         pe[:, :, 1::2] = torch.cos(position * self._div_term)\n",
    "\n",
    "#         if tensor is None:\n",
    "#             return pe        \n",
    "#         return tensor + pe\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# class SinusoidalPositionEncodingmd(nn.Module):\n",
    "#     \"\"\"Multi-dimensional Sinusoidal Position Encoding\"\"\"\n",
    "#     def __init__(self, d_model, pos_scale: float | List[float] = 1e2, n_dims=2):\n",
    "#         super(SinusoidalPositionEncodingmd, self).__init__()\n",
    "#         self.n_dims = n_dims\n",
    "#         self.d_model = d_model\n",
    "#         self.register_buffer('pos_scale', torch.tensor(pos_scale if isinstance(pos_scale, list) else [pos_scale] * n_dims))\n",
    "\n",
    "#         assert d_model % n_dims == 0, \"d_model must be divisible by n_dims\"\n",
    "#         self.register_buffer('_div_term', torch.exp(torch.arange(0, self.d_model // self.n_dims, 2).float() * -(math.log(10000.0) / self.d_model // self.n_dims)))\n",
    "\n",
    "    \n",
    "#     def forward(self, tensor, position = None):\n",
    "#         \"\"\"tensor:[batch_size, seq_len, d_model]  \n",
    "#         position:[batch_size, seq_len, n_dims]  \n",
    "#         tensor is None: return positional encoding   \n",
    "#         position is None: return [0-1]*scale positional encoding for tensor   \n",
    "#         \"\"\"\n",
    "#         assert tensor is not None or position is not None, \"Either tensor or position must be provided\"\n",
    "\n",
    "#         if position is None:\n",
    "#             position = []\n",
    "#             for i in range(self.n_dims):\n",
    "#                 pos = torch.arange(tensor.size(1), device=tensor.device).unsqueeze(0).unsqueeze(-1).float()\n",
    "#                 position.append(pos)\n",
    "#             position = torch.cat(position, dim=-1)\n",
    "#         position = (position - position.min()) / (position.max() - position.min())\n",
    "#         position = position * self.pos_scale\n",
    "\n",
    "#         assert len(position.size()) == 3, \"position must be [batch_size, seq_len, n_dims]\"\n",
    "\n",
    "#         pe = torch.zeros(position.size(0), position.size(1), self.d_model, device=position.device)\n",
    "#         for i in range(self.n_dims): # [aaabbbccc]\n",
    "#             pe[:, :, i * self.d_model // self.n_dims:(i+1) * self.d_model // self.n_dims:2] = torch.sin(position[:, :, i:i+1] * self._div_term)\n",
    "#             pe[:, :, i * self.d_model // self.n_dims + 1:(i+1) * self.d_model // self.n_dims:2] = torch.cos(position[:, :, i:i+1] * self._div_term)\n",
    "\n",
    "#         if tensor is None:\n",
    "#             return pe\n",
    "#         return tensor + pe\n",
    "\n",
    "\n",
    "# class LearnablePositionEncoding(nn.Module):\n",
    "#     \"\"\"Learnable Position Encoding\"\"\"\n",
    "#     def __init__(self, d_model, max_len=512):\n",
    "#         super(LearnablePositionEncoding, self).__init__()\n",
    "#         self.d_model = d_model\n",
    "#         self.max_len = max_len\n",
    "#         self._pe = nn.Parameter(torch.randn(max_len, d_model))\n",
    "\n",
    "#     def forward(self, tensor):\n",
    "#         return tensor + self._pe[:tensor.size(0), :]\n",
    "\n",
    "\n",
    "\n",
    "# import torch\n",
    "# from torch import nn\n",
    "\n",
    "# from einops import rearrange, repeat, pack, unpack\n",
    "# from einops.layers.torch import Rearrange\n",
    "\n",
    "# class FeedForward(nn.Module):\n",
    "#     def __init__(self, dim):\n",
    "#         super().__init__()\n",
    "#         self.net = nn.Sequential(\n",
    "#             nn.Linear(dim, dim),\n",
    "#             nn.GELU(),\n",
    "#             nn.Linear(dim, dim),\n",
    "#         )\n",
    "#     def forward(self, x):\n",
    "#         return self.net(x)\n",
    "\n",
    "\n",
    "# class Attention(nn.Module):\n",
    "#     def __init__(self, dim, heads = 8):\n",
    "#         super().__init__()\n",
    "#         inner_dim = dim *  heads\n",
    "#         dim_head = dim // heads\n",
    "\n",
    "#         self.heads = heads\n",
    "#         self.scale = dim_head ** -0.5\n",
    "\n",
    "#         self.attend = nn.Softmax(dim = -1)\n",
    "\n",
    "#         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)\n",
    "\n",
    "#         self.to_out = nn.Sequential(\n",
    "#             nn.Linear(inner_dim, dim),\n",
    "#         )\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         qkv = self.to_qkv(x).chunk(3, dim = -1)\n",
    "#         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)\n",
    "\n",
    "#         dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n",
    "\n",
    "#         attn = self.attend(dots)\n",
    "\n",
    "#         out = torch.matmul(attn, v)\n",
    "#         out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "#         return self.to_out(out)\n",
    "\n",
    "# class Transformer(nn.Module):\n",
    "#     def __init__(self, dim, depth, heads):\n",
    "#         super().__init__()\n",
    "#         self.norm = nn.LayerNorm(dim, eps = 1e-13)\n",
    "#         self.layers = nn.ModuleList([])\n",
    "#         for _ in range(depth):\n",
    "#             self.layers.append(nn.ModuleList([\n",
    "#                 Attention(dim, heads = heads),\n",
    "#                 FeedForward(dim),\n",
    "#             ]))\n",
    "#     def forward(self, x):\n",
    "#         for attn, ff in self.layers:\n",
    "#             x = attn(x) + x\n",
    "#             x = x * 0.5\n",
    "#             x = ff(x) + x\n",
    "#             x = x * 0.5\n",
    "#             x = self.norm(x)\n",
    "#         return x\n",
    "    \n",
    "# class ViT(nn.Module):\n",
    "#     def __init__(self, *, image_size, patch_size, dim, depth, heads, channels):\n",
    "#         super().__init__()\n",
    "#         image_height, image_width = image_size\n",
    "#         patch_height, patch_width = patch_size\n",
    "\n",
    "#         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "\n",
    "#         num_patches = (image_height // patch_height) * (image_width // patch_width)\n",
    "\n",
    "#         patch_dim = channels * patch_height * patch_width\n",
    "\n",
    "#         dim = patch_dim\n",
    "\n",
    "#         self.to_patch_embedding = nn.Sequential(\n",
    "#             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),\n",
    "#         )\n",
    "\n",
    "#         self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
    "\n",
    "#         self.transformer = Transformer(dim, depth, heads)\n",
    "\n",
    "#         self.mlp_head = nn.Sequential(\n",
    "#             nn.Linear(dim, dim),\n",
    "#             nn.GELU(),\n",
    "#             nn.Linear(dim, dim),\n",
    "#             Rearrange('b (h w) (p1 p2 c) -> b c (h p1) (w p2)', h = image_height // patch_height, w = image_width // patch_width, p1 = patch_height, p2 = patch_width, c = channels),\n",
    "#         )\n",
    "\n",
    "#     def forward(self, series):\n",
    "#         # b c h w \n",
    "#         x = self.to_patch_embedding(series)\n",
    "#         b, n, _ = x.shape\n",
    "#         x += self.pos_embedding[:, :(n)]\n",
    "\n",
    "#         x = self.transformer(x)\n",
    "\n",
    "#         recon = self.mlp_head(x)\n",
    "\n",
    "#         return recon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import models.mynn as mynn\n",
    "# from models.mynn import functional as myf\n",
    "\n",
    "# from torchkbnufft import KbNufft, KbNufftAdjoint, calc_tensor_spmatrix, calc_density_compensation_function, ToepNufft\n",
    "\n",
    "# class MAE(nn.Module):\n",
    "#     def __init__(self, \n",
    "#                  image_size = (320, 320),\n",
    "#                  patch_size = (16,16),\n",
    "#                  mask_ratio = 0.5,\n",
    "#                  in_channels = 16, out_channels = 16,\n",
    "#                 #  d_model_encoder = 1024, d_model_decoder = 1024,\n",
    "#                  nlayers_encoder = 4, nlayers_decoder = 4,\n",
    "#                  nheads_encoder = 8, nheads_decoder = 8,\n",
    "#                  ):\n",
    "#         super().__init__()\n",
    "\n",
    "#         image_height, image_width = image_size\n",
    "#         patch_height, patch_width = patch_size\n",
    "\n",
    "#         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "\n",
    "#         num_patches = (image_height // patch_height) * (image_width // patch_width)\n",
    "#         patch_dim = in_channels * patch_height * patch_width\n",
    "#         d_model_encoder = patch_dim\n",
    "#         d_model_decoder = patch_dim\n",
    "\n",
    "#         self.image_size = image_size\n",
    "#         # self.d_model_encoder = d_model_encoder\n",
    "#         # self.d_model_decoder = d_model_decoder\n",
    "#         self.nlayer_encoder = nlayers_encoder\n",
    "#         self.nlayer_decoder = nlayers_decoder\n",
    "#         self.nheads_encoder = nheads_encoder\n",
    "#         self.nheads_decoder = nheads_decoder\n",
    "#         self.mask_ratio = mask_ratio\n",
    "#         self.in_channels = in_channels\n",
    "#         self.out_channels = out_channels\n",
    "\n",
    "#         ### Encoder\n",
    "#         self.input_embed = nn.Sequential(\n",
    "#             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),\n",
    "#         )\n",
    "\n",
    "#         self.norm_encoder = nn.LayerNorm(d_model_encoder)\n",
    "\n",
    "#         self.pe_encoder = SinusoidalPositionEncoding1d(d_model = d_model_encoder, pos_scale=100)\n",
    "\n",
    "#         self.encoder = Transformer(dim = d_model_encoder, depth = nlayers_encoder, heads = nheads_encoder)\n",
    "\n",
    "\n",
    "#         ### Decoder\n",
    "#         self.decoder_embed = nn.Linear(d_model_encoder, d_model_decoder, bias=True)\n",
    "\n",
    "#         self.pe_decoder = SinusoidalPositionEncoding1d(d_model = d_model_decoder, pos_scale=100)\n",
    "\n",
    "#         self.mask_token = nn.Parameter(torch.zeros(1, 1, d_model_decoder))\n",
    "\n",
    "#         self.decoder = Transformer(dim = d_model_decoder, depth = nlayers_decoder, heads = nheads_decoder)\n",
    "\n",
    "#         self.norm_decoder = nn.LayerNorm(d_model_decoder)\n",
    "\n",
    "#         self.output_embed = nn.Sequential(\n",
    "#             nn.Linear(d_model_decoder, d_model_decoder),\n",
    "#             nn.GELU(),\n",
    "#             nn.Linear(d_model_decoder, d_model_decoder),\n",
    "#             Rearrange('b (h w) (p1 p2 c) -> b c (h p1) (w p2)', h = image_height // patch_height, w = image_width // patch_width, p1 = patch_height, p2 = patch_width, c = out_channels),\n",
    "#         )\n",
    "\n",
    "\n",
    "#     def mean_std_norm_complex(self, data: torch.Tensor, dim = [-1, -2, -3]):\n",
    "#         real = data.real\n",
    "#         imag = data.imag\n",
    "#         real = (real - real.mean(dim=dim, keepdim = True)) / real.std(dim=dim, keepdim = True)\n",
    "#         imag = (imag - imag.mean(dim=dim, keepdim = True)) / imag.std(dim=dim, keepdim = True)\n",
    "#         return real + 1j * imag\n",
    "\n",
    "\n",
    "#     def random_masking(self, x, mask_ratio):\n",
    "#         \"\"\"\n",
    "#         https://github.com/facebookresearch/mae/tree/main\n",
    "\n",
    "#         Perform per-sample random masking by per-sample shuffling.\n",
    "#         Per-sample shuffling is done by argsort random noise.\n",
    "#         x: [N, L, D], sequence\n",
    "#         \"\"\"\n",
    "#         N, L, D = x.shape  # batch, length, dim\n",
    "#         len_keep = int(L * (1 - mask_ratio))\n",
    "        \n",
    "#         noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]\n",
    "        \n",
    "#         # sort noise for each sample\n",
    "#         ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove\n",
    "#         ids_restore = torch.argsort(ids_shuffle, dim=1)\n",
    "\n",
    "#         # keep the first subset\n",
    "#         ids_keep = ids_shuffle[:, :len_keep]\n",
    "#         x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))\n",
    "\n",
    "#         # generate the binary mask: 0 is keep, 1 is remove\n",
    "#         mask = torch.ones([N, L], device=x.device)\n",
    "#         mask[:, :len_keep] = 0\n",
    "#         # unshuffle to get the binary mask\n",
    "#         mask = torch.gather(mask, dim=1, index=ids_restore)\n",
    "\n",
    "#         return x_masked, mask, ids_restore\n",
    "    \n",
    "#     def forward_encoder_train(self, x, mask_ratio):\n",
    "#         # x: [batch, phase, readout * channel * 2]\n",
    "#         # pos: [batch, phase, 1]\n",
    "\n",
    "#         ### Encoder\n",
    "#         # in training mode, the shape of kdata is matched with the shape of ktraj\n",
    "#         # kdata should be masked as same as ktraj\n",
    "#         x = self.pe_encoder(x)\n",
    "#         x, mask, ids_restore = self.random_masking(x, mask_ratio)\n",
    "\n",
    "#         # cls token\n",
    "#         pass\n",
    "\n",
    "#         # encode\n",
    "        \n",
    "#         x = self.encoder(x)\n",
    "#         x = self.norm_encoder(x)\n",
    "\n",
    "#         return x, mask, ids_restore\n",
    "    \n",
    "#     def forward_decoder_train(self, x, ids_restore):\n",
    "#         # x: [batch, phase, readout * channel * 2]\n",
    "#         # pos: [batch, phase, 1]\n",
    "\n",
    "#         ### Decoder\n",
    "#         # embed tokens\n",
    "#         x = self.decoder_embed(x)\n",
    "\n",
    "#         # append mask tokens to sequence\n",
    "#         mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] - x.shape[1], 1)\n",
    "#         x = torch.cat([x, mask_tokens], dim=1)\n",
    "#         x = torch.gather(x, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle\n",
    "\n",
    "#         # pe\n",
    "#         x = self.pe_decoder(x)\n",
    "\n",
    "#         # apply Transformer blocks\n",
    "#         x = self.decoder(x)\n",
    "#         x = self.norm_decoder(x)\n",
    "\n",
    "#         # predictor projection\n",
    "#         x = self.output_embed(x)\n",
    "\n",
    "#         # remove cls token\n",
    "#         pass\n",
    "\n",
    "#         return x\n",
    "\n",
    "\n",
    "#     def forward_train(self, x):\n",
    "#         # seq: [batch, channel, height, weight]\n",
    "#         x = self.input_embed(x) # [batch, seq, feature]\n",
    "#         # encode\n",
    "#         encoder_memory, mask, ids_restore = self.forward_encoder_train(x, self.mask_ratio)\n",
    "        \n",
    "#         # decode\n",
    "#         x = self.forward_decoder_train(encoder_memory, ids_restore)\n",
    "\n",
    "#         return {'pred':x, 'mask': mask}\n",
    "\n",
    "\n",
    "#     def forward_encoder_eval(self, x):\n",
    "#         # x: [batch, phase, readout * channel * 2]\n",
    "#         # pos: [batch, phase, 1]\n",
    "\n",
    "#         ### Encoder\n",
    "#         # in evaluation mode, the shape of kdata is not matched with the shape of ktraj\n",
    "#         # kdata should be masked as same as ktraj\n",
    "#         x = self.pe_encoder(x)\n",
    "\n",
    "#         # cls token\n",
    "#         pass\n",
    "\n",
    "#         # encode\n",
    "#         x = self.encoder(x)\n",
    "#         x = self.norm_encoder(x)\n",
    "\n",
    "#         return x\n",
    "    \n",
    "#     def forward_decoder_eval(self, x):\n",
    "#         # x: [batch, phase, readout * channel * 2]\n",
    "#         # pos: [batch, phase, 1]\n",
    "\n",
    "#         ### Decoder\n",
    "#         # embed tokens\n",
    "#         x = self.decoder_embed(x)\n",
    "\n",
    "#         # pe\n",
    "#         x = self.pe_decoder(x)\n",
    "\n",
    "#         # apply Transformer blocks\n",
    "#         x = self.decoder(x)\n",
    "#         x = self.norm_decoder(x)\n",
    "\n",
    "#         # predictor projection\n",
    "#         x = self.output_embed(x)\n",
    "\n",
    "#         # remove cls token\n",
    "#         pass\n",
    "\n",
    "#         return x\n",
    "    \n",
    "#     def forward_eval(self, x):\n",
    "#         # kdata: [batch, channel, phase, readout]\n",
    "#         # ktraj: [batch, phase, readout]\n",
    "#         x = self.input_embed(x)\n",
    "#         # encode\n",
    "#         encoder_memory = self.forward_encoder_eval(x)\n",
    "\n",
    "#         # decode\n",
    "#         x = self.forward_decoder_eval(encoder_memory)\n",
    "\n",
    "#         return x\n",
    "    \n",
    "\n",
    "#     def forward_loss_train(self, pred, gt):\n",
    "#         pred, mask = pred['pred'], pred['mask']\n",
    "#         pred = self.input_embed(pred)\n",
    "#         gt = self.input_embed(gt)\n",
    "\n",
    "#         loss = (pred - gt) ** 2\n",
    "#         loss = loss.mean(dim=-1)  # [N, L], mean loss per patch\n",
    "\n",
    "#         loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches\n",
    "#         return loss\n",
    "    \n",
    "#     def forward_loss_eval(self, pred, gt):\n",
    "#         loss = nn.functional.mse_loss(pred, gt)\n",
    "#         return loss\n",
    "    \n",
    "#     def forward_loss(self, pred, gt):\n",
    "#         if self.training:\n",
    "#             return self.forward_loss_train(pred, gt)\n",
    "#         else:\n",
    "#             return self.forward_loss_eval(pred, gt)\n",
    "\n",
    "#     # def forward_loss(self, imgs, pred, mask):\n",
    "#     #     \"\"\"\n",
    "#     #     imgs: [N, 3, H, W]\n",
    "#     #     pred: [N, L, p*p*3]\n",
    "#     #     mask: [N, L], 0 is keep, 1 is remove, \n",
    "#     #     \"\"\"\n",
    "#     #     target = self.patchify(imgs)\n",
    "#     #     if self.norm_pix_loss:\n",
    "#     #         mean = target.mean(dim=-1, keepdim=True)\n",
    "#     #         var = target.var(dim=-1, keepdim=True)\n",
    "#     #         target = (target - mean) / (var + 1.e-6)**.5\n",
    "\n",
    "#     #     loss = (pred - target) ** 2\n",
    "#     #     loss = loss.mean(dim=-1)  # [N, L], mean loss per patch\n",
    "\n",
    "#     #     loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches\n",
    "#     #     return loss\n",
    "    \n",
    "\n",
    "#     def forward(self, x):\n",
    "#         if self.training:\n",
    "#             return self.forward_train(x)\n",
    "#         else:\n",
    "#             return self.forward_eval(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# format: filename index_of_slice\n",
    "trainlistpath = \"/home/lyy/moco/datasets/.fastmri_320p/trainlist.txt\"\n",
    "vallistpath = \"/home/lyy/moco/datasets/.fastmri_320p/vallist.txt\"\n",
    "testlistpath = \"/home/lyy/moco/datasets/.fastmri_320p/testlist.txt\"\n",
    "showlistpath = \"/home/lyy/moco/datasets/.fastmri_320p/showlist.txt\"\n",
    "\n",
    "with open(trainlistpath, 'r') as f:\n",
    "    trainlist = f.readlines()\n",
    "    trainlist = [(line.split()[0], int(line.split()[1])) for line in trainlist]\n",
    "\n",
    "with open(vallistpath, 'r') as f:\n",
    "    vallist = f.readlines()\n",
    "    vallist = [(line.split()[0], int(line.split()[1])) for line in vallist]\n",
    "\n",
    "with open(testlistpath, 'r') as f:\n",
    "    testlist = f.readlines()\n",
    "    testlist = [(line.split()[0], int(line.split()[1])) for line in testlist]\n",
    "\n",
    "with open(showlistpath, 'r') as f:\n",
    "    showlist = f.readlines()\n",
    "    showlist = [(line.split()[0], int(line.split()[1])) for line in showlist]\n",
    "\n",
    "print(\"file in trainlist: \", len(trainlist))\n",
    "print(\"file in vallist: \", len(vallist))\n",
    "print(\"file in testlist: \", len(testlist))\n",
    "print(\"file in showlist: \", len(showlist))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loadtransform\n",
    "import torch\n",
    "from torch import nn\n",
    "from datasets import Fastmri_320p\n",
    "\n",
    "\n",
    "load_keys = ['']\n",
    "\n",
    "class Loadtransform(nn.Module):\n",
    "    \"\"\"Load the data to the memory\"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, data):\n",
    "        # input: ismrmrd_header rss csm kspace\n",
    "        return {'pics_image':torch.tensor(data[\"pics_image\"])}\n",
    "\n",
    "loadtransform = Loadtransform()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# All rights reserved.\n",
    "\n",
    "# This source code is licensed under the license found in the\n",
    "# LICENSE file in the root directory of this source tree.\n",
    "# --------------------------------------------------------\n",
    "# References:\n",
    "# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm\n",
    "# DeiT: https://github.com/facebookresearch/deit\n",
    "# --------------------------------------------------------\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from timm.models.vision_transformer import PatchEmbed, Block\n",
    "\n",
    "from utils.mae.util.pos_embed import get_2d_sincos_pos_embed\n",
    "\n",
    "\n",
    "class MaskedAutoencoderViT(nn.Module):\n",
    "    \"\"\" Masked Autoencoder with VisionTransformer backbone\n",
    "    \"\"\"\n",
    "    def __init__(self, img_size=224, patch_size=16, in_chans=3,\n",
    "                 embed_dim=1024, depth=24, num_heads=16,\n",
    "                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,\n",
    "                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):\n",
    "        super().__init__()\n",
    "        self.in_chans = in_chans\n",
    "\n",
    "        # --------------------------------------------------------------------------\n",
    "        # MAE encoder specifics\n",
    "        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)\n",
    "        num_patches = self.patch_embed.num_patches\n",
    "\n",
    "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n",
    "        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding\n",
    "\n",
    "        self.blocks = nn.ModuleList([\n",
    "            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)\n",
    "            for i in range(depth)])\n",
    "        self.norm = norm_layer(embed_dim)\n",
    "        # --------------------------------------------------------------------------\n",
    "\n",
    "        # --------------------------------------------------------------------------\n",
    "        # MAE decoder specifics\n",
    "        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)\n",
    "\n",
    "        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))\n",
    "\n",
    "        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding\n",
    "\n",
    "        self.decoder_blocks = nn.ModuleList([\n",
    "            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)\n",
    "            for i in range(decoder_depth)])\n",
    "\n",
    "        self.decoder_norm = norm_layer(decoder_embed_dim)\n",
    "        self.decoder_pred =  nn.Sequential(\n",
    "            nn.Linear(decoder_embed_dim, decoder_embed_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True),\n",
    "        ) # decoder to patch\n",
    "        \n",
    "        # --------------------------------------------------------------------------\n",
    "\n",
    "        self.norm_pix_loss = norm_pix_loss\n",
    "\n",
    "        self.initialize_weights()\n",
    "\n",
    "    def initialize_weights(self):\n",
    "        # initialization\n",
    "        # initialize (and freeze) pos_embed by sin-cos embedding\n",
    "        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)\n",
    "        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))\n",
    "\n",
    "        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)\n",
    "        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))\n",
    "\n",
    "        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)\n",
    "        w = self.patch_embed.proj.weight.data\n",
    "        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))\n",
    "\n",
    "        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)\n",
    "        torch.nn.init.normal_(self.cls_token, std=.02)\n",
    "        torch.nn.init.normal_(self.mask_token, std=.02)\n",
    "\n",
    "        # initialize nn.Linear and nn.LayerNorm\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            # we use xavier_uniform following official JAX ViT:\n",
    "            torch.nn.init.xavier_uniform_(m.weight)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "\n",
    "    def patchify(self, imgs):\n",
    "        \"\"\"\n",
    "        imgs: (N, self.in_chans, H, W)\n",
    "        x: (N, L, patch_size**2 *self.in_chans)\n",
    "        \"\"\"\n",
    "        p = self.patch_embed.patch_size[0]\n",
    "        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0\n",
    "\n",
    "        h = w = imgs.shape[2] // p\n",
    "        x = imgs.reshape(shape=(imgs.shape[0], self.in_chans, h, p, w, p))\n",
    "        x = torch.einsum('nchpwq->nhwpqc', x)\n",
    "        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * self.in_chans))\n",
    "        return x\n",
    "\n",
    "    def unpatchify(self, x):\n",
    "        \"\"\"\n",
    "        x: (N, L, patch_size**2 *3)\n",
    "        imgs: (N, 3, H, W)\n",
    "        \"\"\"\n",
    "        p = self.patch_embed.patch_size[0]\n",
    "        h = w = int(x.shape[1]**.5)\n",
    "        assert h * w == x.shape[1]\n",
    "        \n",
    "        x = x.reshape(shape=(x.shape[0], h, w, p, p, self.in_chans))\n",
    "        x = torch.einsum('nhwpqc->nchpwq', x)\n",
    "        imgs = x.reshape(shape=(x.shape[0], self.in_chans, h * p, h * p))\n",
    "        return imgs\n",
    "\n",
    "    def random_masking(self, x, mask_ratio):\n",
    "        \"\"\"\n",
    "        Perform per-sample random masking by per-sample shuffling.\n",
    "        Per-sample shuffling is done by argsort random noise.\n",
    "        x: [N, L, D], sequence\n",
    "        \"\"\"\n",
    "        N, L, D = x.shape  # batch, length, dim\n",
    "        len_keep = int(L * (1 - mask_ratio))\n",
    "        \n",
    "        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]\n",
    "        \n",
    "        # sort noise for each sample\n",
    "        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove\n",
    "        ids_restore = torch.argsort(ids_shuffle, dim=1)\n",
    "\n",
    "        # keep the first subset\n",
    "        ids_keep = ids_shuffle[:, :len_keep]\n",
    "        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))\n",
    "\n",
    "        # generate the binary mask: 0 is keep, 1 is remove\n",
    "        mask = torch.ones([N, L], device=x.device)\n",
    "        mask[:, :len_keep] = 0\n",
    "        # unshuffle to get the binary mask\n",
    "        mask = torch.gather(mask, dim=1, index=ids_restore)\n",
    "\n",
    "        return x_masked, mask, ids_restore\n",
    "\n",
    "    def forward_encoder_train(self, x, mask_ratio):\n",
    "        # embed patches\n",
    "        x = self.patch_embed(x)\n",
    "\n",
    "        # add pos embed w/o cls token\n",
    "        x = x + self.pos_embed[:, 1:, :]\n",
    "\n",
    "        # masking: length -> length * mask_ratio\n",
    "        x, mask, ids_restore = self.random_masking(x, mask_ratio)\n",
    "\n",
    "        # append cls token\n",
    "        cls_token = self.cls_token + self.pos_embed[:, :1, :]\n",
    "        cls_tokens = cls_token.expand(x.shape[0], -1, -1)\n",
    "        x = torch.cat((cls_tokens, x), dim=1)\n",
    "\n",
    "        # apply Transformer blocks\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "        x = self.norm(x)\n",
    "\n",
    "        return x, mask, ids_restore\n",
    "\n",
    "    def forward_decoder_train(self, x, ids_restore):\n",
    "        # embed tokens\n",
    "        x = self.decoder_embed(x)\n",
    "\n",
    "        # append mask tokens to sequence\n",
    "        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)\n",
    "        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token\n",
    "        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle\n",
    "        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token\n",
    "\n",
    "        # add pos embed\n",
    "        x = x + self.decoder_pos_embed\n",
    "\n",
    "        # apply Transformer blocks\n",
    "        for blk in self.decoder_blocks:\n",
    "            x = blk(x)\n",
    "        x = self.decoder_norm(x)\n",
    "\n",
    "        # predictor projection\n",
    "        x = self.decoder_pred(x)\n",
    "\n",
    "        # remove cls token\n",
    "        x = x[:, 1:, :]\n",
    "\n",
    "        return x\n",
    "    \n",
    "    def forward_train(self, imgs, mask_ratio=0.75):\n",
    "        latent, mask, ids_restore = self.forward_encoder_train(imgs, mask_ratio)\n",
    "        pred = self.forward_decoder_train(latent, ids_restore)  # [N, L, p*p*3]\n",
    "        return {'pred': pred,'mask': mask}\n",
    "    \n",
    "    def forward_encoder_eval(self, x, mask_ratio):\n",
    "        # embed patches\n",
    "        x = self.patch_embed(x)\n",
    "\n",
    "        # add pos embed w/o cls token\n",
    "        x = x + self.pos_embed[:, 1:, :]\n",
    "\n",
    "        # masking: length -> length * mask_ratio\n",
    "        x, mask, ids_restore = self.random_masking(x, mask_ratio)\n",
    "\n",
    "        # append cls token\n",
    "        cls_token = self.cls_token + self.pos_embed[:, :1, :]\n",
    "        cls_tokens = cls_token.expand(x.shape[0], -1, -1)\n",
    "        x = torch.cat((cls_tokens, x), dim=1)\n",
    "\n",
    "        # apply Transformer blocks\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "        x = self.norm(x)\n",
    "\n",
    "        return x, mask, ids_restore\n",
    "\n",
    "    def forward_decoder_eval(self, x, ids_restore):\n",
    "        # embed tokens\n",
    "        x = self.decoder_embed(x)\n",
    "\n",
    "        # append mask tokens to sequence\n",
    "        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)\n",
    "        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token\n",
    "        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle\n",
    "        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token\n",
    "\n",
    "        # add pos embed\n",
    "        x = x + self.decoder_pos_embed\n",
    "\n",
    "        # apply Transformer blocks\n",
    "        for blk in self.decoder_blocks:\n",
    "            x = blk(x)\n",
    "        x = self.decoder_norm(x)\n",
    "\n",
    "        # predictor projection\n",
    "        x = self.decoder_pred(x)\n",
    "\n",
    "        # remove cls token\n",
    "        x = x[:, 1:, :]\n",
    "\n",
    "        x = self.unpatchify(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    def forward_eval(self, imgs):\n",
    "        latent, _, ids = self.forward_encoder_eval(imgs, mask_ratio = 0.75)\n",
    "        pred = self.forward_decoder_eval(latent, ids)  # [N, L, p*p*3]\n",
    "        return pred\n",
    "\n",
    "    def forward_loss_train(self, pred, gt):\n",
    "        \"\"\"\n",
    "        imgs: [N, self.in_chans, H, W]\n",
    "        pred: [N, L, p*p*self.in_chans]\n",
    "        mask: [N, L], 0 is keep, 1 is remove, \n",
    "        \"\"\"\n",
    "        mask = pred['mask']\n",
    "        pred = pred['pred']\n",
    "        target = self.patchify(gt)\n",
    "        if self.norm_pix_loss:\n",
    "            mean = target.mean(dim=-1, keepdim=True)\n",
    "            var = target.var(dim=-1, keepdim=True)\n",
    "            target = (target - mean) / (var + 1.e-6)**.5\n",
    "\n",
    "        loss = (pred - target) ** 2\n",
    "        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch\n",
    "\n",
    "        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches\n",
    "        return loss\n",
    "    \n",
    "    def forward_loss_eval(self, imgs, pred):\n",
    "        return nn.functional.mse_loss(imgs, pred)\n",
    "    \n",
    "    def forward_loss(self, pred, gt):\n",
    "        if self.training:\n",
    "            return self.forward_loss_train(pred, gt)\n",
    "        else:\n",
    "            return self.forward_loss_eval(pred, gt)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        if self.training:\n",
    "            return self.forward_train(x)\n",
    "        else:\n",
    "            return self.forward_eval(x)\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = (320, 320)\n",
    "\n",
    "model = MaskedAutoencoderViT(\n",
    "    img_size = 320,\n",
    "    patch_size = 16,\n",
    "    in_chans = 1,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# step1 mseloss\n",
    "from datasets.phantom import Phantom\n",
    "from models.mynn.loss.pytorch_msssim.ssim import SSIM\n",
    "\n",
    "\n",
    "mse_fn = nn.MSELoss()\n",
    "\n",
    "ssim_fn = SSIM(data_range=1, size_average=True, channel=1)\n",
    "\n",
    "def loss_fn(x, y):\n",
    "    mse_score = mse_fn(x, y)\n",
    "    if mse_score < 0.01:\n",
    "        return 1 - ssim_fn(x, y)\n",
    "    return mse_score\n",
    "\n",
    "\n",
    "# trainset = Fastmri_320p(trainlist, transform=loadtransform, disk_cache=\"20241225autoencoder\")\n",
    "# valset = Fastmri_320p(vallist, transform=loadtransform, disk_cache=\"20241225autoencoder\")\n",
    "# testset = Fastmri_320p(testlist, transform=loadtransform, disk_cache=\"20241225autoencoder\")\n",
    "# showset = Fastmri_320p(showlist, transform=loadtransform, disk_cache=\"20241225autoencoder\")\n",
    "\n",
    "# trainset.clean_disk_cache()\n",
    "# valset.clean_disk_cache()\n",
    "# testset.clean_disk_cache()\n",
    "# showset.clean_disk_cache()\n",
    "\n",
    "trainset = Phantom(image_size = (320,320))\n",
    "\n",
    "from models.mynn import functional as myf\n",
    "\n",
    "# def data_augmentation(data):\n",
    "#     image = data['pics_image']\n",
    "#     image = myf.complex_to_real(image)\n",
    "#     image = (image - image.mean()) / image.std()\n",
    "#     return {'pics_image':image}\n",
    "\n",
    "# def input_transform(data):\n",
    "#     return data['pics_image'], data['pics_image']\n",
    "\n",
    "# def show_transform(pred, gt, batch = None):\n",
    "#     if batch is None:\n",
    "#         return myf.real_to_complex(pred).abs(), myf.real_to_complex(gt).abs()\n",
    "#     return myf.real_to_complex(pred).abs(), myf.real_to_complex(gt).abs(),myf.real_to_complex(batch).abs()\n",
    "\n",
    "def data_augmentation(data):\n",
    "    image = data\n",
    "    image = (image - image.mean()) / image.std()\n",
    "    return {'pics_image':image}\n",
    "\n",
    "def input_transform(data):\n",
    "    return data['pics_image'], data['pics_image']\n",
    "\n",
    "def show_transform(pred, gt, batch = None):\n",
    "    if batch is None:\n",
    "        return pred, gt\n",
    "    return pred, gt, batch\n",
    "\n",
    "train(model, trainset, model.forward_loss, epochs=300, batch_size=32, device='cuda:5',exp_name=\"phantom-mae\", early_stopping_patience=50,early_stopping_after=20,\n",
    "      epoch_length=640, eval_length_train = 64, eval_length_val = 64, input_transform=input_transform,augmentation_transform=data_augmentation,show_transform=show_transform,\n",
    "      val_dataset=trainset,\n",
    "      imshow_dataset=trainset)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "moco",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}