main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, einsum\n",
    "import torch.nn.functional as F\n",
    "from einops import rearrange, repeat\n",
    "from einops.layers.torch import Rearrange\n",
    "import numpy as np\n",
    "import math\n",
    "    \n",
    "class CoPreNorm(nn.Module):\n",
    "    def __init__(self, dim, fn):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        self.fn = fn\n",
    "    def forward(self, x,x1, **kwargs):\n",
    "        return self.fn(self.norm(x),self.norm(x1), **kwargs)\n",
    "\n",
    "class CoAttention(nn.Module):\n",
    "    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head *  heads\n",
    "        project_out = not (heads == 1 and dim_head == dim)\n",
    "\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\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",
    "            nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "\n",
    "    def forward(self, x,x1):\n",
    "        b, n, _, h = *x.shape, self.heads\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 = h), qkv)\n",
    "        qkv1 = self.to_qkv(x1).chunk(3, dim = -1)\n",
    "        q1, k1, v1 = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv1)\n",
    "        dots = einsum('b h i d, b h j d -> b h i j', q, k1) * self.scale\n",
    "\n",
    "        attn = dots.softmax(dim=-1)\n",
    "        \n",
    "        out = einsum('b h i j, b h j d -> b h i d', attn, v1)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        out =  self.to_out(out)\n",
    "        return out\n",
    "\n",
    "class CoTransformer(nn.Module):\n",
    "    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([])\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                CoPreNorm(dim, CoAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),\n",
    "                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))]))\n",
    "\n",
    "    def forward(self, x,x1):\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x,x1) + x\n",
    "            x = ff(x) + x\n",
    "        return self.norm(x)\n",
    "    \n",
    "class cosa(nn.Module):\n",
    "    def __init__(self, image_size, patch_size, in_channels ,num_frames, depth = 4, heads = 3,num_classes=1, pool = 'cls', \n",
    "                 dim = 8, dim_head = 64, dropout = 0.,emb_dropout = 0., scale_dim = 4, ):\n",
    "        super().__init__()\n",
    "        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "        num_patches = (image_size // patch_size) ** 2\n",
    "        patch_dim = in_channels * patch_size ** 2\n",
    "        self.to_patch_embedding = nn.Sequential(\n",
    "            Rearrange('b t c (h p1) (w p2) -> b (t c) (h w) (p1 p2)', p1 = patch_size, p2 = patch_size),\n",
    "        )\n",
    "        self.pos_embedding = PositionalEmbedding(num_patches*patch_size**2)\n",
    "        self.dropout = nn.Dropout(emb_dropout)\n",
    "        self.transformer = CoTransformer(patch_size**2, depth, heads, dim_head, dim*scale_dim, dropout)\n",
    "        self.img_out = image_size\n",
    "        self.time = num_frames\n",
    "        self.channel = in_channels\n",
    "        self.patch_size = patch_size\n",
    "    def forward(self, x1,x2):\n",
    "        x1,x2 = rearrange(x1, 'b c t n d -> b t c n d'), rearrange(x2, 'b c t n d -> b t c n d')\n",
    "        x1,x2 = self.to_patch_embedding(x1),self.to_patch_embedding(x2)\n",
    "        b, t, n, d = x1.shape\n",
    "        pos1,pos2 = self.pos_embedding(x1,n,d),self.pos_embedding(x2,n,d)\n",
    "        x1 = x1 + pos1\n",
    "        x2 = x2 + pos2\n",
    "        x1 = rearrange(x1, 'b t n d -> (b t) n d')\n",
    "        x2 = rearrange(x2, 'b t n d -> (b t) n d')\n",
    "        x = self.transformer(x1,x2)\n",
    "        patch_height = int(self.img_out/self.patch_size)\n",
    "        x = rearrange(x, '(t c) (ph pw) (p1 p2) -> t c (ph p1) (pw p2)', t = self.time, c = self.channel,\n",
    "                      ph = patch_height, p1 = self.patch_size)\n",
    "        x = x.unsqueeze(0)\n",
    "        x = rearrange(x, 'b t c n d -> b c t n d')\n",
    "        return x\n",
    "    \n",
    "import torch\n",
    "from torch import nn, einsum\n",
    "import torch.nn.functional as F\n",
    "from einops import rearrange, repeat\n",
    "from einops.layers.torch import Rearrange\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "class PreNorm(nn.Module):\n",
    "    def __init__(self, dim, fn):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        self.fn = fn\n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(self.norm(x), **kwargs)\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dim, hidden_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(hidden_dim, dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head *  heads\n",
    "        project_out = not (heads == 1 and dim_head == dim)\n",
    "\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\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",
    "            nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, n, _, h = *x.shape, self.heads\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 = h), qkv)\n",
    "\n",
    "        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale\n",
    "\n",
    "        attn = dots.softmax(dim=-1)\n",
    "\n",
    "        out = einsum('b h i j, b h j d -> b h i d', attn, v)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        out =  self.to_out(out)\n",
    "        return out\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([])\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),\n",
    "                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))\n",
    "            ]))\n",
    "\n",
    "    def forward(self, x):\n",
    "        for attn, ff in self.layers:\n",
    "            x_attn = attn(x) \n",
    "            x = attn(x) + x\n",
    "            x = ff(x) + x\n",
    "        return self.norm(x),self.norm(x_attn)\n",
    "\n",
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, demb):\n",
    "        super(PositionalEmbedding, self).__init__()\n",
    "\n",
    "        self.demb = demb\n",
    "        self.mlp = nn.Sequential(nn.LayerNorm(self.demb),nn.Linear(self.demb, 1))\n",
    "        inv_freq = 1 / (10 ** (torch.arange(0.0, demb, 2.0) / demb))\n",
    "        self.register_buffer('inv_freq', inv_freq)\n",
    "\n",
    "    def forward(self, pos_seq,imgsize,imgdim, bsz=None):\n",
    "        pos_seq = rearrange(pos_seq, 'b t n d -> b t (n d)')\n",
    "        pos_seq = self.mlp(pos_seq)\n",
    "        pos_seq = pos_seq.squeeze(0)\n",
    "        pos_seq = pos_seq.squeeze(1)\n",
    "        #print(pos_seq)\n",
    "        sinusoid_inp = torch.ger(pos_seq, self.inv_freq)\n",
    "        pos_emb = torch.cat([torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)], dim=-1)\n",
    "        if bsz is not None:\n",
    "            pos_emb = pos_emb[:,None,:].expand(-1, bsz, -1)\n",
    "        else:\n",
    "            pos_emb = pos_emb[:,None,:]\n",
    "        #print(pos_emb.shape)\n",
    "        pos_emb = rearrange(pos_emb, 't b (n d) -> b t n d',n=imgsize,d=imgdim)\n",
    "        return pos_emb\n",
    "    \n",
    "class STFH(nn.Module):\n",
    "    def __init__(self, image_size,patch_size, in_channels ,num_frames, depth = 4, heads = 3,num_classes=1, pool = 'cls', \n",
    "                 dim = 8, dim_head = 64, dropout = 0.,emb_dropout = 0., scale_dim = 4, ):\n",
    "        super().__init__()\n",
    "        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "        num_patches = (image_size // patch_size) ** 2\n",
    "        patch_dim = in_channels * patch_size ** 2\n",
    "        self.to_patch_embedding = nn.Sequential(\n",
    "            Rearrange('b t c (h p1) (w p2) -> b (t c) (h w) (p1 p2)', p1 = patch_size, p2 = patch_size),\n",
    "        )\n",
    "        self.pos_embedding = PositionalEmbedding(num_patches*patch_size**2)\n",
    "        self.dropout = nn.Dropout(emb_dropout)\n",
    "        self.spatio_temporal_transformer = Transformer(patch_size**2, depth, heads, dim_head, dim*scale_dim, dropout)\n",
    "        self.img_out = image_size\n",
    "        self.time = num_frames\n",
    "        self.channel = in_channels\n",
    "        self.patch_size = patch_size\n",
    "        \n",
    "       \n",
    "    def forward(self, x):\n",
    "        #print('before p_emb',x.shape)\n",
    "        theta = 5\n",
    "        pool = nn.MaxPool3d(kernel_size=theta, stride=1, padding=(theta - 1) // 2)\n",
    "        x_pool = pool(x)\n",
    "        x = 2*x - x_pool\n",
    "        x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        x = self.to_patch_embedding(x)\n",
    "        b, t, n, d = x.shape\n",
    "        pos = self.pos_embedding(x,n,d)\n",
    "        x = x + pos\n",
    "        x = self.dropout(x)\n",
    "        x = rearrange(x, 'b t n d -> (b t) n d')\n",
    "        x,x_att = self.spatio_temporal_transformer(x)\n",
    "        patch_height = int(self.img_out/self.patch_size)\n",
    "        x = rearrange(x, '(t c) (ph pw) (p1 p2) -> t c (ph p1) (pw p2)', t = self.time, c = self.channel,\n",
    "                      ph = patch_height, p1 = self.patch_size)\n",
    "        x = x.unsqueeze(0)\n",
    "        x = rearrange(x, 'b t c n d -> b c t n d')\n",
    "        x_att = rearrange(x_att, '(t c) (ph pw) (p1 p2) -> t c (ph p1) (pw p2)', t = self.time, c = self.channel,\n",
    "                      ph = patch_height, p1 = self.patch_size)\n",
    "        x_att = x_att.unsqueeze(0)\n",
    "        x_att = rearrange(x_att, 'b t c n d -> b c t n d')\n",
    "        return x,x_att\n",
    "    \n",
    "class Dual_path(nn.Module):\n",
    "    def __init__(self, image_size, in_channels ,num_frames, depth = 4, heads = 3,num_classes=1, pool = 'cls', \n",
    "                 dim = 8, dim_head = 64, dropout = 0.,emb_dropout = 0., scale_dim = 4, ):\n",
    "        super().__init__()\n",
    "        #assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "        #num_patches = (image_size // patch_size) ** 2\n",
    "        #patch_dim = in_channels * patch_size ** 2\n",
    "        self.spatio_stfh = STFH(image_size=image_size, patch_size=image_size//2, in_channels=in_channels ,num_frames=num_frames)\n",
    "        self.temporal_stfh = STFH(image_size=image_size, patch_size=image_size, in_channels=in_channels ,num_frames=num_frames)\n",
    "       \n",
    "    def forward(self, x):\n",
    "        #print('before p_emb',x.shape)\n",
    "        theta = 5\n",
    "        pool = nn.MaxPool3d(kernel_size=theta, stride=1, padding=(theta - 1) // 2)\n",
    "        x_pool = pool(x)\n",
    "        x_bound = 2*x - x_pool\n",
    "        x_bound,x_att_b = self.spatio_stfh(x_bound)\n",
    "        x,x_att_t = self.temporal_stfh(x)\n",
    "        x_croatt = x_att_t*x_att_b\n",
    "        x = x+x_bound+x_croatt\n",
    "        return x\n",
    "\n",
    "class TimeDistributed(nn.Module):\n",
    "    def __init__(self, layer, time_steps):        \n",
    "        super(TimeDistributed, self).__init__()\n",
    "        self.layers = nn.ModuleList([layer for i in range(time_steps)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        batch_size, time_steps, C, H, W = x.size()\n",
    "        output = torch.tensor([]).cuda()\n",
    "        #output = torch.tensor([])\n",
    "        for i in range(time_steps):\n",
    "            output_t = self.layers[i](x[:, i, :, :, :])\n",
    "            output_t  = output_t.unsqueeze(1)\n",
    "            output = torch.cat((output, output_t ), 1)\n",
    "        output = rearrange(output, 'b t c n d -> b c t n d')\n",
    "        return output\n",
    "\n",
    "\n",
    "class EncoderBottleneck3d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1, base_width=64):\n",
    "        super().__init__()\n",
    "\n",
    "        self.downsample = nn.Sequential(\n",
    "            nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=[1,2,2], bias=False),\n",
    "            nn.BatchNorm3d(out_channels)\n",
    "        )\n",
    "\n",
    "        width = int(out_channels * (base_width / 64))\n",
    "\n",
    "        self.conv1 = nn.Conv3d(in_channels, width, kernel_size=1, stride=1, bias=False)\n",
    "        self.norm1 = nn.BatchNorm3d(width)\n",
    "\n",
    "        self.conv2 = nn.Conv3d(width, width, kernel_size=3, stride=[1,2,2], groups=1, padding=1, dilation=1, bias=False)\n",
    "        self.norm2 = nn.BatchNorm3d(width)\n",
    "\n",
    "        self.conv3 = nn.Conv3d(width, out_channels, kernel_size=1, stride=1, bias=False)\n",
    "        self.norm3 = nn.BatchNorm3d(out_channels)\n",
    "\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_down = self.downsample(x)\n",
    "        x = self.conv1(x)\n",
    "        x = self.norm1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.norm2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.norm3(x)\n",
    "        x = x + x_down\n",
    "        x = self.relu(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, img_dim, in_channels, out_channels, head_num, mlp_dim, block_num, patch_dim,seq_frame):\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=7, stride=[1,2,2], padding=3, bias=False)\n",
    "        self.norm1 = nn.BatchNorm3d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "        self.encoder1 = EncoderBottleneck3d(out_channels, out_channels * 2, stride=2)\n",
    "        self.encoder2 = EncoderBottleneck3d(out_channels * 2, out_channels * 4, stride=2)\n",
    "        self.encoder3 = EncoderBottleneck3d(out_channels * 4, out_channels * 8, stride=2)\n",
    "\n",
    "        self.conv2 = nn.Conv3d(out_channels * 8, out_channels * 4, kernel_size=3, stride=1, padding=1)\n",
    "        self.norm2 = nn.BatchNorm3d(out_channels * 4)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = rearrange(x, 'b t c n d -> b c t n d')\n",
    "        x = self.conv1(x)\n",
    "        x = self.norm1(x)\n",
    "        x1 = self.relu(x)\n",
    "        x2 = self.encoder1(x1)\n",
    "        x3 = self.encoder2(x2)\n",
    "        x = self.encoder3(x3)\n",
    "        #x = rearrange(x, \"b t (h w) c ->b t c h w\",h = 8, w=8)\n",
    "\n",
    "        x = self.conv2(x)\n",
    "        x = self.norm2(x)\n",
    "        x = self.relu(x)\n",
    "        #x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        #x1_out = rearrange(x1_out, 'b c t n d -> b t c n d')\n",
    "        #x2_out = rearrange(x2_out, 'b c t n d -> b t c n d')\n",
    "        #x3_out = rearrange(x3_out, 'b c t n d -> b t c n d')\n",
    "        return x, x1, x2, x3\n",
    "\n",
    "class DecoderBottleneck3d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels,seq_frame, scale_factor=2):\n",
    "        super().__init__()\n",
    "\n",
    "        self.upsample = TimeDistributed(nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=True), time_steps = seq_frame)\n",
    "        self.layer = nn.Sequential(\n",
    "            nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
    "            nn.BatchNorm3d(out_channels),\n",
    "            nn.ReLU(inplace=False),\n",
    "            nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
    "            nn.BatchNorm3d(out_channels),\n",
    "            nn.ReLU(inplace=False))\n",
    "\n",
    "    def forward(self, x, x_concat=None):\n",
    "        \n",
    "        x = self.upsample(x)\n",
    "\n",
    "        if x_concat is not None:\n",
    "            x = torch.cat([x_concat, x], dim=1)\n",
    "        x = self.layer(x)\n",
    "        return x\n",
    "\n",
    "class dconv3d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels,d_rate, stride=1, base_width=64):\n",
    "        super().__init__()\n",
    "\n",
    "        width = int(out_channels * (base_width / 64))\n",
    "\n",
    "        self.conv1 = nn.Conv3d(in_channels, width, kernel_size=1, stride=1, bias=False)\n",
    "        self.norm1 = nn.BatchNorm3d(width)\n",
    "\n",
    "        self.conv2 = nn.Conv3d(width, width, kernel_size=3, stride=[1,1,1], groups=1, padding=d_rate, dilation=d_rate, bias=False)\n",
    "        self.norm2 = nn.BatchNorm3d(width)\n",
    "\n",
    "        self.conv3 = nn.Conv3d(width, out_channels, kernel_size=1, stride=1, bias=False)\n",
    "        self.norm3 = nn.BatchNorm3d(out_channels)\n",
    "\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.norm1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.norm2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.norm3(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "class task_specific_decoder(nn.Module):\n",
    "    def __init__(self, out_channels, class_num, drop_rate,seq_frame):\n",
    "        super().__init__()\n",
    "        self.d_rate = drop_rate\n",
    "        self.decoder1 = DecoderBottleneck3d(out_channels * 8, out_channels * 2,seq_frame)\n",
    "        self.decoder2 = DecoderBottleneck3d(out_channels * 4, out_channels,seq_frame)\n",
    "        self.dropout = TimeDistributed(torch.nn.Dropout(p=self.d_rate),time_steps = seq_frame)\n",
    "        \n",
    "    def forward(self, x, x2, x3):\n",
    "        x = self.decoder1(x, x3)\n",
    "        x = self.dropout(x)\n",
    "        x = self.decoder2(x, x2)\n",
    "        #x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        return x\n",
    "\n",
    "    \n",
    "class segment_head(nn.Module):\n",
    "    def __init__(self, out_channels, class_num, drop_rate,seq_frame):\n",
    "        super().__init__()\n",
    "        self.decoder3 = DecoderBottleneck3d(out_channels * 2, int(out_channels * 1 / 2),seq_frame)\n",
    "        self.decoder4 = DecoderBottleneck3d(int(out_channels * 1 / 2), int(out_channels * 1 / 8),seq_frame)\n",
    "        self.conv1 = nn.Conv3d(int(out_channels * 1 / 8), class_num, kernel_size=1)\n",
    "        \n",
    "    def forward(self, x, x1):\n",
    "        x = self.decoder3(x, x1)\n",
    "        x = self.decoder4(x)\n",
    "        x = self.conv1(x)\n",
    "        x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        return x\n",
    "\n",
    "class point_head(nn.Module):\n",
    "    def __init__(self, out_channels, class_num, seq_frame,height,weight):\n",
    "        super().__init__()\n",
    "        self.seq_frame = seq_frame\n",
    "        self.height = height\n",
    "        self.weight = weight\n",
    "        self.decoder3 = DecoderBottleneck3d(out_channels * 2, int(out_channels * 1 / 2),seq_frame)\n",
    "        self.decoder4 = DecoderBottleneck3d(int(out_channels * 1 / 2), int(out_channels * 1 / 8),seq_frame)\n",
    "        self.conv1 = nn.Conv3d(int(out_channels * 1 / 8), class_num, kernel_size=1)\n",
    "        self.conv2 = nn.Conv3d(class_num,1, kernel_size=1)\n",
    "        self.mlp_out = nn.Sequential(\n",
    "            nn.LayerNorm(self.height*self.weight),\n",
    "            nn.Linear(self.height*self.weight, 2))\n",
    "        \n",
    "    def forward(self, x, x1):\n",
    "        x = self.decoder3(x, x1)\n",
    "        x = self.decoder4(x)\n",
    "        x = self.conv1(x)\n",
    "        x_map = self.conv2(x)\n",
    "        x_map = rearrange(x_map, 'b c t n d -> b t c n d')\n",
    "        x = rearrange(x, 'b c t n d -> b t c n d')\n",
    "        x_pot = rearrange(x, 'b t c n d -> b t c (n d)')\n",
    "        x_pot = self.mlp_out(x_pot)\n",
    "        return x_pot,x_map\n",
    "\n",
    "# class classifi_head(nn.Module):\n",
    "#     def __init__(self, out_channels, class_num,seq_frame):\n",
    "#         super().__init__()\n",
    "#         self.out_ch = out_channels\n",
    "#         self.class_num = class_num\n",
    "#         self.seq_frame = seq_frame\n",
    "#         self.avgpool = nn.AvgPool2d(32, stride=1)\n",
    "#         self.fc = nn.Linear(self.out_ch,self.class_num)\n",
    "#         for m in self.modules():\n",
    "#             if isinstance(m, nn.Conv2d):\n",
    "#                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "#             elif isinstance(m, nn.BatchNorm2d):\n",
    "#                 nn.init.constant_(m.weight, 1)\n",
    "#                 nn.init.constant_(m.bias, 0)\n",
    "                \n",
    "#     def forward(self,x):\n",
    "#         x = x[:,:,0,:,:]\n",
    "#         x = self.avgpool(x)\n",
    "#         x = x.view(x.size(0), -1)\n",
    "#         x = self.fc(x)\n",
    "#         return x\n",
    "    \n",
    "class frame_head(nn.Module):\n",
    "    def __init__(self, out_channels, class_num,seq_frame):\n",
    "        super().__init__()\n",
    "        self.out_ch = out_channels\n",
    "        self.class_num = class_num\n",
    "        self.seq_frame = seq_frame\n",
    "        self.dconv3d2 = dconv3d(self.out_ch*1, self.out_ch*2, d_rate = 3, stride=2)\n",
    "        self.avgpool2d2  = TimeDistributed(nn.AdaptiveAvgPool2d(8), time_steps = seq_frame)\n",
    "        self.dconv3d3 = dconv3d(self.out_ch*2, self.out_ch*4, d_rate = 2, stride=2)\n",
    "        self.avgpool2d3  = TimeDistributed(nn.AdaptiveAvgPool2d(2), time_steps = seq_frame)\n",
    "        self.dconv3d4 = dconv3d(self.out_ch*4, self.out_ch*8, d_rate = 1, stride=2)\n",
    "        self.avgpool2d4  = TimeDistributed(nn.AdaptiveAvgPool2d(1), time_steps = seq_frame)\n",
    "        self.mlp_out = nn.Sequential(\n",
    "            nn.LayerNorm(self.out_ch*8),\n",
    "            nn.Linear(self.out_ch*8, self.class_num))\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "                \n",
    "    def forward(self,x):\n",
    "        #print(x.shape)\n",
    "        x = self.dconv3d2(x)\n",
    "        #print(x.shape)\n",
    "        x = self.avgpool2d2(x)\n",
    "        x = self.dconv3d3(x)\n",
    "        x = self.avgpool2d3(x)\n",
    "        x = self.dconv3d4(x)\n",
    "        x = self.avgpool2d4(x)\n",
    "        #print(x.shape)\n",
    "        x = rearrange(x, 'b t c n d -> (c n d)(b t)')\n",
    "        x_reg = self.mlp_out(x)\n",
    "        x_reg.squeeze()\n",
    "        #x_reg = rearrange(x_reg, 'n s -> s n')\n",
    "        return x_reg\n",
    "    \n",
    "class TaskAttention(nn.Module):\n",
    "    def __init__(self, in_ch, ratio=16):\n",
    "        super(TaskAttention, self).__init__()\n",
    "        self.avg_pool = nn.AdaptiveAvgPool3d(1)\n",
    "        self.max_pool = nn.AdaptiveMaxPool3d(1)\n",
    "           \n",
    "        self.fc = nn.Sequential(nn.Conv3d(in_ch*2, in_ch//16, kernel_size=1),\n",
    "                               nn.ReLU(),\n",
    "                               nn.Conv3d(in_ch // 16, in_ch*2, kernel_size=1))\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        self.conv3 = nn.Conv3d(in_ch*2, in_ch, kernel_size=1, stride=1, bias=False)\n",
    "        self.norm3 = nn.BatchNorm3d(in_ch)\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "\n",
    "    def forward(self, x1,x2):\n",
    "        x = torch.cat([x1, x2], dim=1)\n",
    "        avg_out = self.fc(self.avg_pool(x))\n",
    "        max_out = self.fc(self.max_pool(x))\n",
    "        out = avg_out + max_out\n",
    "        x = self.sigmoid(out)+x\n",
    "        x = self.conv3(x)\n",
    "        x = self.norm3(x)\n",
    "        x = self.relu(x)\n",
    "        return x\n",
    "    \n",
    "class pattern_strcture(nn.Module):\n",
    "    def __init__(self, image_size, in_channels ,num_frames, depth = 4, heads = 3,num_classes=1, pool = 'cls', \n",
    "                 dim = 8, dim_head = 64, dropout = 0.,emb_dropout = 0., scale_dim = 4, ):\n",
    "        super().__init__()\n",
    "        self.spatio_stfh = STFH(image_size=image_size, patch_size=image_size//4, in_channels=in_channels ,num_frames=num_frames)\n",
    "        self.ta = TaskAttention(in_channels)\n",
    "        size = 32\n",
    "        self.cosa = cosa(image_size=size,patch_size=int(size//2), in_channels=32 ,num_frames=30)\n",
    "       \n",
    "    def forward(self,x1,x2,x3):\n",
    "        x1_sa,x_att_b1 = self.spatio_stfh(x1)\n",
    "        #print(self.ta(x2,x3,x4).shape)\n",
    "        x2_sa,x_att_b2 = self.spatio_stfh(x2)\n",
    "        #print(x1.shape,x2.shape,x3.shape,x4.shape)\n",
    "        x3_sa,x_att_b3 = self.spatio_stfh(x3)\n",
    "        # x4_sa,x_att_b4 = self.spatio_stfh(x4)\n",
    "        beta = 0.1\n",
    "        # x1 = beta*self.ta(x2_sa,x3_sa)+(1-beta)*x1_sa\n",
    "        # x2 = beta*self.ta(x1_sa,x3_sa)+(1-beta)*x2_sa\n",
    "        # x3 = beta*self.ta(x2_sa,x1_sa)+(1-beta)*x3_sa\n",
    "        #print(self.ta(x2_sa,x3_sa).shape,x1_sa.shape)\n",
    "        x1 = self.cosa(beta*self.ta(x2_sa,x3_sa),(1-beta)*x1_sa)\n",
    "        x2 = self.cosa(beta*self.ta(x1_sa,x3_sa),(1-beta)*x2_sa)\n",
    "        x3 = self.cosa(beta*self.ta(x2_sa,x1_sa),(1-beta)*x3_sa)\n",
    "        return x1,x2,x3\n",
    "\n",
    "seq_frame = 30\n",
    "\n",
    "class Multivit_net(nn.Module):\n",
    "    def __init__(self, img_dim, in_channels, out_channels, head_num, mlp_dim, block_num, patch_dim, class_num, drop_rate,seq_frame\n",
    "                 ,mode,height,weight):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = Encoder(img_dim, in_channels, out_channels,\n",
    "                               head_num, mlp_dim, block_num, patch_dim,seq_frame)\n",
    "        self.seg_decoder = task_specific_decoder(out_channels, class_num, drop_rate,seq_frame)\n",
    "        self.pot_decoder = task_specific_decoder(out_channels, class_num, drop_rate,seq_frame)\n",
    "        self.frm_decoder = task_specific_decoder(out_channels, class_num, drop_rate,seq_frame)\n",
    "        # self.cls_decoder = task_specific_decoder(out_channels, class_num, drop_rate,seq_frame)\n",
    "        self.seg_head = segment_head(out_channels, class_num, drop_rate,seq_frame)\n",
    "        # self.cls_head = classifi_head(out_channels, class_num=4,seq_frame=seq_frame)\n",
    "        self.pot_head = point_head(out_channels, class_num=4,seq_frame=seq_frame,height=height,weight=weight)\n",
    "        self.frm_head = frame_head(out_channels, class_num=3,seq_frame=seq_frame)\n",
    "        self.mode = mode\n",
    "        self.stps = pattern_strcture(image_size=int(img_dim/4), in_channels=out_channels ,num_frames=seq_frame)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x, x1, x2, x3 = self.encoder(x)\n",
    "        if self.mode == 'seg':\n",
    "            x_seg = self.seg_decoder(x, x2, x3)\n",
    "            x_seg = self.seg_head(x_seg,x1)\n",
    "            return x_seg\n",
    "        elif self.mode == 'pot':\n",
    "            x_pot = self.pot_decoder(x, x2, x3)\n",
    "            x_pot = self.pot_head(x_pot,x1)\n",
    "            return x_pot\n",
    "        elif self.mode == 'frm':\n",
    "            x_frm = self.frm_decoder(x, x2, x3)\n",
    "            x_frm = self.frm_head(x_frm)\n",
    "            return x_frm\n",
    "        # elif self.mode == 'cls':\n",
    "        #     x_cls = self.cls_decoder(x, x2, x3)\n",
    "        #     x_cls = self.cls_head(x_cls)\n",
    "            return x_cls\n",
    "        elif self.mode == 'mtl':\n",
    "            x_seg = self.seg_decoder(x, x2, x3)\n",
    "            x_pot = self.pot_decoder(x, x2, x3)\n",
    "            x_frm = self.frm_decoder(x, x2, x3)\n",
    "            x_seg,x_pot,x_frm = self.stps(x_seg,x_pot,x_frm)\n",
    "            x_seg = self.seg_head(x_seg,x1)\n",
    "            x_pot,x_potmap = self.pot_head(x_pot,x1)\n",
    "            x_frm = self.frm_head(x_frm)\n",
    "        return x_pot,x_potmap,x_frm,x_seg  #x_pot=(1,30,4,2),x_frm(30,3),x_seg(1,30,1,128,128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbbedingSpace(nn.Module):\n",
    "    def __init__(self, dim, frame ,mode,class_num=1):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.frame = frame\n",
    "        self.mode = mode\n",
    "        self.class_num = class_num\n",
    "        self.mlp_out = nn.Sequential(\n",
    "            nn.LayerNorm(self.dim),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(self.dim, self.class_num))\n",
    "    def forward(self, x):\n",
    "        if self.mode == 'seg':\n",
    "            x = rearrange(x, 'b t c n d -> (b t) (c n d)')\n",
    "        elif self.mode == 'pot':\n",
    "            x = rearrange(x, 'b t n d -> (b t) (n d)')\n",
    "        elif self.mode == 'frm':\n",
    "            x = x\n",
    "        x = self.mlp_out(x)\n",
    "        return x\n",
    "    \n",
    "class Embbeding_net(nn.Module):\n",
    "    def __init__(self, dim, frame,device,class_num=1):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.frame = frame\n",
    "        self.class_num = class_num\n",
    "        self.device = device\n",
    "        self.pot_emb = EmbbedingSpace(8,frame,'pot').to(self.device)\n",
    "        self.frm_emb = EmbbedingSpace(3,frame,'frm').to(self.device)\n",
    "        self.seg_emb = EmbbedingSpace(128*128,frame,'seg').to(self.device)\n",
    "    def forward(self, x_pot,x_frm,x_seg):\n",
    "        x_frm = one_hot(x_frm.long(),3).to(device)\n",
    "        x_frm = tr_oh(x_frm,30,3)\n",
    "        xp_emb = self.pot_emb(x_pot.to(torch.float32)).to(self.device)\n",
    "        xf_emb = self.frm_emb(x_frm.to(torch.float32)).to(self.device)\n",
    "        xs_emb = self.seg_emb(x_seg.to(torch.float32)).to(self.device)\n",
    "        return xp_emb,xf_emb,xs_emb\n",
    "\n",
    "def one_hot(label, n_classes, requires_grad=True):\n",
    "    \"\"\"Return One Hot Label\"\"\"\n",
    "    divce = label.device\n",
    "    one_hot_label = torch.eye(n_classes, requires_grad=requires_grad)[label]\n",
    "    # one_hot_label = one_hot_label.transpose(1, 3).transpose(2, 3)\n",
    "    return one_hot_label\n",
    "    \n",
    "def trend(frm, f, k):\n",
    "    frm1 = torch.zeros_like(frm)\n",
    "    for i in range(f-k):\n",
    "        for kt in range(k):\n",
    "            # print(k-kt,i+kt)\n",
    "            frm1[i] = frm1[i]+(k-kt)*frm[i+kt]\n",
    "    for i in range(k):\n",
    "        i = i+f-k\n",
    "        for kt in range(k):\n",
    "            # print(k-kt,k-(f-i)+kt)\n",
    "            frm1[i] = frm1[i]+(k-kt)*frm[k-(f-i)+kt]\n",
    "    return frm1/k\n",
    "\n",
    "def tr_oh(frm, f, k):\n",
    "    frm[:,0],frm[:,1],frm[:,2] = trend(frm[:,0],f,k),trend(frm[:,1],f,k),trend(frm[:,2],f,k)\n",
    "    return frm\n",
    "\n",
    "def cosine_loss(input,target):\n",
    "    # print(input.squeeze(1).shape,target.shape)\n",
    "    input,target = input.squeeze(1),target.squeeze(1)\n",
    "    sim = torch.cosine_similarity(input.unsqueeze(0),target.unsqueeze(0))\n",
    "    return 1-sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.ndimage import morphology\n",
    "def coefficients(gt, pred, smooth=1e-12):\n",
    "    pred = torch.sigmoid(pred)\n",
    "    pred = torch.gt(pred, 0.5)\n",
    "    pred = pred.type(torch.float32)\n",
    "    intersection = torch.sum(gt * pred)\n",
    "    gt, pred = torch.sum(gt), torch.sum(pred)\n",
    "    union = gt + pred - intersection\n",
    "\n",
    "    precision = intersection / (pred + smooth)\n",
    "    recall = intersection / (gt + smooth)\n",
    "\n",
    "    beta_square = 0.3\n",
    "    f_beta_coeff = (1 + beta_square) * precision * recall / (beta_square * precision + recall + smooth)\n",
    "    dice_coeff = (2. * intersection) / (union + intersection + smooth)\n",
    "    jaccard_coeff = intersection / (union + smooth)\n",
    "    return dice_coeff, jaccard_coeff, f_beta_coeff\n",
    "\n",
    "def acc(pred,gt):\n",
    "    right,error = 0,0\n",
    "    right_es,error_es = 0,0\n",
    "    right_ed,error_ed = 0,0\n",
    "    for i in range(30):\n",
    "        if float(gt[i:i+1]) == 0:\n",
    "            if float(gt[i:i+1]) == float(torch.max(pred[i:i+1], 1)[1]):\n",
    "                    right_es = right_es+1\n",
    "            else:\n",
    "                    error_es = error_es+1\n",
    "        elif float(gt[i:i+1]) == 1:\n",
    "            if float(gt[i:i+1]) == float(torch.max(pred[i:i+1], 1)[1]):\n",
    "                    right = right+1\n",
    "            else:\n",
    "                    error = error+1\n",
    "        elif float(gt[i:i+1]) == 2:\n",
    "            if float(gt[i:i+1]) == float(torch.max(pred[i:i+1], 1)[1]):\n",
    "                    right_ed = right_ed+1\n",
    "            else:\n",
    "                    error_ed = error_ed+1\n",
    "    return right/(right+error),right_es/(right_es+error_es+0.002),right_ed/(right_ed+error_ed+0.002)\n",
    "    \n",
    "\n",
    "def get_hausdorff(gt, pred, sampling=0.3, connectivity=1):\n",
    "    pred = torch.sigmoid(pred)\n",
    "    pred = torch.gt(pred, 0.5)\n",
    "    input1 = gt\n",
    "    input2 = pred\n",
    "    input1 = np.array(input1.cpu().clone()) \n",
    "    input2 = np.array(input2.cpu().clone()) \n",
    "    input_1 = np.atleast_1d(input1.astype(np.bool))\n",
    "    input_2 = np.atleast_1d(input2.astype(np.bool))\n",
    "\n",
    "    conn = morphology.generate_binary_structure(input_1.ndim, connectivity)\n",
    "\n",
    "    S = input_1 ^ morphology.binary_erosion(input_1, conn)\n",
    "    Sprime = input_2 ^ morphology.binary_erosion(input_2, conn)\n",
    "\n",
    "    dta = morphology.distance_transform_edt(~S, sampling)\n",
    "    dtb = morphology.distance_transform_edt(~Sprime, sampling)\n",
    "\n",
    "    sds = np.concatenate([np.ravel(dta[Sprime != 0]), np.ravel(dtb[S != 0])])\n",
    "    hausdorff_distance = sds.max()\n",
    "    mean_abs_distance = np.abs(sds).mean()\n",
    "    return hausdorff_distance, mean_abs_distance\n",
    "    \n",
    "def seg_loss(y_pred,y_true):\n",
    "    y_pred = torch.sigmoid(y_pred)\n",
    "    smooth       = 1e-12\n",
    "    y_true_back  = 1 - y_true\n",
    "    y_pred_back  = 1 - y_pred\n",
    "    alpha        = 1 / (torch.pow(torch.sum(y_true), 2) + smooth)\n",
    "    beta         = 1 / (torch.pow(torch.sum(y_true_back), 2) + smooth)\n",
    "    numerater    = alpha * torch.sum(y_true * y_pred) + beta * torch.sum(y_true_back * y_pred_back)\n",
    "    denominator  = alpha * torch.sum(y_true + y_pred) + beta * torch.sum(y_true_back + y_pred_back)\n",
    "    dice_loss    = 1 - (2. * numerater) / (denominator + smooth)\n",
    "    mae_loss     = torch.mean(torch.log(1 + torch.exp(torch.abs(y_pred - y_true))))\n",
    "    w            = (img_size * img_size - torch.sum(y_pred)) / (torch.sum(y_pred) + smooth)\n",
    "    key_w        = 0.003\n",
    "    crossentropy = - torch.mean(key_w * w * y_true * torch.log(y_pred + smooth) + y_true_back * torch.log(y_pred_back + smooth))\n",
    "    #print(crossentropy)\n",
    "    return crossentropy + dice_loss + mae_loss\n",
    "\n",
    "def one_hot(label, n_classes, requires_grad=True):\n",
    "    \"\"\"Return One Hot Label\"\"\"\n",
    "    divce = label.device\n",
    "    one_hot_label = torch.eye(n_classes, device=device, requires_grad=requires_grad)[label]\n",
    "    # one_hot_label = one_hot_label.transpose(1, 3).transpose(2, 3)\n",
    "    return one_hot_label\n",
    "\n",
    "def boundary_cos_loss(gt , pred):\n",
    "    pred = torch.sigmoid(pred)\n",
    "    #gt dimension (B,T,C,H,W)\n",
    "    b,t,c,h,w = gt.shape\n",
    "    for i in range(t):\n",
    "        gt_frame = gt[0,i:i+1,:,:,:]\n",
    "        pred_frame = pred[0,i:i+1,:,:,:]\n",
    "        theta0 = 3\n",
    "        gt_cont = F.max_pool2d(1 - gt_frame, kernel_size=theta0, stride=1, padding=(theta0 - 1) // 2)\n",
    "        gt_cont -= 1 - gt_frame\n",
    "        pred_cont = F.max_pool2d(1 - pred_frame, kernel_size=theta0, stride=1, padding=(theta0 - 1) // 2)\n",
    "        pred_cont -= 1 - pred_frame\n",
    "        sim = torch.cosine_similarity(gt_cont.squeeze(0).squeeze(0),pred_cont.squeeze(0).squeeze(0))\n",
    "        sim_norm = torch.sum(sim)/(h)\n",
    "        sim_loss = 1 - sim_norm*2\n",
    "        if i == 0:\n",
    "            loss = sim_loss\n",
    "        else:\n",
    "            loss = loss + sim_loss\n",
    "    return loss / t\n",
    "\n",
    "def cos_sim_loss(gt , pred):\n",
    "    pred = torch.sigmoid(pred)\n",
    "    #gt dimension (B,T,C,H,W)\n",
    "    b,t,c,h,w = gt.shape\n",
    "    for i in range(t):\n",
    "        gt_frame = gt[0,i:i+1,:,:,:]\n",
    "        pred_frame = pred[0,i:i+1,:,:,:]\n",
    "        theta0 = 3\n",
    "        sim = torch.cosine_similarity(gt_frame.squeeze(0).squeeze(0),pred_frame.squeeze(0).squeeze(0))\n",
    "        sim_norm = torch.sum(sim)/(h)\n",
    "        sim_loss = 1 - sim_norm*2\n",
    "        if i == 0:\n",
    "            loss = sim_loss\n",
    "        else:\n",
    "            loss = loss + sim_loss\n",
    "    return loss / t\n",
    "\n",
    "def corr_loss(pred,gt):\n",
    "    pred_lva,pred_mvd,pred_lvd = pred[0,:],pred[1,:],pred[2,:]\n",
    "    gt_lva,  gt_mvd,  gt_lvd   = gt[0,:],  gt[1,:],  gt[2,:]\n",
    "    pred,gt = pred_lva, gt_lva\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_lva = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / ((\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2)))+1e-12)\n",
    "    pred,gt = pred_mvd, gt_mvd\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_mvd = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / ((\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2)))+1e-12)\n",
    "    pred,gt = pred_lvd, gt_lvd\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_lvd = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / ((\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2)))+1e-12)\n",
    "    #print('corr1:',corr_lva,'corr2:',corr_mvd,'corr3:',corr_lvd)\n",
    "    corr = corr_lva+2*corr_mvd+2*corr_lvd+1e-12\n",
    "    return 5-corr\n",
    "\n",
    "def mae_loss(pred,gt):\n",
    "    mae1 = torch.mean(torch.abs(pred[0,:]-gt[0,:]))\n",
    "    mae2 = torch.mean(torch.abs(pred[1,:]-gt[1,:]))\n",
    "    mae3 = torch.mean(torch.abs(pred[2,:]-gt[2,:]))\n",
    "    mae = mae2+mae3*2\n",
    "    #mae = torch.mean((pred-gt)* torch.tanh(pred-gt))\n",
    "    #logcosh = torch.mean(torch.log(torch.cosh((pred-gt) + 1e-12)))\n",
    "    mae_mean = torch.mean(torch.abs(torch.mean(pred)-torch.mean(gt)))\n",
    "    return mae\n",
    "\n",
    "def mae_point(pred,gt):\n",
    "    mae = torch.mean(torch.abs(pred-gt))\n",
    "    #mae = torch.mean((pred-gt)* torch.tanh(pred-gt))\n",
    "    logcosh = torch.mean(torch.log(torch.cosh((pred-gt) + 1e-12)))\n",
    "    return mae\n",
    "\n",
    "def mae_cal(pred,gt):\n",
    "    mae1 = torch.mean(torch.abs(pred[0,:]-gt[0,:]))\n",
    "    mae2 = torch.mean(torch.abs(pred[1,:]-gt[1,:]))\n",
    "    mae3 = torch.mean(torch.abs(pred[2,:]-gt[2,:]))\n",
    "    mae = mae2+mae3*2\n",
    "    #mae = torch.mean((pred-gt)* torch.tanh(pred-gt))\n",
    "    logcosh = torch.mean(torch.log(torch.cosh((pred-gt) + 1e-12)))\n",
    "    return mae\n",
    "\n",
    "\n",
    "def person_corr(pred, gt):#皮尔森相关系数\n",
    "    pred_lva,pred_mvd,pred_lvd = pred[0,:],pred[1,:],pred[2,:]\n",
    "    gt_lva,  gt_mvd,  gt_lvd   = gt[0,:],  gt[1,:],  gt[2,:]\n",
    "    pred,gt = pred_lva, gt_lva\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_lva = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / (\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2))+1e-12)\n",
    "    pred,gt = pred_mvd, gt_mvd\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_mvd = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / (\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2))+1e-12)\n",
    "    pred,gt = pred_lvd, gt_lvd\n",
    "    pred_mean, gt_mean = torch.mean(pred), torch.mean(gt)\n",
    "    corr_lvd = (torch.sum((pred - pred_mean) * (gt - gt_mean))) / (\n",
    "                torch.sqrt(torch.sum((pred - pred_mean) ** 2)) * torch.sqrt(torch.sum((gt - gt_mean) ** 2))+1e-12)\n",
    "    return corr_lva, corr_mvd,corr_lvd\n",
    "\n",
    "\n",
    "def point2linear(point):\n",
    "    mvdt = []\n",
    "    lvdt = []\n",
    "    #point:1,30,4,2\n",
    "    for t in range(30):\n",
    "        mvd = ((point[t,0,0]-point[t,1,0])**2+(point[t,0,1]-point[t,1,1])**2)**(1/2)\n",
    "        lvd = ((point[t,2,0]-point[t,3,0])**2+(point[t,2,1]-point[t,3,1])**2)**(1/2)\n",
    "        mvdt.append(mvd)\n",
    "        lvdt.append(lvd)\n",
    "        if t>0:\n",
    "            if mvdt[t]<mvdt[t-1]*0.6:\n",
    "                mvdt[t] = mvdt[t-1]\n",
    "            # elif mvdt[t]>mvdt[t-1]*1.8:\n",
    "            #     mvdt[t] = mvdt[t-1]\n",
    "    return mvdt,lvdt\n",
    "\n",
    "def pot2index(srcnpy):\n",
    "    mvd_gt,lvd_gt  = [],[]\n",
    "    # print(srcnpy.shape)#1,30,4,2\n",
    "    # for i in range(srcnpy.shape[0]):\n",
    "    mvd1,lvd1 = point2linear(srcnpy[0])\n",
    "    mvd_gt.append(mvd1)\n",
    "    lvd_gt.append(lvd1)\n",
    "    #print(lvd,lvd1)\n",
    "    lvd_gt = np.array(lvd_gt)[:,:,np.newaxis]\n",
    "    mvd_gt = np.array(mvd_gt)[:,:,np.newaxis]\n",
    "    index  = np.concatenate((lvd_gt, mvd_gt), axis=2)\n",
    "    return index\n",
    "\n",
    "def mae_div(pred,gt):\n",
    "    pred = pred.cpu().detach().numpy()\n",
    "    gt = gt.cpu().detach().numpy()\n",
    "    pred_ind = pot2index(pred)\n",
    "    gt_ind   = pot2index(gt)\n",
    "    # print(gt_ind.shape)\n",
    "    mae1 = np.mean(np.abs(pred_ind[:,:,0]-gt_ind[:,:,0]))\n",
    "    mae2 = np.mean(np.abs(pred_ind[:,:,1]-gt_ind[:,:,1]))\n",
    "    mae = mae1+mae2\n",
    "    #mae = torch.mean((pred-gt)* torch.tanh(pred-gt))\n",
    "    # logcosh = torch.mean(torch.log(torch.cosh((pred-gt) + 1e-12)))\n",
    "    return mae1,mae2\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def l2_norm(seg , pot):\n",
    "    seg = torch.sigmoid(seg)\n",
    "    pot = torch.sigmoid(pot)\n",
    "    #gt dimension (B,T,C,H,W)\n",
    "    b,t,c,h,w = seg.shape\n",
    "    for i in range(t):\n",
    "        seg_frame = seg[0,i:i+1,:,:,:]\n",
    "        pot_frame = pot[0,i:i+1,:,:,:]\n",
    "        theta0 = 9\n",
    "        seg_cont = F.max_pool2d(1 - seg_frame, kernel_size=theta0, stride=1, padding=(theta0 - 1) // 2)\n",
    "        seg_cont -= 1 - seg_frame\n",
    "        pot_cont = F.max_pool2d(1 - pot_frame, kernel_size=theta0, stride=1, padding=(theta0 - 1) // 2)\n",
    "        pot_cont -= 1 - pot_frame\n",
    "        potmine = (pot_frame.squeeze(0).squeeze(0).to(torch.float32)/3-\n",
    "                                                                    seg_cont.squeeze(0).squeeze(0).to(torch.float32))\n",
    "        potmine = 1-torch.gt(potmine,0).to(torch.float32)\n",
    "        potsum = torch.sum(potmine)\n",
    "        segmine = torch.gt(seg_cont,0).to(torch.float32)\n",
    "        segsum = torch.sum(segmine)\n",
    "    return potsum/segsum\n",
    "\n",
    "def cal_a(mask):\n",
    "    lab = torch.reshape(mask,(mask.shape[0],30,1,128*128))\n",
    "    lab_line = torch.sum(lab,dim=3)\n",
    "    lab_line1 = lab_line - torch.mean(lab_line)\n",
    "    lab_line1 = lab_line1/torch.std(lab_line)\n",
    "    return lab_line1\n",
    "\n",
    "\n",
    "def cosine_loss(input,target):\n",
    "    # print(input.squeeze(1).shape,target.shape)\n",
    "    input,target = input.squeeze(1),target.squeeze(1)\n",
    "    sim = torch.cosine_similarity(input.unsqueeze(0),target.unsqueeze(0))\n",
    "    return 1-sim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#traning set\n",
    "import torch\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import OrdinalEncoder\n",
    "import numpy.matlib\n",
    "def landmark(center_x,center_y,IMAGE_HEIGHT, IMAGE_WIDTH):\n",
    "    R = np.sqrt(2**2 + 2**2)\n",
    "    Gauss_map = np.zeros((IMAGE_HEIGHT, IMAGE_WIDTH))\n",
    "    mask_x = np.matlib.repmat(center_x, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "    mask_y = np.matlib.repmat(center_y, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "    x1 = np.arange(IMAGE_WIDTH)\n",
    "    x_map = np.matlib.repmat(x1, IMAGE_HEIGHT, 1)\n",
    "    y1 = np.arange(IMAGE_HEIGHT)\n",
    "    y_map = np.matlib.repmat(y1, IMAGE_WIDTH, 1)\n",
    "    y_map = np.transpose(y_map)\n",
    "    Gauss_map = np.sqrt((x_map-mask_x)**2+(y_map-mask_y)**2)\n",
    "    Gauss_map = np.exp(-0.5*Gauss_map/R)\n",
    "    return Gauss_map\n",
    "\n",
    "def locmap(pot):\n",
    "    gauss_batch = []\n",
    "    for i in range(0,pot.shape[0]):\n",
    "        gauss_tp = []\n",
    "        for j in range(0,pot.shape[1]):\n",
    "            g_map1 = landmark(pot[i, j, 0, 0],pot[i, j, 0, 1],128,128)\n",
    "            g_map2 = landmark(pot[i, j, 1, 0],pot[i, j, 1, 1],128,128)\n",
    "            g_map3 = landmark(pot[i, j, 2, 0],pot[i, j, 2, 1],128,128)\n",
    "            g_map4 = landmark(pot[i, j, 3, 0],pot[i, j, 3, 1],128,128)\n",
    "            Gauss_map = (g_map1+g_map2+g_map3+g_map4)/4\n",
    "            gauss_tp.append(Gauss_map)\n",
    "        gauss_batch.append(gauss_tp)\n",
    "    gauss_batch = np.array(gauss_batch)[:, :, :, :, np.newaxis]\n",
    "    return gauss_batch\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/nature_data'\n",
    "train_set_down = 0\n",
    "train_set_up = 1800\n",
    "nat_a4c_ims  = np.load(data_path + '/' + 'ims_a4c_1.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "nat_a4c_gts  = np.load(data_path + '/' + 'gts_a4c_1.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "nat_a4c_cls  = 4*np.ones((train_set_up))\n",
    "nat_a4c_pot  = np.load(data_path + '/' + 'pot_a4c_1.npy')[train_set_down:train_set_up, :]\n",
    "nat_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c_1.npy')[train_set_down:train_set_up, :]\n",
    "nat_a4c_frm  = np.load(data_path + '/' + 'frm_a4c_1.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load nature===============')\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/HMC_QU_data'\n",
    "train_set_down = 0\n",
    "train_set_up = 100\n",
    "hmc_a4c_ims  = np.load(data_path + '/' + 'ims_a4c_hmc.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "hmc_a4c_gts  = np.load(data_path + '/' + 'gts_a4c_hmc.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "hmc_a4c_cls  = 4*np.ones((train_set_up))\n",
    "hmc_a4c_pot  = np.load(data_path + '/' + 'pot_a4c_hmc.npy')[train_set_down:train_set_up, :]\n",
    "hmc_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c_hmc.npy')[train_set_down:train_set_up, :]\n",
    "hmc_a4c_frm  = np.load(data_path + '/' + 'frm_a4c_hmc.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load hmc===============')\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/camus_data'\n",
    "train_set_down = 0\n",
    "train_set_up = 500\n",
    "camus_a2c_ims  = np.load(data_path + '/' + 'ims_a2c_camus.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "camus_a4c_ims  = np.load(data_path + '/' + 'ims_a4c_camus.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "camus_a2c_gts  = np.load(data_path + '/' + 'gts_a2c_camus.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "camus_a4c_gts  = np.load(data_path + '/' + 'gts_a4c_camus.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "camus_a2c_cls  = 2*np.ones((train_set_up))\n",
    "camus_a4c_cls  = 4*np.ones((train_set_up))\n",
    "camus_a2c_pot  = np.load(data_path + '/' + 'pot_a2c_camus.npy')[train_set_down:train_set_up, :]\n",
    "camus_a4c_pot  = np.load(data_path + '/' + 'pot_a4c_camus.npy')[train_set_down:train_set_up, :]\n",
    "camus_a2c_potmap  = np.load(data_path + '/' + 'potmap_a2c_camus.npy')[train_set_down:train_set_up, :]\n",
    "camus_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c_camus.npy')[train_set_down:train_set_up, :]\n",
    "camus_a2c_frm  = np.load(data_path + '/' + 'frm_a2c_camus.npy')[train_set_down:train_set_up, :]\n",
    "camus_a4c_frm  = np.load(data_path + '/' + 'frm_a4c_camus.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load camus===============')\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/lm_data'\n",
    "train_set_down = 0\n",
    "train_set_up = 122\n",
    "lm_a2c_ims  = np.load(data_path + '/' + 'ims_a2c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a3c_ims  = np.load(data_path + '/' + 'ims_a3c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a4c_ims  = np.load(data_path + '/' + 'ims_a4c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a2c_gts  = np.load(data_path + '/' + 'gts_a2c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a3c_gts  = np.load(data_path + '/' + 'gts_a3c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a4c_gts  = np.load(data_path + '/' + 'gts_a4c_lm.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "lm_a2c_cls  = 2*np.ones((train_set_up))\n",
    "lm_a3c_cls  = 3*np.ones((train_set_up))\n",
    "lm_a4c_cls  = 4*np.ones((train_set_up))\n",
    "lm_a2c_pot  = np.load(data_path + '/' + 'pot_a2c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a3c_pot  = np.load(data_path + '/' + 'pot_a3c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a4c_pot  = np.load(data_path + '/' + 'pot_a4c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a2c_potmap  = np.load(data_path + '/' + 'potmap_a2c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a3c_potmap  = np.load(data_path + '/' + 'potmap_a3c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a2c_frm  = np.load(data_path + '/' + 'frm_a2c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a3c_frm  = np.load(data_path + '/' + 'frm_a3c_lm.npy')[train_set_down:train_set_up, :]\n",
    "lm_a4c_frm  = np.load(data_path + '/' + 'frm_a4c_lm.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load lm===============')\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/mx_data'\n",
    "train_set_down = 0\n",
    "train_set_up = 80\n",
    "mx_a2c_ims  = np.load(data_path + '/' + 'ims_a2c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a3c_ims  = np.load(data_path + '/' + 'ims_a3c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a4c_ims  = np.load(data_path + '/' + 'ims_a4c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a2c_gts  = np.load(data_path + '/' + 'gts_a2c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a3c_gts  = np.load(data_path + '/' + 'gts_a3c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a4c_gts  = np.load(data_path + '/' + 'gts_a4c_mx.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "mx_a2c_cls  = 2*np.ones((train_set_up))\n",
    "mx_a3c_cls  = 3*np.ones((train_set_up))\n",
    "mx_a4c_cls  = 4*np.ones((train_set_up))\n",
    "mx_a2c_pot  = np.load(data_path + '/' + 'pot_a2c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a3c_pot  = np.load(data_path + '/' + 'pot_a3c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a4c_pot  = np.load(data_path + '/' + 'pot_a4c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a2c_potmap  = np.load(data_path + '/' + 'potmap_a2c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a3c_potmap  = np.load(data_path + '/' + 'potmap_a3c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a2c_frm  = np.load(data_path + '/' + 'frm_a2c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a3c_frm  = np.load(data_path + '/' + 'frm_a3c_mx.npy')[train_set_down:train_set_up, :]\n",
    "mx_a4c_frm  = np.load(data_path + '/' + 'frm_a4c_mx.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load mx===============')\n",
    "\n",
    "data_path = '/data/zhangzhenxuan/szkid'\n",
    "train_set_down = 0\n",
    "train_set_up = 80\n",
    "sz_a2c_ims  = np.load(data_path + '/' + 'ims_a2c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a3c_ims  = np.load(data_path + '/' + 'ims_a3c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a4c_ims  = np.load(data_path + '/' + 'ims_a4c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_asc_ims  = np.load(data_path + '/' + 'ims_asc.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a2c_gts  = np.load(data_path + '/' + 'gts_a2c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a3c_gts  = np.load(data_path + '/' + 'gts_a3c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a4c_gts  = np.load(data_path + '/' + 'gts_a4c.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_asc_gts  = np.load(data_path + '/' + 'gts_asc.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_out_gts  = np.load(data_path + '/' + 'gts_out.npy')[train_set_down:train_set_up, :, :, :, np.newaxis]\n",
    "sz_a2c_cls  = 2*np.ones((train_set_up))\n",
    "sz_a3c_cls  = 3*np.ones((train_set_up))\n",
    "sz_a4c_cls  = 4*np.ones((train_set_up))\n",
    "sz_asc_cls  = 1*np.ones((train_set_up))\n",
    "sz_a2c_pot  = np.load(data_path + '/' + 'pot_a2c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a3c_pot  = np.load(data_path + '/' + 'pot_a3c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a4c_pot  = np.load(data_path + '/' + 'pot_a4c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a2c_potmap  = np.load(data_path + '/' + 'potmap_a2c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a3c_potmap  = np.load(data_path + '/' + 'potmap_a3c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a4c_potmap  = np.load(data_path + '/' + 'potmap_a4c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a2c_frm  = np.load(data_path + '/' + 'frm_a2c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a3c_frm  = np.load(data_path + '/' + 'frm_a3c.npy')[train_set_down:train_set_up, :]\n",
    "sz_a4c_frm  = np.load(data_path + '/' + 'frm_a4c.npy')[train_set_down:train_set_up, :]\n",
    "print('==============load sz===============')\n",
    "\n",
    "# ims         = np.concatenate((hmc_a4c_ims,camus_a2c_ims, camus_a4c_ims,mx_a2c_ims, mx_a3c_ims, mx_a4c_ims,\n",
    "#                               lm_a2c_ims, lm_a3c_ims, lm_a4c_ims,sz_a2c_ims, sz_a3c_ims, sz_a4c_ims, sz_asc_ims), axis=0)\n",
    "# gts         = np.concatenate((hmc_a4c_gts,camus_a2c_gts, camus_a4c_gts,mx_a2c_gts, mx_a3c_gts, mx_a4c_gts,\n",
    "#                               lm_a2c_gts, lm_a3c_gts, lm_a4c_gts,sz_a2c_gts, sz_a3c_gts, sz_a4c_gts, sz_out_gts), axis=0)\n",
    "# clss         = np.concatenate((hmc_a4c_cls,camus_a2c_cls, camus_a4c_cls,mx_a2c_cls, mx_a3c_cls, mx_a4c_cls,\n",
    "#                              lm_a2c_cls, lm_a3c_cls, lm_a4c_cls,sz_a2c_cls, sz_a3c_cls, sz_a4c_cls, sz_asc_cls), axis=0)\n",
    "# ims         = np.concatenate((hmc_a4c_ims,camus_a2c_ims, camus_a4c_ims,mx_a2c_ims, mx_a3c_ims, mx_a4c_ims,\n",
    "#                               lm_a2c_ims, lm_a3c_ims, lm_a4c_ims,sz_a2c_ims, sz_a3c_ims, sz_a4c_ims), axis=0)\n",
    "# gts         = np.concatenate((hmc_a4c_gts,camus_a2c_gts, camus_a4c_gts,mx_a2c_gts, mx_a3c_gts, mx_a4c_gts,\n",
    "#                               lm_a2c_gts, lm_a3c_gts, lm_a4c_gts,sz_a2c_gts, sz_a3c_gts, sz_a4c_gts), axis=0)\n",
    "# clss        = np.concatenate((hmc_a4c_cls,camus_a2c_cls, camus_a4c_cls,mx_a2c_cls, mx_a3c_cls, mx_a4c_cls,\n",
    "#                               lm_a2c_cls, lm_a3c_cls, lm_a4c_cls,sz_a2c_cls, sz_a3c_cls, sz_a4c_cls), axis=0)\n",
    "# reg         = np.concatenate((hmc_a4c_reg,camus_a2c_reg, camus_a4c_reg,mx_a2c_reg, mx_a3c_reg, mx_a4c_reg,\n",
    "#                               lm_a2c_reg, lm_a3c_reg, lm_a4c_reg,sz_a2c_reg, sz_a3c_reg, sz_a4c_reg), axis=0)\n",
    "\n",
    "data_mode = 'nat'\n",
    "if data_mode == 'psax':\n",
    "    ims         = np.concatenate((sz_a2c_ims, sz_a3c_ims, sz_a4c_ims, sz_asc_ims), axis=0)\n",
    "    gts         = np.concatenate((sz_a2c_gts, sz_a3c_gts, sz_a4c_gts, sz_out_gts), axis=0)\n",
    "    clss         = np.concatenate((sz_a2c_cls, sz_a3c_cls, sz_a4c_cls, sz_asc_cls), axis=0)\n",
    "    pot         = np.concatenate((sz_a2c_pot, sz_a3c_pot, sz_a4c_pot, sz_asc_pot), axis=0)\n",
    "    frm         = np.concatenate((sz_a2c_frm, sz_a3c_frm, sz_a4c_frm, sz_asc_frm), axis=0)\n",
    "elif data_mode == 'apic_1':\n",
    "    ims         = np.concatenate((sz_a2c_ims, sz_a3c_ims, sz_a4c_ims), axis=0)\n",
    "    gts         = np.concatenate((sz_a2c_gts, sz_a3c_gts, sz_a4c_gts), axis=0)\n",
    "    clss         = np.concatenate((sz_a2c_cls, sz_a3c_cls, sz_a4c_cls), axis=0)\n",
    "    pot         = np.concatenate((sz_a2c_pot, sz_a3c_pot, sz_a4c_pot), axis=0)\n",
    "    potmap      = locmap(pot)\n",
    "    frm         = np.concatenate((sz_a2c_frm, sz_a3c_frm, sz_a4c_frm), axis=0) \n",
    "elif data_mode == 'apic_2':\n",
    "    ims         = np.concatenate((hmc_a4c_ims,camus_a2c_ims, camus_a4c_ims,mx_a2c_ims, mx_a3c_ims, mx_a4c_ims,\n",
    "                                  lm_a2c_ims, lm_a3c_ims, lm_a4c_ims,sz_a2c_ims, sz_a3c_ims, sz_a4c_ims), axis=0)\n",
    "    print(ims.shape)\n",
    "    gts         = np.concatenate((hmc_a4c_gts,camus_a2c_gts, camus_a4c_gts,mx_a2c_gts, mx_a3c_gts, mx_a4c_gts,\n",
    "                                  lm_a2c_gts, lm_a3c_gts, lm_a4c_gts,sz_a2c_gts, sz_a3c_gts, sz_a4c_gts), axis=0)\n",
    "    print(gts.shape)\n",
    "    clss        = np.concatenate((hmc_a4c_cls,camus_a2c_cls, camus_a4c_cls,mx_a2c_cls, mx_a3c_cls, mx_a4c_cls,\n",
    "                                  lm_a2c_cls, lm_a3c_cls, lm_a4c_cls,sz_a2c_cls, sz_a3c_cls, sz_a4c_cls), axis=0)\n",
    "    print(clss.shape)\n",
    "    pot         = np.concatenate((hmc_a4c_pot,camus_a2c_pot, camus_a4c_pot,mx_a2c_pot, mx_a3c_pot, mx_a4c_pot,\n",
    "                                  lm_a2c_pot, lm_a3c_pot, lm_a4c_pot,sz_a2c_pot, sz_a3c_pot, sz_a4c_pot), axis=0)\n",
    "    potmap      = np.concatenate((hmc_a4c_potmap,camus_a2c_potmap, camus_a4c_potmap,mx_a2c_potmap, mx_a3c_potmap, mx_a4c_potmap,\n",
    "                                  lm_a2c_potmap, lm_a3c_potmap, lm_a4c_potmap,sz_a2c_potmap, sz_a3c_potmap, sz_a4c_potmap), axis=0)\n",
    "    print(pot.shape,potmap.shape)\n",
    "    #potmap      = locmap(pot)\n",
    "    frm         = np.concatenate((hmc_a4c_frm,camus_a2c_frm, camus_a4c_frm,mx_a2c_frm, mx_a3c_frm, mx_a4c_frm,\n",
    "                                  lm_a2c_frm, lm_a3c_frm, lm_a4c_frm,sz_a2c_frm, sz_a3c_frm, sz_a4c_frm), axis=0) \n",
    "elif data_mode == 'nat':\n",
    "    ims         = np.concatenate((nat_a4c_ims,hmc_a4c_ims,camus_a2c_ims, camus_a4c_ims,mx_a2c_ims, mx_a3c_ims, mx_a4c_ims,\n",
    "                                  lm_a2c_ims, lm_a3c_ims, lm_a4c_ims,sz_a2c_ims, sz_a3c_ims, sz_a4c_ims), axis=0)\n",
    "    print(ims.shape)\n",
    "    gts         = np.concatenate((nat_a4c_gts,hmc_a4c_gts,camus_a2c_gts, camus_a4c_gts,mx_a2c_gts, mx_a3c_gts, mx_a4c_gts,\n",
    "                                  lm_a2c_gts, lm_a3c_gts, lm_a4c_gts,sz_a2c_gts, sz_a3c_gts, sz_a4c_gts), axis=0)\n",
    "    print(gts.shape)\n",
    "    clss        = np.concatenate((nat_a4c_cls,hmc_a4c_cls,camus_a2c_cls, camus_a4c_cls,mx_a2c_cls, mx_a3c_cls, mx_a4c_cls,\n",
    "                                  lm_a2c_cls, lm_a3c_cls, lm_a4c_cls,sz_a2c_cls, sz_a3c_cls, sz_a4c_cls), axis=0)\n",
    "    print(clss.shape)\n",
    "    pot         = np.concatenate((nat_a4c_pot,hmc_a4c_pot,camus_a2c_pot, camus_a4c_pot,mx_a2c_pot, mx_a3c_pot, mx_a4c_pot,\n",
    "                                  lm_a2c_pot, lm_a3c_pot, lm_a4c_pot,sz_a2c_pot, sz_a3c_pot, sz_a4c_pot), axis=0)\n",
    "    potmap      = np.concatenate((nat_a4c_potmap,hmc_a4c_potmap,camus_a2c_potmap, camus_a4c_potmap,mx_a2c_potmap, mx_a3c_potmap, mx_a4c_potmap,\n",
    "                                  lm_a2c_potmap, lm_a3c_potmap, lm_a4c_potmap,sz_a2c_potmap, sz_a3c_potmap, sz_a4c_potmap), axis=0)\n",
    "    print(pot.shape,potmap.shape)\n",
    "    #potmap      = locmap(pot)\n",
    "    frm         = np.concatenate((nat_a4c_frm,hmc_a4c_frm,camus_a2c_frm, camus_a4c_frm,mx_a2c_frm, mx_a3c_frm, mx_a4c_frm,\n",
    "                                  lm_a2c_frm, lm_a3c_frm, lm_a4c_frm,sz_a2c_frm, sz_a3c_frm, sz_a4c_frm), axis=0) \n",
    "    print(frm.shape) \n",
    "    \n",
    "# ims         = camus_a4c_ims\n",
    "# gts         = camus_a4c_gts\n",
    "#oe = OrdinalEncoder()\n",
    "#clss = oe.fit_transform(clss.reshape(-1, 1)).ravel()\n",
    "#clss = np.eye(4)[np.array(clss, dtype=np.int32)]\n",
    "print(ims.shape, gts.shape, clss.shape,pot.shape,potmap.shape,frm.shape)\n",
    "#ims = ims[:,:,:,:,:]\n",
    "#gts = gts[1:91,:,:,:,:]\n",
    "image = torch.from_numpy(ims)\n",
    "image = image.permute(0, 1, 4, 2, 3)\n",
    "label = torch.from_numpy(gts)\n",
    "label = label.permute(0, 1, 4, 2, 3)\n",
    "classi = torch.from_numpy(clss)\n",
    "point = torch.from_numpy(pot)\n",
    "pointmap = torch.from_numpy(potmap)\n",
    "pointmap = pointmap.permute(0, 1, 4, 2, 3)\n",
    "frame = torch.from_numpy(frm)\n",
    "#image = image[0:15, :, :,:,:]\n",
    "#label = label[0:15, :, :,:,:]\n",
    "pi = np.random.permutation(image.shape[0])\n",
    "image = image[pi, :, :,:,:]\n",
    "label = label[pi, :, :,:,:]\n",
    "classi = classi[pi]-1\n",
    "point = point[pi]\n",
    "pointmap = pointmap[pi]\n",
    "frame = frame[pi]\n",
    "#split train-test\n",
    "sp = 3500\n",
    "sup = 3700\n",
    "image_test = image[sp:sup,:,:,:,:]\n",
    "label_test = label[sp:sup,:,:,:,:]\n",
    "classi_test = classi[sp:sup]\n",
    "point_test = point[sp:sup,:,:]\n",
    "pointmap_test = pointmap[sp:sup,:,:,:,:]\n",
    "frame_test = frame[sp:sup,:]\n",
    "image = image[:sp,:,:,:,:]\n",
    "label = label[:sp:,:,:,:,:]\n",
    "classi = classi[:sp]\n",
    "pointmap = pointmap[:sp:,:,:,:,:]\n",
    "frame = frame[:sp,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "#训练分割分支可视化\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision\n",
    "import torch\n",
    "from torchvision.transforms import transforms\n",
    "from torch import nn, optim\n",
    "import timeit\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from tqdm import trange\n",
    "\n",
    "img_size = 128\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "torch.backends.cudnn.benchmark = True\n",
    "l_r = 0.0002  #0.0002\n",
    "device = torch.device(\"cuda\")\n",
    "train_mode = 'mtl'\n",
    "model = Multivit_net(img_dim=128,in_channels=1,out_channels=32,head_num=4,mlp_dim=512,block_num=8,\n",
    "                     patch_dim=16,class_num=1,drop_rate = 0.2,seq_frame=30,mode =train_mode,height=128,weight=128).to(device)\n",
    "#model.load_state_dict(torch.load('./weight/mlt_weights_beta0.1.pth'),True)\n",
    "model.load_state_dict(torch.load('./weight/mlt_weights_base.pth'),True)\n",
    "enet = Embbeding_net(dim=1, frame=30,device = device).to(device)\n",
    "enet.load_state_dict(torch.load( './weight_emb/mlt_emb_min4.pth'),True)\n",
    "param_optim = []\n",
    "layers = []\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=l_r)\n",
    "dice_less = 0.85\n",
    "cor_less = 0.85\n",
    "mae_less = 30\n",
    "pot_less = 30\n",
    "frm_less = 30\n",
    "cls_less = 0.8\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "criterion_f = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array([5/13,3/13,5/13])).float(),size_average=True)\n",
    "criterion_f.to(device)\n",
    "crit_pot = torch.nn.SmoothL1Loss()\n",
    "kl_loss = nn.KLDivLoss(reduction=\"batchmean\")\n",
    "pot_loss = nn.MSELoss(size_average=True)\n",
    "lossseg,losspot,lossfrm,losscls = [],[],[],[]\n",
    "for t in range(50):\n",
    "    # Forward pass: Compute predicted y by passing x to the model\n",
    "    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [5, 10,15], 0.1)\n",
    "    lens = image.shape[0]\n",
    "    print('epoch=',t)\n",
    "    dice,jcd, fb = 0,0,0\n",
    "    cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "    mae_all = 0\n",
    "    mae_all_pot = 0\n",
    "    mae_all_frm = 0\n",
    "    right,error = 0,0\n",
    "    image = image\n",
    "    label = label\n",
    "    classi = classi\n",
    "    point = point\n",
    "    pointmap = pointmap\n",
    "    frame = frame\n",
    "    lens2 = image_test.shape[0]\n",
    "    image_test = image_test.to(device)\n",
    "    label_test = label_test.to(device)\n",
    "    classi_test = classi_test.to(device)\n",
    "    point_test = point_test.to(device)\n",
    "    frame_test = frame_test.to(device)\n",
    "    model.train()\n",
    "    print('training')\n",
    "    with trange(lens) as tr:\n",
    "        for i in tr:\n",
    "            video = image[i:i+1,:,:,:,:].to(device)\n",
    "            #print('batch=',i,' of ',lens)\n",
    "            labelv = label[i:i+1,:,:,:,:].to(device)\n",
    "            classiv = classi[i:i+1].to(device)\n",
    "            pointi = point[i:i+1].to(device)\n",
    "            pointmapi = pointmap[i:i+1].to(device)\n",
    "            framei = frame[i].to(device)\n",
    "            if train_mode=='cls':\n",
    "                pred_cls = model(video)\n",
    "                # Compute and print loss\n",
    "                loss_cls = criterion(pred_cls,classiv.long())\n",
    "                loss = loss_cls\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]))\n",
    "                if float(classiv) == float(torch.max(pred_cls, 1)[1]):\n",
    "                    right = right+1\n",
    "                else:\n",
    "                    error = error+1\n",
    "            elif train_mode=='pot':\n",
    "                pred_pot,pred_potmap = model(video)\n",
    "                #print(pred_pot.shape,pointi.shape)\n",
    "                loss_pot = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                # loss_pot_map = kl_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                loss_pot_map = pot_loss(pointmapi.to(torch.float32),pred_potmap.to(torch.float32))\n",
    "                loss = loss_pot_map+loss_pot\n",
    "                mae = mae_point(pred_pot,pointi)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae=float(mae),loss_map=float(loss_pot_map),loss_pot=float(loss_pot))\n",
    "                mae_all+=float(mae)\n",
    "            elif train_mode=='frm':\n",
    "                pred_frm = model(video)\n",
    "                # pred_frm = torch.tanh(pred_frm)\n",
    "                # loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                framei = framei+1\n",
    "                loss_frm = criterion_f(pred_frm,framei.long())\n",
    "                loss = loss_frm\n",
    "                mae = acc(pred_frm,framei)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                mae_all+=float(mae)\n",
    "            elif train_mode=='seg':\n",
    "                pred_seg = model(video)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                dice = dice + float(dice_coeff)\n",
    "                loss1 = seg_loss(pred_seg, labelv)\n",
    "                loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                a = t/50\n",
    "                loss = (1-a)*loss1+a*loss2\n",
    "                hd,md = get_hausdorff(labelv,pred_seg)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),dice=float(dice_coeff))\n",
    "            elif train_mode == 'mtl':\n",
    "                pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                # loss_cls = criterion(pred_cls,classiv.long())\n",
    "                loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                loss_pot_map = pot_loss(torch.sigmoid(pred_potmap).to(torch.float32), pointmapi.to(torch.float32))\n",
    "                #print(pred_potmap.shape, pointmapi.shape)\n",
    "                #loss_pot_map = criterion(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                loss_pot = loss_pot1+loss_pot_map\n",
    "                # pred_frm = torch.tanh(pred_frm)\n",
    "                framei = framei+1\n",
    "                # print(framei.shape)\n",
    "                loss_frm = criterion_f(pred_frm,framei.long())\n",
    "                loss1 = seg_loss(pred_seg, labelv)\n",
    "                loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                loss_con1 = l2_norm(pred_seg,pointmapi)\n",
    "                loss_con2 = kl_loss(F.log_softmax(cal_a(pred_seg)[0,:,0], dim=0), F.softmax(torch.topk(pred_frm, 1)[1].squeeze(1).float(), dim=0))\n",
    "                a = t/50\n",
    "                #embspace\n",
    "                pot_emb,frm_emb,seg_emb = enet(pred_pot,framei,pred_seg)\n",
    "                loss_emb = cosine_loss(pot_emb,framei.unsqueeze(1))+cosine_loss(seg_emb,framei.unsqueeze(1))+crit_pot(pot_emb,seg_emb)\n",
    "                loss_seg = (1-a)*loss1+a*loss2\n",
    "                w_1,w_2,w_3 = 1.0,1.0,1.0\n",
    "                loss = w_1*loss_pot+w_2*loss_frm+w_3*loss_seg+loss_emb\n",
    "                # loss = loss_con1+loss_con2\n",
    "                mae_lvd,mae_mvd = mae_div(pred_pot,pointi)\n",
    "                mae_frm,mae_es,mae_ed = acc(pred_frm,framei)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                mae_all_frm+=float(mae_frm)\n",
    "                mae_all_pot+=float(mae_lvd)\n",
    "                dice = dice + float(dice_coeff)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae_lvd=float(mae_lvd),mae_mvd=float(mae_mvd)\n",
    "                              ,mae_frm=float(mae_frm),mae_ed = float(mae_ed),mae_es = float(mae_es),\n",
    "                               dice=float(dice_coeff),loss_seg=float(loss_seg),loss_frm=float(loss_frm),loss_pot=float(loss_pot))\n",
    "                # if float(classiv) == float(torch.max(pred_cls, 1)[1]):\n",
    "                #     right = right+1\n",
    "                # else:\n",
    "                #     error = error+1\n",
    "                lossseg.append(float(loss_seg))\n",
    "                lossfrm.append(float(loss_frm))\n",
    "                losspot.append(float(loss_pot))\n",
    "                # losscls.append(float(loss_cls))\n",
    "            # Zero gradients, perform a backward pass, and update the weights.\n",
    "            optimizer.zero_grad() # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度\n",
    "            loss.backward() # loss反向传播\n",
    "            optimizer.step() # 反向传播后参数更新 \n",
    "    #print('loss=', loss)\n",
    "    if train_mode=='cls':\n",
    "        print('epoch_train_acc=',right/(right+error))\n",
    "    elif train_mode=='pot':\n",
    "        print('epoch_train_acc=',mae_all/lens)\n",
    "    elif train_mode=='frm':\n",
    "        print('epoch_train_acc=',mae_all/lens)\n",
    "    elif train_mode=='seg':\n",
    "        print('epoch_train_acc=',dice/lens)\n",
    "    elif train_mode=='mtl':\n",
    "        print('epoch_train_frame=',mae_all_frm/lens,\n",
    "              'epoch_train_point=',mae_all_pot/lens,'epoch_train_dice=',dice/lens)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    #测试阶段\n",
    "        #model.eval()\n",
    "        print('evaluating')\n",
    "        dice,jcd, fb = 0,0,0\n",
    "        cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "        mae_all = 0\n",
    "        mae_all_pot = 0\n",
    "        mae_all_frm = 0\n",
    "        right,error = 0,0\n",
    "        with trange(lens2) as tr:\n",
    "            for i in tr:\n",
    "                video_test = image_test[i:i+1,:,:,:,:]\n",
    "                #print('batch=',i,' of ',lens)\n",
    "                labelv_test = label_test[i:i+1,:,:,:,:]\n",
    "                classiv_test = classi_test[i:i+1]\n",
    "                pointi_test = point_test[i:i+1]\n",
    "                framei_test = frame_test[i]\n",
    "                if train_mode == 'cls':\n",
    "                    pred_test_cls = model(video_test)\n",
    "                    tr.set_postfix(loss=float(loss),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]))\n",
    "                    if float(classiv_test) == float(torch.max(pred_test_cls, 1)[1]):\n",
    "                        right = right+1\n",
    "                    else:\n",
    "                        error = error+1\n",
    "                elif train_mode == 'pot':\n",
    "                    pred_pot_test,_ = model(video_test)\n",
    "                    mae = mae_point(pred_pot_test,pointi_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                    mae_all+=float(mae) \n",
    "                elif train_mode == 'frm':\n",
    "                    pred_frm_test = model(video_test)\n",
    "                    framei_test = framei_test+1\n",
    "                    mae = acc(pred_frm_test,framei_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                    mae_all+=float(mae) \n",
    "                elif train_mode=='seg':\n",
    "                    pred_seg_test = model(video_test)\n",
    "                    dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv_test,pred_seg_test) # 计算指标\n",
    "                    dice = dice + float(dice_coeff)\n",
    "                    hd,md = get_hausdorff(labelv,pred_seg_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),dice=float(dice_coeff))\n",
    "                elif train_mode=='mtl':\n",
    "                    pred_pot_test,_,pred_frm_test,pred_seg_test = model(video_test)\n",
    "                    pred_frm_test = torch.tanh(pred_frm_test)\n",
    "                    mae_lvd_test,mae_mvd_test = mae_div(pred_pot_test,pointi_test)\n",
    "                    framei_test = framei_test+1\n",
    "                    mae_frm_test,mae_es_test,mae_ed_test = acc(pred_frm_test,framei_test)\n",
    "                    # mae_frm_test = mae_point(pred_frm_test,framei_test)\n",
    "                    dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv_test,pred_seg_test) # 计算指标\n",
    "                    mae_all_frm+=float(mae_frm_test)\n",
    "                    mae_all_pot+=float(mae_lvd_test)\n",
    "                    dice = dice + float(dice_coeff)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),\n",
    "                                   mae_lvd=float(mae_lvd_test),mae_mvd=float(mae_mvd_test),mae_frm=float(mae_frm_test),\n",
    "                                   mae_ed=float(mae_ed_test),mae_es=float(mae_es_test),dice=float(dice_coeff))\n",
    "                    # if float(classiv_test) == float(torch.max(pred_cls_test, 1)[1]):\n",
    "                    #     right = right+1\n",
    "                    # else:\n",
    "                    #     error = error+1\n",
    "                    \n",
    "            if train_mode=='cls':\n",
    "                print('epoch_train_acc=',right/(right+error))\n",
    "            elif train_mode=='pot':\n",
    "                print('epoch_train_acc=',mae_all/lens2)\n",
    "            elif train_mode=='frm':\n",
    "                print('epoch_train_acc=',mae_all/lens2)\n",
    "            elif train_mode=='seg':\n",
    "                print('epoch_train_acc=',dice/lens2)\n",
    "            elif train_mode=='mtl':\n",
    "                print('epoch_train_point=',mae_all_frm/lens2,\n",
    "                      'epoch_train_frame=',mae_all_pot/lens2,'epoch_train_dice=',dice/lens2)\n",
    "                \n",
    "    if train_mode == 'cls':\n",
    "        cls_acc =   right/(right+error)  \n",
    "        if cls_acc >= cls_less:\n",
    "            print('save model')\n",
    "            cls_less = cls_acc\n",
    "            torch.save(model.state_dict(), './weight_cls/cls_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',cls_less)\n",
    "    elif train_mode == 'pot':\n",
    "        pot_acc =   mae_all/lens \n",
    "        if pot_acc <= pot_less:\n",
    "            print('save model')\n",
    "            pot_less = pot_acc\n",
    "            torch.save(model.state_dict(), './weight_pot/pot_weights4.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',pot_less)\n",
    "    elif train_mode == 'frm':\n",
    "        frm_acc =   mae_all/lens2 \n",
    "        if frm_acc >= 0.85:\n",
    "            print('save model')\n",
    "            frm_less = frm_acc\n",
    "            torch.save(model.state_dict(), './weight_frm/frm_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',frm_less)\n",
    "    elif train_mode == 'seg':\n",
    "        seg_acc =   dice/lens2\n",
    "        if seg_acc >= dice_less:\n",
    "            print('save model')\n",
    "            dice_less = seg_acc\n",
    "            torch.save(model.state_dict(), './weight_seg/seg_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',dice_less)\n",
    "    elif train_mode == 'mtl':\n",
    "        seg_acc =   dice/lens2\n",
    "        if seg_acc >= dice_less:\n",
    "            print('save model')\n",
    "            dice_less = seg_acc\n",
    "            torch.save(model.state_dict(), './weight/mlt_weights_base2.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',dice_less)\n",
    "scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.topk(pred_frm, 1)[1].squeeze(1)\n",
    "cal_a(pred_seg).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    i = 23\n",
    "    video_test = image_test[i:i+1,:,:,:,:]\n",
    "    framei_test = frame_test[i:i+1]\n",
    "    b = model(video_test)\n",
    "    print(torch.max(b, 1)[1])\n",
    "    c = framei_test+1\n",
    "    print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import os\n",
    "\n",
    "class MinNormSolver:\n",
    "    MAX_ITER = 250\n",
    "    STOP_CRIT = 1e-5\n",
    "\n",
    "    def _min_norm_element_from2(v1v1, v1v2, v2v2):\n",
    "        \"\"\"\n",
    "        Analytical solution for min_{c} |cx_1 + (1-c)x_2|_2^2\n",
    "        d is the distance (objective) optimzed\n",
    "        v1v1 = <x1,x1>\n",
    "        v1v2 = <x1,x2>\n",
    "        v2v2 = <x2,x2>\n",
    "        \"\"\"\n",
    "        if v1v2 >= v1v1:\n",
    "            # Case: Fig 1, third column\n",
    "            gamma = 0.999\n",
    "            cost = v1v1\n",
    "            return gamma, cost\n",
    "        if v1v2 >= v2v2:\n",
    "            # Case: Fig 1, first column\n",
    "            gamma = 0.001\n",
    "            cost = v2v2\n",
    "            return gamma, cost\n",
    "        # Case: Fig 1, second column\n",
    "        gamma = -1.0 * ( (v1v2 - v2v2) / (v1v1+v2v2 - 2*v1v2) )\n",
    "        cost = v2v2 + gamma*(v1v2 - v2v2)\n",
    "        return gamma, cost\n",
    "\n",
    "    def _min_norm_2d(vecs, dps):\n",
    "        \"\"\"\n",
    "        Find the minimum norm solution as combination of two points\n",
    "        This is correct only in 2D\n",
    "        ie. min_c |\\sum c_i x_i|_2^2 st. \\sum c_i = 1 , 1 >= c_1 >= 0 for all i, c_i + c_j = 1.0 for some i, j\n",
    "        \"\"\"\n",
    "        dmin = 1e8\n",
    "        for i in range(len(vecs)):\n",
    "            for j in range(i+1,len(vecs)):\n",
    "                if (i,j) not in dps:\n",
    "                    dps[(i, j)] = 0.0\n",
    "                    for k in range(len(vecs[i])):\n",
    "                        dps[(i,j)] += torch.dot(vecs[i][k], vecs[j][k]).item()#torch.dot(vecs[i][k], vecs[j][k]).data[0]\n",
    "                    dps[(j, i)] = dps[(i, j)]\n",
    "                if (i,i) not in dps:\n",
    "                    dps[(i, i)] = 0.0\n",
    "                    for k in range(len(vecs[i])):\n",
    "                        dps[(i,i)] += torch.dot(vecs[i][k], vecs[i][k]).item()#torch.dot(vecs[i][k], vecs[i][k]).data[0]\n",
    "                if (j,j) not in dps:\n",
    "                    dps[(j, j)] = 0.0   \n",
    "                    for k in range(len(vecs[i])):\n",
    "                        dps[(j, j)] += torch.dot(vecs[j][k], vecs[j][k]).item()#torch.dot(vecs[j][k], vecs[j][k]).data[0]\n",
    "                c,d = MinNormSolver._min_norm_element_from2(dps[(i,i)], dps[(i,j)], dps[(j,j)])\n",
    "                if d < dmin:\n",
    "                    dmin = d\n",
    "                    sol = [(i,j),c,d]\n",
    "        return sol, dps\n",
    "\n",
    "    def _projection2simplex(y):\n",
    "        \"\"\"\n",
    "        Given y, it solves argmin_z |y-z|_2 st \\sum z = 1 , 1 >= z_i >= 0 for all i\n",
    "        \"\"\"\n",
    "        m = len(y)\n",
    "        sorted_y = np.flip(np.sort(y), axis=0)\n",
    "        tmpsum = 0.0\n",
    "        tmax_f = (np.sum(y) - 1.0)/m\n",
    "        for i in range(m-1):\n",
    "            tmpsum+= sorted_y[i]\n",
    "            tmax = (tmpsum - 1)/ (i+1.0)\n",
    "            if tmax > sorted_y[i+1]:\n",
    "                tmax_f = tmax\n",
    "                break\n",
    "        return np.maximum(y - tmax_f, np.zeros(y.shape))\n",
    "    \n",
    "    def _next_point(cur_val, grad, n):\n",
    "        proj_grad = grad - ( np.sum(grad) / n )\n",
    "        tm1 = -1.0*cur_val[proj_grad<0]/proj_grad[proj_grad<0]\n",
    "        tm2 = (1.0 - cur_val[proj_grad>0])/(proj_grad[proj_grad>0])\n",
    "        \n",
    "        skippers = np.sum(tm1<1e-7) + np.sum(tm2<1e-7)\n",
    "        t = 1\n",
    "        if len(tm1[tm1>1e-7]) > 0:\n",
    "            t = np.min(tm1[tm1>1e-7])\n",
    "        if len(tm2[tm2>1e-7]) > 0:\n",
    "            t = min(t, np.min(tm2[tm2>1e-7]))\n",
    "\n",
    "        next_point = proj_grad*t + cur_val\n",
    "        next_point = MinNormSolver._projection2simplex(next_point)\n",
    "        return next_point\n",
    "\n",
    "    def find_min_norm_element(vecs):\n",
    "        \"\"\"\n",
    "        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull\n",
    "        as min |u|_2 st. u = \\sum c_i vecs[i] and \\sum c_i = 1.\n",
    "        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})\n",
    "        Hence, we find the best 2-task solution, and then run the projected gradient descent until convergence\n",
    "        \"\"\"\n",
    "        # Solution lying at the combination of two points\n",
    "        dps = {}\n",
    "        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)\n",
    "        \n",
    "        n=len(vecs)\n",
    "        sol_vec = np.zeros(n)\n",
    "        sol_vec[init_sol[0][0]] = init_sol[1]\n",
    "        sol_vec[init_sol[0][1]] = 1 - init_sol[1]\n",
    "\n",
    "        if n < 3:\n",
    "            # This is optimal for n=2, so return the solution\n",
    "            return sol_vec , init_sol[2]\n",
    "    \n",
    "        iter_count = 0\n",
    "\n",
    "        grad_mat = np.zeros((n,n))\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                grad_mat[i,j] = dps[(i, j)]\n",
    "                \n",
    "\n",
    "        while iter_count < MinNormSolver.MAX_ITER:\n",
    "            grad_dir = -1.0*np.dot(grad_mat, sol_vec)\n",
    "            new_point = MinNormSolver._next_point(sol_vec, grad_dir, n)\n",
    "            # Re-compute the inner products for line search\n",
    "            v1v1 = 0.0\n",
    "            v1v2 = 0.0\n",
    "            v2v2 = 0.0\n",
    "            for i in range(n):\n",
    "                for j in range(n):\n",
    "                    v1v1 += sol_vec[i]*sol_vec[j]*dps[(i,j)]\n",
    "                    v1v2 += sol_vec[i]*new_point[j]*dps[(i,j)]\n",
    "                    v2v2 += new_point[i]*new_point[j]*dps[(i,j)]\n",
    "            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)\n",
    "            new_sol_vec = nc*sol_vec + (1-nc)*new_point\n",
    "            change = new_sol_vec - sol_vec\n",
    "            if np.sum(np.abs(change)) < MinNormSolver.STOP_CRIT:\n",
    "                return sol_vec, nd\n",
    "            sol_vec = new_sol_vec\n",
    "\n",
    "    def find_min_norm_element_FW(vecs):\n",
    "        \"\"\"\n",
    "        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull\n",
    "        as min |u|_2 st. u = \\sum c_i vecs[i] and \\sum c_i = 1.\n",
    "        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})\n",
    "        Hence, we find the best 2-task solution, and then run the Frank Wolfe until convergence\n",
    "        \"\"\"\n",
    "        # Solution lying at the combination of two points\n",
    "        dps = {}\n",
    "        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)\n",
    "\n",
    "        n=len(vecs)\n",
    "        sol_vec = np.zeros(n)\n",
    "        sol_vec[init_sol[0][0]] = init_sol[1]\n",
    "        sol_vec[init_sol[0][1]] = 1 - init_sol[1]\n",
    "\n",
    "        if n < 3:\n",
    "            # This is optimal for n=2, so return the solution\n",
    "            return sol_vec , init_sol[2]\n",
    "\n",
    "        iter_count = 0\n",
    "\n",
    "        grad_mat = np.zeros((n,n))\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                grad_mat[i,j] = dps[(i, j)]\n",
    "\n",
    "        while iter_count < MinNormSolver.MAX_ITER:\n",
    "            t_iter = np.argmin(np.dot(grad_mat, sol_vec))\n",
    "\n",
    "            v1v1 = np.dot(sol_vec, np.dot(grad_mat, sol_vec))\n",
    "            v1v2 = np.dot(sol_vec, grad_mat[:, t_iter])\n",
    "            v2v2 = grad_mat[t_iter, t_iter]\n",
    "\n",
    "            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)\n",
    "            new_sol_vec = nc*sol_vec\n",
    "            new_sol_vec[t_iter] += 1 - nc\n",
    "\n",
    "            change = new_sol_vec - sol_vec\n",
    "            if np.sum(np.abs(change)) < MinNormSolver.STOP_CRIT:\n",
    "                return sol_vec, nd\n",
    "            sol_vec = new_sol_vec\n",
    "\n",
    "\n",
    "def gradient_normalizers(grads, losses, normalization_type):\n",
    "    gn = {}\n",
    "    if normalization_type == 'l2':\n",
    "        for t in grads:\n",
    "            gn[t] = np.sqrt(np.sum([gr.pow(2).sum().data[0] for gr in grads[t]]))\n",
    "    elif normalization_type == 'loss':\n",
    "        for t in grads:\n",
    "            gn[t] = losses[t]\n",
    "    elif normalization_type == 'loss+':\n",
    "        for t in grads:\n",
    "            gn[t] = losses[t] * np.sqrt(np.sum([gr.pow(2).sum().data[0] for gr in grads[t]]))\n",
    "    elif normalization_type == 'none':\n",
    "        for t in grads:\n",
    "            gn[t] = 1.0\n",
    "    else:\n",
    "        print('ERROR: Invalid Normalization Type')\n",
    "    return gn\n",
    "\n",
    "def get_d_paretomtl_init(grads,value,weights,i):\n",
    "    \"\"\" \n",
    "    calculate the gradient direction for ParetoMTL initialization \n",
    "    \"\"\"\n",
    "    \n",
    "    flag = False\n",
    "    nobj = value.shape\n",
    "   \n",
    "    # check active constraints\n",
    "    current_weight = weights[i]\n",
    "    rest_weights = weights\n",
    "    w = rest_weights - current_weight\n",
    "    \n",
    "    gx =  torch.matmul(w,value/torch.norm(value))\n",
    "    idx = gx >  0\n",
    "   \n",
    "    # calculate the descent direction\n",
    "    if torch.sum(idx) <= 0:\n",
    "        flag = True\n",
    "        return flag, torch.zeros(nobj)\n",
    "    if torch.sum(idx) == 1:\n",
    "        sol = torch.ones(1).cuda().float()\n",
    "    else:\n",
    "        vec =  torch.matmul(w[idx],grads)\n",
    "        sol, nd = MinNormSolver.find_min_norm_element([[vec[t]] for t in range(len(vec))])\n",
    "\n",
    "\n",
    "    weight0 =  torch.sum(torch.stack([sol[j] * w[idx][j ,0] for j in torch.arange(0, torch.sum(idx))]))\n",
    "    weight1 =  torch.sum(torch.stack([sol[j] * w[idx][j ,1] for j in torch.arange(0, torch.sum(idx))]))\n",
    "    weight2 =  torch.sum(torch.stack([sol[j] * w[idx][j ,2] for j in torch.arange(0, torch.sum(idx))]))\n",
    "    weight3 =  torch.sum(torch.stack([sol[j] * w[idx][j ,3] for j in torch.arange(0, torch.sum(idx))]))\n",
    "    weight = torch.stack([weight0,weight1,weight2,weight3])\n",
    "   \n",
    "    \n",
    "    return flag, weight\n",
    "\n",
    "\n",
    "def get_d_paretomtl(grads,value,weights,i):\n",
    "    \"\"\" calculate the gradient direction for ParetoMTL \"\"\"\n",
    "    \n",
    "    # check active constraints\n",
    "    current_weight = weights[i]\n",
    "    rest_weights = weights \n",
    "    w = rest_weights - current_weight\n",
    "    \n",
    "    gx =  torch.matmul(w,value/torch.norm(value))\n",
    "    idx = gx >  0\n",
    "    \n",
    "\n",
    "    # calculate the descent direction\n",
    "    if torch.sum(idx) <= 0:\n",
    "        sol, nd = MinNormSolver.find_min_norm_element([[grads[t]] for t in range(len(grads))])\n",
    "        return torch.tensor(sol).cuda().float()\n",
    "\n",
    "\n",
    "    vec =  torch.cat((grads, torch.matmul(w[idx],grads)))\n",
    "    sol, nd = MinNormSolver.find_min_norm_element([[vec[t]] for t in range(len(vec))])\n",
    "\n",
    "\n",
    "    weight0 =  sol[0] + torch.sum(torch.stack([sol[j] * w[idx][j - 2 ,0] for j in torch.arange(2, 2 + torch.sum(idx))]))\n",
    "    weight1 =  sol[1] + torch.sum(torch.stack([sol[j] * w[idx][j - 2 ,1] for j in torch.arange(2, 2 + torch.sum(idx))]))\n",
    "    weight2 =  sol[2] + torch.sum(torch.stack([sol[j] * w[idx][j - 2 ,2] for j in torch.arange(2, 2 + torch.sum(idx))]))\n",
    "    weight3 =  sol[3] + torch.sum(torch.stack([sol[j] * w[idx][j - 2 ,3] for j in torch.arange(2, 2 + torch.sum(idx))]))\n",
    "    weight = torch.stack([weight0,weight1,weight2,weight3])\n",
    "    \n",
    "    return weight\n",
    "\n",
    "\n",
    "def circle_points(r, n):\n",
    "    \"\"\"\n",
    "    generate evenly distributed unit preference vectors for two tasks\n",
    "    \"\"\"\n",
    "    circles = []\n",
    "    for r, n in zip(r, n):\n",
    "        t = np.linspace(0, 0.5 * np.pi, n)\n",
    "        x = r * np.cos(t)\n",
    "        y = 2/3*r * np.cos(t)+1/3*r * np.sin(t)\n",
    "        z = 1/3*r * np.cos(t)+2/3*r * np.sin(t)\n",
    "        k = r * np.sin(t)\n",
    "        circles.append(np.c_[x, y, z, k])\n",
    "    return circles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#训练分割分支可视化\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision\n",
    "import torch\n",
    "from torchvision.transforms import transforms\n",
    "from torch import nn, optim\n",
    "import timeit\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from tqdm import trange\n",
    "\n",
    "img_size = 128\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "torch.backends.cudnn.benchmark = True\n",
    "l_r = 0.0002  #0.0002\n",
    "device = torch.device(\"cuda\")\n",
    "train_mode = 'mtl'\n",
    "model = Multivit_net(img_dim=128,in_channels=1,out_channels=128,head_num=4,mlp_dim=512,block_num=8,\n",
    "                     patch_dim=16,class_num=1,drop_rate = 0.2,seq_frame=30,mode =train_mode,height=128,weight=128).to(device)\n",
    "#model.load_state_dict(torch.load('./weight/mlt_weights_beta0.1.pth'),True)\n",
    "#model.load_state_dict(torch.load('./weight_seg/seg_weights.pth'),True)\n",
    "param_optim = []\n",
    "layers = []\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=l_r)\n",
    "dice_less = 0.85\n",
    "cor_less = 0.85\n",
    "mae_less = 30\n",
    "pot_less = 30\n",
    "frm_less = 30\n",
    "cls_less = 0.8\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "crit_pot = torch.nn.SmoothL1Loss()\n",
    "kl_loss = nn.KLDivLoss(reduction=\"mean\")\n",
    "lossseg,losspot,lossfrm,losscls = [],[],[],[]\n",
    "for t in range(10):\n",
    "    # Forward pass: Compute predicted y by passing x to the model\n",
    "    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [25, 30], 0.1)\n",
    "    lens = image.shape[0]\n",
    "    print('epoch=',t)\n",
    "    dice,jcd, fb = 0,0,0\n",
    "    cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "    mae_all = 0\n",
    "    mae_all_pot = 0\n",
    "    mae_all_frm = 0\n",
    "    right,error = 0,0\n",
    "    image = image.to(device)\n",
    "    label = label.to(device)\n",
    "    classi = classi.to(device)\n",
    "    point = point.to(device)\n",
    "    pointmap = pointmap.to(device)\n",
    "    frame = frame.to(device)\n",
    "    lens2 = image_test.shape[0]\n",
    "    image_test = image_test.to(device)\n",
    "    label_test = label_test.to(device)\n",
    "    classi_test = classi_test.to(device)\n",
    "    point_test = point_test.to(device)\n",
    "    frame_test = frame_test.to(device)\n",
    "    model.train()\n",
    "    print('training')\n",
    "    with trange(lens) as tr:\n",
    "        for i in tr:\n",
    "            video = image[i:i+1,:,:,:,:]\n",
    "            #print('batch=',i,' of ',lens)\n",
    "            labelv = label[i:i+1,:,:,:,:]\n",
    "            classiv = classi[i:i+1]\n",
    "            pointi = point[i]\n",
    "            pointmapi = pointmap[i:i+1]\n",
    "            framei = frame[i]\n",
    "            if train_mode=='cls':\n",
    "                pred_cls = model(video)\n",
    "                # Compute and print loss\n",
    "                loss_cls = criterion(pred_cls,classiv.long())\n",
    "                loss = loss_cls\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]))\n",
    "                if float(classiv) == float(torch.max(pred_cls, 1)[1]):\n",
    "                    right = right+1\n",
    "                else:\n",
    "                    error = error+1\n",
    "            elif train_mode=='pot':\n",
    "                pred_pot,pred_potmap = model(video)\n",
    "                #print(pred_pot.shape,pointi.shape)\n",
    "                loss_pot = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                loss_pot_map = kl_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                loss = loss_pot+loss_pot_map\n",
    "                mae = mae_point(pred_pot,pointi)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                mae_all+=float(mae)\n",
    "            elif train_mode=='frm':\n",
    "                pred_frm = model(video)\n",
    "                pred_frm = torch.tanh(pred_frm)\n",
    "                loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                loss = loss_frm\n",
    "                mae = mae_point(pred_frm,framei)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                mae_all+=float(mae)\n",
    "            elif train_mode=='seg':\n",
    "                pred_seg = model(video)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                dice = dice + float(dice_coeff)\n",
    "                loss1 = seg_loss(pred_seg, labelv)\n",
    "                loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                a = t/50\n",
    "                loss = (1-a)*loss1+a*loss2\n",
    "                hd,md = get_hausdorff(labelv,pred_seg)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),dice=float(dice_coeff))\n",
    "            elif train_mode == 'mtl':\n",
    "                pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                # loss_cls = criterion(pred_cls,classiv.long())\n",
    "                loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                loss_pot_map = kl_loss(torch.sigmoid(pred_potmap).log().to(torch.float32), pointmapi.to(torch.float32))\n",
    "                #print(pred_potmap.shape, pointmapi.shape)\n",
    "                #loss_pot_map = criterion(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                loss_pot = loss_pot1+loss_pot_map\n",
    "                pred_frm = torch.tanh(pred_frm)\n",
    "                loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                loss1 = seg_loss(pred_seg, labelv)\n",
    "                loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                a = t/50\n",
    "                loss_seg = (1-a)*loss1+a*loss2\n",
    "                loss = loss_pot+loss_frm+loss_seg\n",
    "                mae_pot = mae_point(pred_pot,pointi)\n",
    "                mae_frm = mae_point(pred_frm,framei)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                mae_all_frm+=float(mae_frm)\n",
    "                mae_all_pot+=float(mae_pot)\n",
    "                dice = dice + float(dice_coeff)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss),mae_pot=float(mae_pot)\n",
    "                              ,mae_frm=float(mae_frm),dice=float(dice_coeff),loss_seg=float(loss_seg),loss_frm=float(loss_frm),\n",
    "                              loss_pot=float(loss_pot),loss_cls=float(loss_cls))\n",
    "                # if float(classiv) == float(torch.max(pred_cls, 1)[1]):\n",
    "                #     right = right+1\n",
    "                # else:\n",
    "                #     error = error+1\n",
    "                lossseg.append(float(loss_seg))\n",
    "                lossfrm.append(float(loss_frm))\n",
    "                losspot.append(float(loss_pot))\n",
    "                # losscls.append(float(loss_cls))\n",
    "            # Zero gradients, perform a backward pass, and update the weights.\n",
    "            optimizer.zero_grad() # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度\n",
    "            loss.backward() # loss反向传播\n",
    "            optimizer.step() # 反向传播后参数更新 \n",
    "    #print('loss=', loss)\n",
    "    if train_mode=='cls':\n",
    "        print('epoch_train_acc=',right/(right+error))\n",
    "    elif train_mode=='pot':\n",
    "        print('epoch_train_acc=',mae_all/lens)\n",
    "    elif train_mode=='frm':\n",
    "        print('epoch_train_acc=',mae_all/lens)\n",
    "    elif train_mode=='seg':\n",
    "        print('epoch_train_acc=',dice/lens)\n",
    "    elif train_mode=='mtl':\n",
    "        print('epoch_train_frame=',mae_all_frm/lens,\n",
    "              'epoch_train_point=',mae_all_pot/lens,'epoch_train_dice=',dice/lens)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    #测试阶段\n",
    "        #model.eval()\n",
    "        print('evaluating')\n",
    "        dice,jcd, fb = 0,0,0\n",
    "        cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "        mae_all = 0\n",
    "        mae_all_pot = 0\n",
    "        mae_all_frm = 0\n",
    "        right,error = 0,0\n",
    "        with trange(lens2) as tr:\n",
    "            for i in tr:\n",
    "                video_test = image_test[i:i+1,:,:,:,:]\n",
    "                #print('batch=',i,' of ',lens)\n",
    "                labelv_test = label_test[i:i+1,:,:,:,:]\n",
    "                classiv_test = classi_test[i:i+1]\n",
    "                pointi_test = point_test[i:i+1]\n",
    "                framei_test = frame_test[i]\n",
    "                if train_mode == 'cls':\n",
    "                    pred_test_cls = model(video_test)\n",
    "                    tr.set_postfix(loss=float(loss),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]))\n",
    "                    if float(classiv_test) == float(torch.max(pred_test_cls, 1)[1]):\n",
    "                        right = right+1\n",
    "                    else:\n",
    "                        error = error+1\n",
    "                elif train_mode == 'pot':\n",
    "                    pred_pot_test,_ = model(video)\n",
    "                    mae = mae_point(pred_pot_test,pointi_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                    mae_all+=float(mae) \n",
    "                elif train_mode == 'frm':\n",
    "                    pred_frm_test = model(video)\n",
    "                    mae = mae_point(pred_frm_test,framei_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),mae=float(mae))\n",
    "                    mae_all+=float(mae) \n",
    "                elif train_mode=='seg':\n",
    "                    pred_seg_test = model(video)\n",
    "                    dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv_test,pred_seg_test) # 计算指标\n",
    "                    dice = dice + float(dice_coeff)\n",
    "                    hd,md = get_hausdorff(labelv,pred_seg_test)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),dice=float(dice_coeff))\n",
    "                elif train_mode=='mtl':\n",
    "                    pred_pot_test,_,pred_frm_test,pred_seg_test = model(video_test)\n",
    "                    pred_frm_test = torch.tanh(pred_frm_test)\n",
    "                    mae_pot_test = mae_point(pred_pot_test,pointi_test)\n",
    "                    mae_frm_test = mae_point(pred_frm_test,framei_test)\n",
    "                    dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv_test,pred_seg_test) # 计算指标\n",
    "                    mae_all_frm+=float(mae_frm_test)\n",
    "                    mae_all_pot+=float(mae_pot_test)\n",
    "                    dice = dice + float(dice_coeff)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(loss=float(loss),\n",
    "                                   mae_pot=float(mae_pot_test),mae_frm=float(mae_frm_test),dice=float(dice_coeff))\n",
    "                    # if float(classiv_test) == float(torch.max(pred_cls_test, 1)[1]):\n",
    "                    #     right = right+1\n",
    "                    # else:\n",
    "                    #     error = error+1\n",
    "                    \n",
    "            if train_mode=='cls':\n",
    "                print('epoch_train_acc=',right/(right+error))\n",
    "            elif train_mode=='pot':\n",
    "                print('epoch_train_acc=',mae_all/lens2)\n",
    "            elif train_mode=='frm':\n",
    "                print('epoch_train_acc=',mae_all/lens2)\n",
    "            elif train_mode=='seg':\n",
    "                print('epoch_train_acc=',dice/lens2)\n",
    "            elif train_mode=='mtl':\n",
    "                print('epoch_train_point=',mae_all_frm/lens2,\n",
    "                      'epoch_train_frame=',mae_all_pot/lens2,'epoch_train_dice=',dice/lens2)\n",
    "                \n",
    "    if train_mode == 'cls':\n",
    "        cls_acc =   right/(right+error)  \n",
    "        if cls_acc >= cls_less:\n",
    "            print('save model')\n",
    "            cls_less = cls_acc\n",
    "            torch.save(model.state_dict(), './weight_cls/cls_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',cls_less)\n",
    "    elif train_mode == 'pot':\n",
    "        pot_acc =   mae_all/lens \n",
    "        if pot_acc <= pot_less:\n",
    "            print('save model')\n",
    "            pot_less = pot_acc\n",
    "            torch.save(model.state_dict(), './weight_pot/pot_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',pot_less)\n",
    "    elif train_mode == 'frm':\n",
    "        frm_acc =   mae_all/lens2 \n",
    "        if frm_acc <= frm_less:\n",
    "            print('save model')\n",
    "            frm_less = frm_acc\n",
    "            torch.save(model.state_dict(), './weight_frm/frm_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',frm_less)\n",
    "    elif train_mode == 'seg':\n",
    "        seg_acc =   dice/lens2\n",
    "        if seg_acc >= dice_less:\n",
    "            print('save model')\n",
    "            dice_less = seg_acc\n",
    "            torch.save(model.state_dict(), './weight_seg/seg_weights3.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',dice_less)\n",
    "    elif train_mode == 'mtl':\n",
    "        seg_acc =   dice/lens2\n",
    "        if seg_acc >= dice_less:\n",
    "            print('save model')\n",
    "            dice_less = seg_acc\n",
    "            torch.save(model.state_dict(), './weight/mlt_weights_beta1.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',dice_less)\n",
    "scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "#训练分割分支可视化\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision\n",
    "import torch\n",
    "from torchvision.transforms import transforms\n",
    "from torch import nn, optim\n",
    "import timeit\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from tqdm import trange\n",
    "\n",
    "img_size = 128\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "torch.backends.cudnn.benchmark = True\n",
    "l_r = 0.0002  #0.0002\n",
    "device = torch.device(\"cuda\")\n",
    "train_mode = 'mtl'\n",
    "model = Multivit_net(img_dim=128,in_channels=1,out_channels=16,head_num=4,mlp_dim=512,block_num=8,\n",
    "                     patch_dim=16,class_num=1,drop_rate = 0.2,seq_frame=30,mode =train_mode,height=128,weight=128).to(device)\n",
    "#model.load_state_dict(torch.load('./weight_frm/frm_weights.pth'),True)\n",
    "#model.load_state_dict(torch.load('./weight_seg/seg_weights.pth'),True)\n",
    "param_optim = []\n",
    "layers = []\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=l_r)\n",
    "dice_less = 0.85\n",
    "cor_less = 0.85\n",
    "mae_less = 30\n",
    "pot_less = 30\n",
    "frm_less = 30\n",
    "cls_less = 0.8\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "crit_pot = torch.nn.SmoothL1Loss()\n",
    "for t in range(50):\n",
    "    # Forward pass: Compute predicted y by passing x to the model\n",
    "    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [25, 30], 0.1)\n",
    "    lens = image.shape[0]\n",
    "    print('epoch=',t)\n",
    "    dice,jcd, fb = 0,0,0\n",
    "    cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "    mae_all = 0\n",
    "    mae_all_pot = 0\n",
    "    mae_all_frm = 0\n",
    "    right,error = 0,0\n",
    "    image = image\n",
    "    label = label\n",
    "    classi = classi\n",
    "    point = point\n",
    "    pointmap = pointmap\n",
    "    frame = frame\n",
    "    lens2 = image_test.shape[0]\n",
    "    image_test = image_test\n",
    "    label_test = label_test\n",
    "    classi_test = classi_test\n",
    "    point_test = point_test\n",
    "    frame_test = frame_test\n",
    "    model.train()\n",
    "    print('training')\n",
    "    with trange(lens) as tr:\n",
    "        for i in tr:\n",
    "            video = image[i:i+1,:,:,:,:].to(device)\n",
    "            #print('batch=',i,' of ',lens)\n",
    "            labelv = label[i:i+1,:,:,:,:].to(device)\n",
    "            classiv = classi[i:i+1].to(device)\n",
    "            pointi = point[i].to(device)\n",
    "            pointmapi = pointmap[i].to(device)\n",
    "            framei = frame[i].to(device)\n",
    "            if train_mode == 'mtl':\n",
    "                grads = {}\n",
    "                losses_vec = []\n",
    "                n_tasks = 4\n",
    "                # obtain and store the gradient value\n",
    "                for i in range(n_tasks):\n",
    "                    pred_cls,pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                    loss_cls = criterion(pred_cls,classiv.long())\n",
    "                    loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                    loss_pot_map = seg_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                    loss_pot = loss_pot1+loss_pot_map\n",
    "                    pred_frm = torch.tanh(pred_frm)\n",
    "                    loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                    loss1 = seg_loss(pred_seg, labelv)\n",
    "                    loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                    a = t/50\n",
    "                    loss_seg = (1-a)*loss1+a*loss2\n",
    "                    #loss = loss_cls+loss_pot+loss_frm+loss_seg\n",
    "                    task_loss = [loss_cls,loss_pot,loss_frm,loss_seg]\n",
    "                    optimizer.zero_grad()\n",
    "                    losses_vec.append(task_loss[i].data)\n",
    "                    task_loss[i].backward()\n",
    "                    grads[i] = []\n",
    "                    # can use scalable method proposed in the MOO-MTL paper for large scale problem\n",
    "                    # but we keep use the gradient of all parameters in this experiment\n",
    "                    for param in model.parameters():\n",
    "                        if param.grad is not None:\n",
    "                            grads[i].append(Variable(param.grad.data.clone().flatten(), requires_grad=False))\n",
    "                #print(len(grads[1]),len(grads[2]))\n",
    "                grads_list = [torch.cat(grads[i]) for i in range(len(grads))]\n",
    "                #grads = torch.stack(grads_list)\n",
    "\n",
    "                # calculate the weights\n",
    "                losses_vec = torch.stack(losses_vec)\n",
    "                #print('losses_vec',losses_vec,'grad',grads.shape)\n",
    "                npref = 5\n",
    "                ref_vec = torch.tensor(circle_points([1], [npref])[0]).cuda().float()\n",
    "                flag, weight_vec = get_d_paretomtl_init(grads_list,losses_vec,ref_vec,2)\n",
    "\n",
    "                # early stop once a feasible solution is obtained\n",
    "                if flag == True:\n",
    "                    print(\"fealsible solution is obtained.\")\n",
    "                    break\n",
    "\n",
    "                # optimization step\n",
    "                optimizer.zero_grad()\n",
    "                for i in range(len(task_loss)):\n",
    "                    pred_cls,pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                    loss_cls = criterion(pred_cls,classiv.long())\n",
    "                    loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                    loss_pot_map = seg_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                    loss_pot = loss_pot1+loss_pot_map\n",
    "                    pred_frm = torch.tanh(pred_frm)\n",
    "                    loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                    loss1 = seg_loss(pred_seg, labelv)\n",
    "                    loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                    a = t/50\n",
    "                    loss_seg = (1-a)*loss1+a*loss2\n",
    "                    #loss = loss_cls+loss_pot+loss_frm+loss_seg\n",
    "                    task_loss = [loss_cls,loss_pot,loss_frm,loss_seg]\n",
    "                    if i == 0:\n",
    "                        loss_total = weight_vec[i] * task_loss[i]\n",
    "                    else:\n",
    "                        loss_total = loss_total + weight_vec[i] * task_loss[i]\n",
    "\n",
    "                loss_total.backward()\n",
    "                optimizer.step()\n",
    "                mae_pot = mae_point(pred_pot,pointi)\n",
    "                mae_frm = mae_point(pred_frm,framei)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                mae_all_frm+=float(mae_frm)\n",
    "                mae_all_pot+=float(mae_pot)\n",
    "                dice = dice + float(dice_coeff)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss_total),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]),mae_pot=float(mae_pot)\n",
    "                              ,mae_frm=float(mae_frm),dice=float(dice_coeff),task_loss=task_loss,weight=weight_vec)\n",
    "            else:\n",
    "            # continue if no feasible solution is found\n",
    "                continue\n",
    "            # break the loop once a feasible solutions is found\n",
    "            break\n",
    "                \n",
    "    \n",
    "    with trange(lens) as tr:\n",
    "        for i in tr:\n",
    "            video = image[i:i+1,:,:,:,:].to(device)\n",
    "            #print('batch=',i,' of ',lens)\n",
    "            labelv = label[i:i+1,:,:,:,:].to(device)\n",
    "            classiv = classi[i:i+1].to(device)\n",
    "            pointi = point[i].to(device)\n",
    "            pointmapi = pointmap[i].to(device)\n",
    "            framei = frame[i].to(device)\n",
    "            if train_mode == 'mtl':\n",
    "                grads = {}\n",
    "                losses_vec = []\n",
    "                n_tasks = 4\n",
    "                # obtain and store the gradient value\n",
    "                for i in range(n_tasks):\n",
    "                    pred_cls,pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                    loss_cls = criterion(pred_cls,classiv.long())\n",
    "                    loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                    loss_pot_map = seg_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                    loss_pot = loss_pot1+loss_pot_map\n",
    "                    pred_frm = torch.tanh(pred_frm)\n",
    "                    loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                    loss1 = seg_loss(pred_seg, labelv)\n",
    "                    loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                    a = t/50\n",
    "                    loss_seg = (1-a)*loss1+a*loss2\n",
    "                    #loss = loss_cls+loss_pot+loss_frm+loss_seg\n",
    "                    task_loss = [loss_cls,loss_pot,loss_frm,loss_seg]\n",
    "                    optimizer.zero_grad()\n",
    "                    losses_vec.append(task_loss[i].data)\n",
    "                    task_loss[i].backward()\n",
    "                    grads[i] = []\n",
    "                    # can use scalable method proposed in the MOO-MTL paper for large scale problem\n",
    "                    # but we keep use the gradient of all parameters in this experiment\n",
    "                    for param in model.parameters():\n",
    "                        if param.grad is not None:\n",
    "                            grads[i].append(Variable(param.grad.data.clone().flatten(), requires_grad=False))\n",
    "                #print(len(grads[1]),len(grads[2]))\n",
    "                grads_list = [torch.cat(grads[i]) for i in range(len(grads))]\n",
    "                grads = torch.stack(grads_list)\n",
    "\n",
    "                # calculate the weights\n",
    "                losses_vec = torch.stack(losses_vec)\n",
    "                #print('losses_vec',losses_vec,'grad',grads.shape)\n",
    "                npref = 5\n",
    "                ref_vec = torch.tensor(circle_points([1], [npref])[0]).cuda().float()\n",
    "                weight_vec = get_d_paretomtl(grads,losses_vec,ref_vec,2)\n",
    "                normalize_coeff = n_tasks / torch.sum(torch.abs(weight_vec))\n",
    "                weight_vec = weight_vec * normalize_coeff\n",
    "\n",
    "                # optimization step\n",
    "                optimizer.zero_grad()\n",
    "                for i in range(len(task_loss)):\n",
    "                    pred_cls,pred_pot,pred_potmap,pred_frm,pred_seg = model(video)\n",
    "                    loss_cls = criterion(pred_cls,classiv.long())\n",
    "                    loss_pot1 = crit_pot(pred_pot[0].to(torch.float32),pointi.to(torch.float32))\n",
    "                    loss_pot_map = seg_loss(pred_potmap.to(torch.float32), pointmapi.to(torch.float32))\n",
    "                    loss_pot = loss_pot1+loss_pot_map\n",
    "                    pred_frm = torch.tanh(pred_frm)\n",
    "                    loss_frm = crit_pot(pred_frm.to(torch.float32),framei.to(torch.float32))\n",
    "                    loss1 = seg_loss(pred_seg, labelv)\n",
    "                    loss2 = boundary_cos_loss(labelv,pred_seg)\n",
    "                    a = t/50\n",
    "                    loss_seg = (1-a)*loss1+a*loss2\n",
    "                    #loss = loss_cls+loss_pot+loss_frm+loss_seg\n",
    "                    task_loss = [loss_cls,loss_pot,loss_frm,loss_seg]\n",
    "                    if i == 0:\n",
    "                        loss_total = weight_vec[i] * task_loss[i]\n",
    "                    else:\n",
    "                        loss_total = loss_total + weight_vec[i] * task_loss[i]\n",
    "\n",
    "                loss_total.backward()\n",
    "                optimizer.step()\n",
    "                mae_pot = mae_point(pred_pot,pointi)\n",
    "                mae_frm = mae_point(pred_frm,framei)\n",
    "                dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv,pred_seg) # 计算指标\n",
    "                mae_all_frm+=float(mae_frm)\n",
    "                mae_all_pot+=float(mae_pot)\n",
    "                dice = dice + float(dice_coeff)\n",
    "                tr.set_description('batch= %i' % i)\n",
    "                tr.set_postfix(loss=float(loss_total),class_ori=classiv,class_pred=float(torch.max(pred_cls, 1)[1]),mae_pot=float(mae_pot)\n",
    "                              ,mae_frm=float(mae_frm),dice=float(dice_coeff),weight=weight_vec)\n",
    "                if float(classiv) == float(torch.max(pred_cls, 1)[1]):\n",
    "                    right = right+1\n",
    "                else:\n",
    "                    error = error+1\n",
    "            # Zero gradients, perform a backward pass, and update the weights.\n",
    "            # optimizer.zero_grad() # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度\n",
    "            # loss.backward() # loss反向传播\n",
    "            # optimizer.step() # 反向传播后参数更新 \n",
    "    #print('loss=', loss)\n",
    "    if train_mode=='mtl':\n",
    "        print('epoch_train_class=',right/(right+error),'epoch_train_point=',mae_all_frm/lens,\n",
    "              'epoch_train_frame=',mae_all_pot/lens,'epoch_train_dice=',dice/lens)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    #测试阶段\n",
    "        #model.eval()\n",
    "        print('evaluating')\n",
    "        dice,jcd, fb = 0,0,0\n",
    "        cor_lva ,cor_mvd,cor_lvd = 0,0,0\n",
    "        mae_all = 0\n",
    "        mae_all_pot = 0\n",
    "        mae_all_frm = 0\n",
    "        right,error = 0,0\n",
    "        with trange(lens2) as tr:\n",
    "            for i in tr:\n",
    "                video_test = image_test[i:i+1,:,:,:,:].to(device)\n",
    "                #print('batch=',i,' of ',lens)\n",
    "                labelv_test = label_test[i:i+1,:,:,:,:].to(device)\n",
    "                classiv_test = classi_test[i:i+1].to(device)\n",
    "                pointi_test = point_test[i].to(device)\n",
    "                framei_test = frame_test[i].to(device)\n",
    "                if train_mode=='mtl':\n",
    "                    pred_cls_test,pred_pot_test,_,pred_frm_test,pred_seg_test = model(video_test)\n",
    "                    pred_frm_test = torch.tanh(pred_frm_test)\n",
    "                    mae_pot_test = mae_point(pred_pot_test,pointi_test)\n",
    "                    mae_frm_test = mae_point(pred_frm_test,framei_test)\n",
    "                    dice_coeff, jaccard_coeff, f_beta_coeff = coefficients(labelv_test,pred_seg_test) # 计算指标\n",
    "                    mae_all_frm+=float(mae_frm_test)\n",
    "                    mae_all_pot+=float(mae_pot_test)\n",
    "                    dice = dice + float(dice_coeff)\n",
    "                    tr.set_description('batch= %i' % i)\n",
    "                    tr.set_postfix(class_ori=classiv_test,class_pred=float(torch.max(pred_cls_test, 1)[1]),\n",
    "                                   mae_pot=float(mae_pot_test),mae_frm=float(mae_frm_test),dice=float(dice_coeff))\n",
    "                    if float(classiv_test) == float(torch.max(pred_cls_test, 1)[1]):\n",
    "                        right = right+1\n",
    "                    else:\n",
    "                        error = error+1\n",
    "                    \n",
    "            if train_mode=='mtl':\n",
    "                print('epoch_train_class=',right/(right+error),'epoch_train_point=',mae_all_frm/lens2,\n",
    "                      'epoch_train_frame=',mae_all_pot/lens2,'epoch_train_dice=',dice/lens2)\n",
    "                \n",
    "    if train_mode == 'mtl':\n",
    "        seg_acc =   dice/lens2\n",
    "        if seg_acc >= dice_less:\n",
    "            print('save model')\n",
    "            dice_less = seg_acc\n",
    "            torch.save(model.state_dict(), './weight/mlt_weights.pth')    \n",
    "        else:\n",
    "            print('not save,the best cor is:',dice_less)\n",
    "scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "#traning set\n",
    "import torch\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import OrdinalEncoder\n",
    "import numpy.matlib\n",
    "from torchvision.transforms import transforms\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn, optim\n",
    "import timeit\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from tqdm import trange\n",
    "import os\n",
    "import sys\n",
    "import random\n",
    "import itertools\n",
    "import colorsys\n",
    "\n",
    "import numpy as np\n",
    "from skimage.measure import find_contours\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import patches,  lines\n",
    "from matplotlib.patches import Polygon\n",
    "import IPython.display\n",
    "\n",
    "from torchvision.transforms import transforms\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "unloader = transforms.ToPILImage()\n",
    "\n",
    "def tensor_to_img(tensor_pred , data_num, frame_num,title=None):\n",
    "    #image = unloader(image)\n",
    "    #tensor_pred = tensor_pred.type(torch.float32)\n",
    "    #pred = tensor_pred.cpu().clone()  # we clone the tensor to not do changes on it\n",
    "    pred = tensor_pred\n",
    "    pred = unloader(pred)\n",
    "    #fig = plt.figure()\n",
    "    pred = np.array(pred)\n",
    "    pred = cv2.resize(pred, (256,256) , interpolation=cv2.INTER_AREA) \n",
    "    #plt.imshow(pred, cmap = 'gray')\n",
    "    return pred\n",
    "\n",
    "def tensor_to_lb(tensor_pred , data_num, frame_num,title=None):\n",
    "    #image = unloader(image)\n",
    "    tensor_pred = torch.sigmoid(tensor_pred)\n",
    "    tensor_pred = torch.gt(tensor_pred, 0.5)\n",
    "    tensor_pred = tensor_pred.type(torch.float32)\n",
    "    pred = tensor_pred.cpu().clone()  # we clone the tensor to not do changes on it\n",
    "    pred = unloader(pred)\n",
    "    #fig = plt.figure()\n",
    "    pred = np.array(pred)\n",
    "    pred = cv2.resize(pred, (256,256) , interpolation=cv2.INTER_AREA) \n",
    "    #plt.imshow(pred, cmap = 'gray')\n",
    "    return pred\n",
    "\n",
    "\n",
    "def random_colors(N, bright=True):\n",
    "    \"\"\"\n",
    "    Generate random colors.\n",
    "    To get visually distinct colors, generate them in HSV space then\n",
    "    convert to RGB.\n",
    "    \"\"\"\n",
    "    brightness = 1.0 if bright else 0.7\n",
    "    hsv = [(i / N, 1, brightness) for i in range(N)]\n",
    "    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))\n",
    "    random.shuffle(colors)\n",
    "    return colors\n",
    "\n",
    "\n",
    "def apply_mask(image, mask, color, alpha=0.4):\n",
    "    \"\"\"Apply the given mask to the image.\n",
    "    \"\"\"\n",
    "    image = image.copy()\n",
    "    mask_out = cv2.Canny(mask.astype(np.uint8),0,1)\n",
    "    kernel = np.ones((2, 2), dtype=np.uint8)\n",
    "    mask_out = cv2.dilate(mask_out, kernel, 1)\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(mask >= 0.5,\n",
    "                                  image[:, :, c] *\n",
    "                                  (1 - alpha) + alpha * color[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(mask_out >= 0.5,\n",
    "                                  color[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    return image\n",
    "\n",
    "unloader = transforms.ToPILImage()\n",
    "train_mode = 'mtl'\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "torch.backends.cudnn.benchmark = True\n",
    "l_r = 0.0002  #0.0002\n",
    "device = torch.device(\"cuda\")\n",
    "def fill_contour(img):\n",
    "    contours, _ = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    n = len(contours)  # 轮廓的个数\n",
    "    max_area = 0\n",
    "    for i  in range(n):\n",
    "        if cv2.contourArea(contours[i]) > max_area:\n",
    "            max_area = cv2.contourArea(contours[i])\n",
    "    cv_contours = []\n",
    "    if n != 1:\n",
    "        for contour in contours:\n",
    "            area = cv2.contourArea(contour)\n",
    "            if area < max_area:\n",
    "                cv_contours.append(contour)\n",
    "                x, y, w, h = cv2.boundingRect(contour)\n",
    "                img[y:y + h, x:x + w] = 1\n",
    "            else:\n",
    "                continue\n",
    "    else:\n",
    "        pass\n",
    "    \n",
    "    return img\n",
    "\n",
    "def tensor_im(tensor_pred , data_num, frame_num,title=None):\n",
    "    #tensor_pred = torch.sigmoid(tensor_pred)\n",
    "    #tensor_pred = torch.gt(tensor_pred, 0.5)\n",
    "    tensor_pred = tensor_pred.type(torch.float32)\n",
    "    pred = tensor_pred.cpu().clone()  # we clone the tensor to not do changes on it\n",
    "    pred = unloader(pred)\n",
    "    pred = np.array(pred)\n",
    "    #plt.imshow(pred)\n",
    "    return pred\n",
    "    \n",
    "def tensor_save(tensor_pred , data_num, frame_num,title=None):\n",
    "    tensor_pred = torch.sigmoid(tensor_pred)\n",
    "    tensor_pred = torch.gt(tensor_pred, 0.5)\n",
    "    tensor_pred = tensor_pred.type(torch.float32)\n",
    "    pred = tensor_pred.cpu().clone()  # we clone the tensor to not do changes on it\n",
    "    pred = unloader(pred)\n",
    "    kernel1 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5, 5))\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(10, 10))\n",
    "    pred = np.array(pred)\n",
    "    #plt.imshow(pred)\n",
    "    return pred\n",
    "\n",
    "def landmark(center_x,center_y,IMAGE_HEIGHT, IMAGE_WIDTH):\n",
    "    R = np.sqrt(1**1 + 1**1)\n",
    "    Gauss_map = np.zeros((IMAGE_HEIGHT, IMAGE_WIDTH))\n",
    "    # 直接利用矩阵运算实现\n",
    "    mask_x = np.matlib.repmat(center_x, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "    mask_y = np.matlib.repmat(center_y, IMAGE_HEIGHT, IMAGE_WIDTH)\n",
    "    x1 = np.arange(IMAGE_WIDTH)\n",
    "    x_map = np.matlib.repmat(x1, IMAGE_HEIGHT, 1)\n",
    "    y1 = np.arange(IMAGE_HEIGHT)\n",
    "    y_map = np.matlib.repmat(y1, IMAGE_WIDTH, 1)\n",
    "    y_map = np.transpose(y_map)\n",
    "    Gauss_map = np.sqrt((x_map-mask_x)**2+(y_map-mask_y)**2)\n",
    "    Gauss_map = np.exp(-0.5*Gauss_map/R)\n",
    "    return Gauss_map\n",
    "\n",
    "def norm(img):\n",
    "    img = np.array(img, dtype=np.float32)\n",
    "    img -= np.mean(img)\n",
    "    img /= (np.std(img) + 1e-12)\n",
    "    return img\n",
    "    \n",
    "def locmap(pot):\n",
    "    gauss_batch = []\n",
    "    for i in range(0,pot.shape[0]):\n",
    "        gauss_tp = []\n",
    "        for j in range(0,pot.shape[1]):\n",
    "            g_map1 = landmark(pot[i, j, 0, 0],pot[i, j, 0, 1],128,128)\n",
    "            g_map2 = landmark(pot[i, j, 1, 0],pot[i, j, 1, 1],128,128)\n",
    "            g_map3 = landmark(pot[i, j, 2, 0],pot[i, j, 2, 1],128,128)\n",
    "            g_map4 = landmark(pot[i, j, 3, 0],pot[i, j, 3, 1],128,128)\n",
    "            Gauss_map = [g_map1,g_map2,g_map3,g_map4]\n",
    "            Gauss_map = norm(Gauss_map)\n",
    "            gauss_tp.append(Gauss_map)\n",
    "        gauss_batch.append(gauss_tp)\n",
    "    gauss_batch = np.array(gauss_batch)[:, :, :, :, :]\n",
    "    return gauss_batch\n",
    "\n",
    "def point_color(potmap):\n",
    "    image = np.zeros((128,128,3),dtype=int)\n",
    "    kernel = np.ones((2, 2), dtype=np.uint8)\n",
    "    color1 = [1,1,0]\n",
    "    color2 = [0,1,1]\n",
    "    color3 = [1,0,1]\n",
    "    color4 = [1,0,0]\n",
    "    alpha = 1\n",
    "    k = 10\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(potmap[0, :, :] >= k,alpha * color1[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(potmap[1, :, :] >= k,alpha * color2[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(potmap[2, :, :] >= k,alpha * color3[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    for c in range(3):\n",
    "        image[:, :, c] = np.where(potmap[3, :, :] >= k,alpha * color4[c] * 255,\n",
    "                                  image[:, :, c])\n",
    "    return image\n",
    "    \n",
    "model = Multivit_net(img_dim=128,in_channels=1,out_channels=128,head_num=4,mlp_dim=512,block_num=8,\n",
    "                     patch_dim=16,class_num=1,drop_rate = 0.2,seq_frame=30,mode =train_mode,height=128,weight=128).to(device)\n",
    "model.load_state_dict(torch.load('./weight/mlt_weights_beta0.5.pth'),True)\n",
    "model.train()\n",
    "size = 128\n",
    "ki = image_test.to(device)\n",
    "kg = label_test.to(device)\n",
    "kp = point_test.to(device)\n",
    "kf = frame_test.to(device)\n",
    "for i in range(99,109):\n",
    "    a = ki[i:i+1]\n",
    "    d = kg[i:i+1]\n",
    "    c = kp[i:i+1]\n",
    "    e = kf[i:i+1]\n",
    "    plt.subplots(figsize=(15,40))\n",
    "    color_lvla = [(1,0,0),(0,1,0)]\n",
    "    with torch.no_grad():\n",
    "        b,b1,b2,b3,b4 = model(a)\n",
    "        #torch.tanh(b)\n",
    "        print(a.shape,c.shape,b.shape,b1.shape,b2.shape,b3.shape,b4.shape)\n",
    "        print(b1[0,0,:,:],c[0,0,:,:])\n",
    "        print(e,b3)\n",
    "        for k in range(2):\n",
    "            j = k\n",
    "            plt.subplot(161)\n",
    "            img = tensor_im(a[0,j,0,:,:].to(torch.uint8),0,1).astype(np.uint8)\n",
    "            alp,bet = 53,11\n",
    "            plt.imshow(img*alp+bet,cmap = 'gray')\n",
    "            plt.axis('off')\n",
    "            image_ori = np.zeros((128,128,3),dtype=int)\n",
    "            image_ori[:,:,0] = (img*alp+bet)[:,:]\n",
    "            image_ori[:,:,1] = (img*alp+bet)[:,:]\n",
    "            image_ori[:,:,2] = (img*alp+bet)[:,:]\n",
    "            plt.subplot(162)\n",
    "            gt = tensor_save(d[0,j,0,:,:].to(torch.float32),0,1)\n",
    "            gt = apply_mask(image_ori,gt,color_lvla[1])\n",
    "            plt.imshow(gt)\n",
    "            plt.axis('off')\n",
    "            plt.subplot(163)\n",
    "            pred = tensor_save(b4[0,j,0,:,:].to(torch.float32),0,1)\n",
    "            pred = apply_mask(image_ori,pred,color_lvla[1])\n",
    "            plt.imshow(pred)\n",
    "            plt.axis('off')\n",
    "            plt.subplot(164)\n",
    "            x = locmap(np.array(c.cpu()))[0,j,:,:]\n",
    "            max_index = np.unravel_index(np.argmax(x, axis=None), x.shape)\n",
    "            max_value = x[max_index]\n",
    "            #print(max_index,max_value)\n",
    "            comap = point_color(locmap(np.array(c.cpu()))[0,j,:,:])\n",
    "            plt.imshow(comap)\n",
    "            plt.axis('off')\n",
    "            plt.subplot(165)\n",
    "            b1map = point_color(locmap(np.array(b1.cpu()))[0,j,:,:])\n",
    "            plt.imshow(b1map)\n",
    "            plt.axis('off')\n",
    "            plt.subplot(166)\n",
    "            b1map = point_color(locmap(np.array(b1.cpu()))[0,j,:,:])\n",
    "            # plt.imshow(b1map)\n",
    "            # plt.axis('off')\n",
    "            peak_list = np.array(e.cpu()[0])\n",
    "            # print(peak_list.shape)\n",
    "            if peak_list[j] == -1:\n",
    "                plt.imshow(img,cmap = 'PuRd_r')\n",
    "            elif peak_list[j] == 0:\n",
    "                plt.imshow(img)\n",
    "            else:\n",
    "                plt.imshow(img,cmap = 'BrBG_r')\n",
    "            plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}