liversegmentation_daTest_0416_selfsup.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/kevin801221/CNN-and-computer-vision/blob/master/liversegmentation_daTest_0416_selfsup.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0503cc44",
      "metadata": {
        "id": "0503cc44"
      },
      "outputs": [],
      "source": [
        "# !pip install -q --user albumentations\n",
        "# !pip3 install pydicom\n",
        "# !git clone https://github.com/June103310110/Image_Segmentation.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bc8a8a40",
      "metadata": {
        "id": "bc8a8a40"
      },
      "outputs": [],
      "source": [
        "# import argparse\n",
        "# parser = argparse.ArgumentParser()\n",
        "# parser.add_argument(\"--lamb\",type=float)\n",
        "# parser.add_argument(\"--train_size\", type=float)\n",
        "# parser.add_argument(\"--epoch\", type=int)\n",
        "# args = parser.parse_args()\n",
        "# lamb_dc = args.lamb\n",
        "# train_size = args.train_size\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d37f81bb",
      "metadata": {
        "id": "d37f81bb"
      },
      "outputs": [],
      "source": [
        "BATCH_SIZE = 8 # 8 for 256x256/ 16 for 128x128\n",
        "NUM_LABELS = 1\n",
        "WIDTH = 256\n",
        "HEIGHT = 256 \n",
        "MULTI_CHANNELS = False\n",
        "device = 'cuda:0'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0272c6c4",
      "metadata": {
        "id": "0272c6c4"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import os\n",
        "sys.path.append('Image_Segmentation/PART2')\n",
        "os.makedirs('saved', exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "06e17e51",
      "metadata": {
        "id": "06e17e51"
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "08f06e22",
      "metadata": {
        "id": "08f06e22"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import numpy as np\n",
        "import torch # 1.9\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import albumentations as A\n",
        "import os\n",
        "import torch.optim as optim\n",
        "import time\n",
        "import ipywidgets as widgets\n",
        "import pickle\n",
        "\n",
        "# 從repo裡面導入套件\n",
        "from utils import show_image_mask, mask_CutMix#, patience\n",
        "from unet import UNet\n",
        "\n",
        "\n",
        "# 導入dicom套件\n",
        "from pydicom import dcmread\n",
        "from pydicom.data import get_testdata_files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "42238a87",
      "metadata": {
        "scrolled": true,
        "id": "42238a87"
      },
      "outputs": [],
      "source": [
        "# !pip install --upgrade gdown\n",
        "# import gdown \n",
        "# url = \"https://drive.google.com/u/1/uc?id=15EVEIB2o-sJJd0qOjoFbZrz4briIYoeF&export=download\"\n",
        "# output = \"full_data.zip\"\n",
        "# gdown.download(url, output)\n",
        "# !rm -rf data\n",
        "# !mkdir data\n",
        "# os.system('unzip -o full_data.zip -d ./data')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "676eefc1",
      "metadata": {
        "id": "676eefc1"
      },
      "outputs": [],
      "source": [
        "def show_image(*img_):\n",
        "    for i in img_:\n",
        "        assert i.__class__.__name__ == 'ndarray', 'imput data type should be ndarray'\n",
        "\n",
        "    plt.figure(figsize=(10,3))\n",
        "    for i, img in enumerate(list(img_), 1):\n",
        "        plt.subplot(1,len(img_),i)\n",
        "\n",
        "        if len(np.shape(img)) == 2 or np.shape(img)[-1] == 1:\n",
        "            plt.imshow(img, cmap='gray')\n",
        "        elif len(np.shape(img)) == 3:\n",
        "            plt.imshow(img)\n",
        "    plt.show()\n",
        "    plt.close()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9ed02d72",
      "metadata": {
        "id": "9ed02d72"
      },
      "outputs": [],
      "source": [
        "# root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/1_Domain_Gernalization_dataset/Train_Sets/MRI/1/T2SPIR/'\n",
        "# dcm_lis = sorted(os.listdir(root+'DICOM_anon/'))\n",
        "# png_lis = sorted(os.listdir(root+'Ground/'))\n",
        "\n",
        "# for dcm, png in zip(dcm_lis, png_lis):\n",
        "#     print(dcm)\n",
        "#     # print(dcm)\n",
        "#     ds = dcmread(root+'DICOM_anon/'+dcm)\n",
        "#     img = ds.pixel_array\n",
        "\n",
        "#     png = cv2.imread(root+'Ground/'+png)[...,0]\n",
        "#     show_image(img, png)\n",
        "#     del img, png\n",
        "#     break\n",
        "# # lis"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8b004bb8",
      "metadata": {
        "id": "8b004bb8"
      },
      "source": [
        "### 取得image list\n",
        "輸出: data_dic (字典)\n",
        "- key: X_train, X_test, y_train, y_test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2c998230",
      "metadata": {
        "id": "2c998230",
        "outputId": "ddd14d8c-370c-46c9-c1d1-e316d03ca01a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/\n"
          ]
        }
      ],
      "source": [
        "# root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/1_Domain_Gernalization_dataset/Train_Sets/CT/'\n",
        "root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/'\n",
        "print(root)\n",
        "dic = {}\n",
        "for a,b,c in os.walk(root, topdown=True):\n",
        "    if len(c)>0: # 當前目錄內包含檔案\n",
        "        if not a.__contains__('OutPhase'):\n",
        "            dic[a] = c\n",
        "dataset = {}\n",
        "lis = ['CT']\n",
        "for task in lis:\n",
        "    class_lis = []\n",
        "    for sub_folder in dic.keys():\n",
        "        if task in sub_folder.split('/'):\n",
        "            class_lis+=[sub_folder+'/'+filename for filename in dic[sub_folder]]\n",
        "    dataset[task] = class_lis\n",
        "    \n",
        "dataset['CT_train'] = sorted([i for i in dataset['CT'] if 'dcm' in i])\n",
        "dataset['CT_mask'] = sorted([i for i in dataset['CT'] if 'png' in i])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cfff55e9",
      "metadata": {
        "id": "cfff55e9",
        "outputId": "d9d42dd3-4eff-4125-f5dd-a6493c68420c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/\n"
          ]
        }
      ],
      "source": [
        "# root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/1_Domain_Gernalization_dataset/Train_Sets/MRI/'\n",
        "root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/'\n",
        "print(root)\n",
        "dic = {}\n",
        "for a,b,c in os.walk(root, topdown=True):\n",
        "    if len(c)>0: # 當前目錄內包含檔案\n",
        "        if not a.__contains__('OutPhase'):\n",
        "            dic[a] = c\n",
        "            \n",
        "lis = ['MRI']\n",
        "for task in lis:\n",
        "    class_lis = []\n",
        "    for sub_folder in dic.keys():\n",
        "        if task in sub_folder.split('/'):\n",
        "            class_lis+=[sub_folder+'/'+filename for filename in dic[sub_folder]]\n",
        "    dataset[task] = class_lis\n",
        "\n",
        "# dataset['CT_train'] = sorted([i for i in dataset['CT'] if 'dcm' in i])\n",
        "# dataset['CT_mask'] = sorted([i for i in dataset['CT'] if 'png' in i])\n",
        "\n",
        "dataset['MRI_DICOM_anon'] = sorted([i for i in dataset['MRI'] if 'dcm' in i])\n",
        "dataset['MRI_Ground'] = sorted([i for i in dataset['MRI'] if 'png' in i])\n",
        "\n",
        "# dataset['MRI_T1DUAL_train'] = sorted([i for i in dataset['MRI_DICOM_anon'] if 'T1DUAL' in i])\n",
        "# dataset['MRI_T1DUAL_mask'] = sorted([i for i in dataset['MRI_Ground'] if 'T1DUAL' in i])\n",
        "\n",
        "dataset['MRI_T2SPIR_train'] = sorted([i for i in dataset['MRI_DICOM_anon'] if 'T2SPIR' in i])\n",
        "dataset['MRI_T2SPIR_mask'] = sorted([i for i in dataset['MRI_Ground'] if 'T2SPIR' in i])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bf83ac87",
      "metadata": {
        "id": "bf83ac87",
        "outputId": "f14d92fc-d447-48a1-9e5b-679ce6de588b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_nonLabel/\n"
          ]
        }
      ],
      "source": [
        "root = '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_nonLabel/'\n",
        "print(root)\n",
        "dic = {}\n",
        "for a,b,c in os.walk(root, topdown=True):\n",
        "    if len(c)>0: # 當前目錄內包含檔案\n",
        "        if not a.__contains__('OutPhase'):\n",
        "            dic[a] = c\n",
        "            \n",
        "lis = ['MRI']\n",
        "for task in lis:\n",
        "    class_lis = []\n",
        "    for sub_folder in dic.keys():\n",
        "        if task in sub_folder.split('/'):\n",
        "            class_lis+=[sub_folder+'/'+filename for filename in dic[sub_folder]]\n",
        "    dataset[task] = class_lis\n",
        "\n",
        "dataset['MRI_DICOM_anon'] = sorted([i for i in dataset['MRI'] if 'dcm' in i])\n",
        "\n",
        "dataset['MRI_T2SPIR_nonLabel_train'] = sorted([i for i in dataset['MRI_DICOM_anon'] if 'T2SPIR' in i])\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0c7bbfcf",
      "metadata": {
        "id": "0c7bbfcf",
        "outputId": "6a374fe0-d2f9-4ebe-d5ee-69843559aab8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/DICOM_anon/i0000,0000b.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/DICOM_anon/i0001,0000b.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/DICOM_anon/i0002,0000b.dcm']\n",
            "['/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/Ground/liver_GT_000.png', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/Ground/liver_GT_001.png', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/1/Ground/liver_GT_002.png']\n",
            "['/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/DICOM_anon/IMG-0002-00001.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/DICOM_anon/IMG-0002-00002.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/DICOM_anon/IMG-0002-00003.dcm']\n",
            "['/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/Ground/IMG-0002-00001.png', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/Ground/IMG-0002-00002.png', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/1/T2SPIR/Ground/IMG-0002-00003.png']\n",
            "['/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_nonLabel/10/T2SPIR/DICOM_anon/IMG-0043-00001.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_nonLabel/10/T2SPIR/DICOM_anon/IMG-0043-00002.dcm', '/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_nonLabel/10/T2SPIR/DICOM_anon/IMG-0043-00003.dcm']\n"
          ]
        }
      ],
      "source": [
        "print(dataset['CT_train'][:3])\n",
        "print(dataset['CT_mask'][:3])\n",
        "print(dataset['MRI_T2SPIR_train'][:3])\n",
        "print(dataset['MRI_T2SPIR_mask'][:3])\n",
        "print(dataset['MRI_T2SPIR_nonLabel_train'][:3])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e2004ecd",
      "metadata": {
        "id": "e2004ecd"
      },
      "outputs": [],
      "source": [
        "domains = ['CT', 'MRI_T2SPIR']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "15769e5f",
      "metadata": {
        "id": "15769e5f"
      },
      "outputs": [],
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "dic = {}\n",
        "for domain in domains:\n",
        "    train_key = f'{domain}_train'\n",
        "    mask_key = f'{domain}_mask'\n",
        "\n",
        "    train, test = train_test_split(list(zip(dataset[train_key], dataset[mask_key])),\n",
        "                                 train_size = 0.8)\n",
        "    dic[domain] = {'train':train, 'test':test}\n",
        "dataset['TTsplit'] = dic\n",
        "\n",
        "  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4db86dbc",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4db86dbc",
        "outputId": "d6cf277d-1ab0-4af8-af3d-1d2b7489d613"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "dict_keys(['CT', 'MRI_T2SPIR'])\n"
          ]
        }
      ],
      "source": [
        "# T1DUAL 我們不使用\n",
        "# 任務: liver segmentation, 從CT轉到MRI(T1DUAL)\n",
        "print(dataset['TTsplit'].keys())\n",
        "MRI_train = dataset['TTsplit']['MRI_T2SPIR']['train']\n",
        "MRI_test = dataset['TTsplit']['MRI_T2SPIR']['test']\n",
        "\n",
        "CT_train = dataset['TTsplit']['CT']['train']\n",
        "CT_test = dataset['TTsplit']['CT']['test']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0486bee4",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0486bee4",
        "outputId": "039ba98f-a2da-45aa-f713-f3613ee369c0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "針對影像/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/CT/6/Ground/liver_GT_035.png做檢查，確定為什麼ground Truth是3通道\n",
            "(512, 512, 3)\n",
            "0, 1通道是否完全相同 True\n",
            "0, 2通道是否完全相同 True\n",
            "針對影像/home/jovyan/DA/DATA/ST_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/MRI/MRI_Label/3/T2SPIR/Ground/IMG-0002-00022.png做檢查，確定為什麼ground Truth是3通道\n",
            "(320, 320, 3)\n",
            "0, 1通道是否完全相同 True\n",
            "0, 2通道是否完全相同 True\n",
            "MRI [0, 63, 252, 126, 189]\n",
            "CT [0, 255]\n"
          ]
        }
      ],
      "source": [
        "png = MRI_train[0][1]\n",
        "png_lis = [CT_train[0][1], MRI_train[0][1]]\n",
        "for png in png_lis:\n",
        "    print(f'針對影像{png}做檢查，確定為什麼ground Truth是3通道')\n",
        "    mask = cv2.imread(png)\n",
        "    print(mask.shape)\n",
        "    for i in range(1,3):\n",
        "        print(f'0, {i}通道是否完全相同', (mask[0]==mask[i]).any())\n",
        "\n",
        "'''\n",
        "Ground Truth中的不同數值代表一種被標註的器官的代號，\n",
        "我們的專題只使用liver，他的代號是63\n",
        "'''\n",
        "unique_ele = []\n",
        "for dcm, png in [*MRI_train[:]]:\n",
        "    mask = cv2.imread(png)[...,0]\n",
        "    unique_ele += [i for i in np.unique(mask) if not unique_ele.__contains__(i)]\n",
        "print('MRI', unique_ele)\n",
        "\n",
        "'''\n",
        "CT的影像只有0或255，\n",
        "0是背景，255是前景，也就是肝臟\n",
        "'''\n",
        "unique_ele = []\n",
        "for dcm, png in [*CT_train[:]]:\n",
        "    mask = cv2.imread(png)[...,0]\n",
        "    unique_ele += [i for i in np.unique(mask) if not unique_ele.__contains__(i)]\n",
        "print('CT', unique_ele)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6ba607b9",
      "metadata": {
        "id": "6ba607b9"
      },
      "source": [
        "#### MRI影像"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7852426a",
      "metadata": {
        "id": "7852426a"
      },
      "outputs": [],
      "source": [
        "# from pydicom import dcmread\n",
        "# from pydicom.data import get_testdata_files\n",
        "# import cv2\n",
        "\n",
        "# for dcm, png in [*MRI_train[:3]]:\n",
        "#     print(dcm, png)\n",
        "#     ds = dcmread(dcm)\n",
        "#     img = ds.pixel_array\n",
        "\n",
        "#     mask = cv2.imread(png)[...,0] # 3通道完全相同\n",
        "#     print(img.shape, mask.shape) \n",
        "\n",
        "#     liver_mask = mask.copy()\n",
        "#     liver_value = 63 # 數值63的部分是肝臟(liver)/ 126右腎/ 189左腎/ 252脾臟\n",
        "#     liver_mask[liver_mask!=liver_value] = 0 \n",
        "#     liver_mask[liver_mask==liver_value] = 1\n",
        "#     show_image(img, mask, liver_mask)\n",
        "#     del img, mask, liver_mask"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f7b537ca",
      "metadata": {
        "id": "f7b537ca"
      },
      "source": [
        "#### CT影像"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fa4a1fb1",
      "metadata": {
        "id": "fa4a1fb1"
      },
      "outputs": [],
      "source": [
        "# for dcm, png in [*CT_train[:3]]:\n",
        "#     print(dcm, png)\n",
        "#     ds = dcmread(dcm)\n",
        "#     img = ds.pixel_array\n",
        "\n",
        "#     mask = cv2.imread(png)[...,0] # 3通道完全相同\n",
        "#     '''\n",
        "#     CT的影像只有0或255，\n",
        "#     0是背景，255是前景，也就是肝臟\n",
        "#     '''\n",
        "#     print(img.shape, mask.shape) \n",
        "\n",
        "#     show_image(img, mask)\n",
        "#     del img, mask"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "391ee489",
      "metadata": {
        "id": "391ee489"
      },
      "source": [
        "### Build torch dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e896e954",
      "metadata": {
        "id": "e896e954"
      },
      "outputs": [],
      "source": [
        "#  https://pytorch.org/tutorials/beginner/basics/data_tutorial.html\n",
        "\n",
        "class CTMRI_ImageDataset(Dataset):\n",
        "    def __init__(self, imgs_anno_path_list, dtype, transform=None):\n",
        "        self.imgs_anno_path_list = imgs_anno_path_list\n",
        "        if type(imgs_anno_path_list) != list:\n",
        "            raise ValueError('Need Input a list')\n",
        "        self.transform = transform\n",
        "        self.dtype = dtype\n",
        "#   \n",
        "    def __len__(self):\n",
        "        return len(self.imgs_anno_path_list)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # now = time.time()\n",
        "        img_anno_path = self.imgs_anno_path_list[idx]\n",
        "        # print(img_anno_path)\n",
        "        # file_name = img_path.split('/')[-1].split('.')[0]\n",
        "        img_path = img_anno_path[0]\n",
        "        mask_path = img_anno_path[1]\n",
        "        # mask_path = [i for i in self.anno_path_list if i.__contains__(file_name)][0]\n",
        "\n",
        "        '''\n",
        "        根據獲得的img/ mask路徑讀取檔案\n",
        "        在讀取之前先確定讀到的是.dcm，否則raise Error例外\n",
        "        '''\n",
        "        if img_path.__contains__('.dcm'):  \n",
        "          # pydcm read image\n",
        "            ds = dcmread(img_path)\n",
        "            image = ds.pixel_array\n",
        "            image = image.astype('uint8') # 調整格式以配合albumentation套件需求\n",
        "        else:\n",
        "            raise ValueError(f'img path: {img_path} unknown')\n",
        "\n",
        "        '''\n",
        "        在這邊讀取mask，不論是CT或是MRI讀取到的都是png，\n",
        "        要注意的是cv2預設讀近來是float16，\n",
        "        '''\n",
        "        # cv2 read mask(Ground)\n",
        "        mask = cv2.imread(mask_path)[...,0]\n",
        "        mask = mask.astype('float32') # 調整格式以配合albumentation套件需求\n",
        "\n",
        "        '''\n",
        "        如果讀到MRI，我們只考慮所有value=63的像素，他代表肝臟的標示\n",
        "        其他像素點設置成0\n",
        "        '''\n",
        "        # check mask type\n",
        "        if self.dtype == 'MRI':\n",
        "            tmp = mask.copy()\n",
        "            tmp[tmp!=63] = 0\n",
        "            tmp[tmp==63] = 1\n",
        "            mask = tmp\n",
        "        elif self.dtype == 'CT':\n",
        "            mask /= 255\n",
        "        else:\n",
        "            raise ValueError('Non-support dtype')\n",
        "\n",
        "\n",
        "        if self.transform:\n",
        "            transformed = self.transform(image=image, mask=mask)\n",
        "            image = transformed['image']\n",
        "            mask = transformed['mask']\n",
        "\n",
        "\n",
        "        image = np.expand_dims(image, axis=0)\n",
        "        image = np.concatenate((image, image, image), axis=0)\n",
        "\n",
        "        image = torch.Tensor(image)\n",
        "\n",
        "        mask = torch.Tensor(mask) \n",
        "        mask = mask.unsqueeze(0)\n",
        "\n",
        "        return image, mask\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bfe3e403",
      "metadata": {
        "id": "bfe3e403"
      },
      "outputs": [],
      "source": [
        "#  https://pytorch.org/tutorials/beginner/basics/data_tutorial.html\n",
        "\n",
        "class ImageOnly_Dataset(Dataset):\n",
        "    def __init__(self, imgs_path_list, dtype, transform=None):\n",
        "        self.imgs_path_list = imgs_path_list\n",
        "        if type(imgs_path_list) != list:\n",
        "            raise ValueError('Need Input a list')\n",
        "        self.transform = transform\n",
        "        self.dtype = dtype\n",
        "#   \n",
        "    def __len__(self):\n",
        "        return len(self.imgs_path_list)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # now = time.time()\n",
        "        img_path = self.imgs_path_list[idx]\n",
        "        # print(img_anno_path)\n",
        "        # file_name = img_path.split('/')[-1].split('.')[0]\n",
        "#         img_path = img_anno_path[0]\n",
        "#         mask_path = img_anno_path[1]\n",
        "        # mask_path = [i for i in self.anno_path_list if i.__contains__(file_name)][0]\n",
        "\n",
        "        '''\n",
        "        根據獲得的img/ mask路徑讀取檔案\n",
        "        在讀取之前先確定讀到的是.dcm，否則raise Error例外\n",
        "        '''\n",
        "        if img_path.__contains__('.dcm'):  \n",
        "          # pydcm read image\n",
        "            ds = dcmread(img_path)\n",
        "            image = ds.pixel_array\n",
        "            image = image.astype('uint8') # 調整格式以配合albumentation套件需求\n",
        "        else:\n",
        "            raise ValueError(f'img path: {img_path} unknown')\n",
        "\n",
        "\n",
        "        if self.transform:\n",
        "            transformed = self.transform(image=image)\n",
        "            image = transformed['image']\n",
        "\n",
        "\n",
        "        image = np.expand_dims(image, axis=0)\n",
        "        image = np.concatenate((image, image, image), axis=0)\n",
        "\n",
        "        image = torch.Tensor(image)\n",
        "\n",
        "        return image"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5ed780c5",
      "metadata": {
        "id": "5ed780c5"
      },
      "source": [
        "#### 補充: 如何從dataset抽image, mask出來\n",
        "```\n",
        "MRI_train = dataset['TTsplit']['MRI_T2SPIR']['train']\n",
        "MRI_test = dataset['TTsplit']['MRI_T2SPIR']['test']\n",
        "\n",
        "CT_train = dataset['TTsplit']['CT']['train']\n",
        "CT_test = dataset['TTsplit']['CT']['test']\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ecb1387c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ecb1387c",
        "outputId": "2641296d-e9ad-4f9f-bbcc-faee19c524bb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(320, 320) (320, 320)\n",
            "(320, 320) (320, 320)\n",
            "(320, 320) (320, 320)\n"
          ]
        }
      ],
      "source": [
        "a = CTMRI_ImageDataset(MRI_train, dtype='MRI')\n",
        "\n",
        "'''\n",
        "在這邊你可以仔細檢查，然後會發現mask的shape其實不是固定的，等一下再transform的時候要強制對所有mask resize\n",
        "'''\n",
        "for idx in range(a.__len__()):\n",
        "    test_image, test_mask = a.__getitem__(idx)\n",
        "    test_image = test_image[0].numpy()\n",
        "    test_mask = test_mask.squeeze(0).numpy()\n",
        "    print(test_image.shape, test_mask.shape)\n",
        "    if idx==2: break"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6a65dad",
      "metadata": {
        "id": "f6a65dad"
      },
      "source": [
        "#### 使用albumentations進行資料擴增"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e2a7d87",
      "metadata": {
        "id": "8e2a7d87"
      },
      "outputs": [],
      "source": [
        "# https://albumentations.ai/docs/getting_started/mask_augmentation/\n",
        "\n",
        "transform = A.Compose([\n",
        "    A.HorizontalFlip(p=0.5),\n",
        "    A.RandomBrightnessContrast(brightness_limit=[-0.05, 0.05], p=0.2),\n",
        "    A.Rotate((-30, 30), interpolation=0), \n",
        "    A.RandomContrast(limit=0.2, p=1), \n",
        "\n",
        "    A.Normalize(p=1, mean=(0.485), std=(0.229)),\n",
        "    A.ToFloat(always_apply=True),\n",
        "    A.Resize(WIDTH, HEIGHT),\n",
        "])\n",
        "\n",
        "target_transform = A.Compose([\n",
        "    A.Normalize(p=1, mean=(0.485), std=(0.229)),                         \n",
        "    A.ToFloat(always_apply=True),\n",
        "    A.Resize(WIDTH, HEIGHT),\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "263dd3dc",
      "metadata": {
        "id": "263dd3dc"
      },
      "source": [
        "### 建立DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb480c0e",
      "metadata": {
        "id": "eb480c0e"
      },
      "outputs": [],
      "source": [
        "# 建議同時間只有8個(256,256)的sample進行計算 (Total = BATCH_SIZE*MULTIPLE_BATCH)\n",
        "\n",
        "dataset_train = CTMRI_ImageDataset(MRI_train, dtype='MRI', transform=transform)\n",
        "dataloader_train = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)\n",
        "\n",
        "dataset_test = CTMRI_ImageDataset(MRI_test, dtype='MRI', transform=target_transform) # **如果要正式使用要記得把這裡換成X_test\n",
        "dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)\n",
        "\n",
        "CT_dataset_train = CTMRI_ImageDataset(CT_train, dtype='CT', transform=transform)\n",
        "CT_dataloader_train = torch.utils.data.DataLoader(CT_dataset_train, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)\n",
        "\n",
        "CT_dataset_test = CTMRI_ImageDataset(CT_test, dtype='CT', transform=target_transform)\n",
        "CT_dataloader_test = torch.utils.data.DataLoader(CT_dataset_test, batch_size=BATCH_SIZE, shuffle=False)\n",
        "\n",
        "\n",
        "MRI_image_dataset_train = ImageOnly_Dataset(dataset['MRI_T2SPIR_nonLabel_train'], dtype='MRI', transform=transform)\n",
        "MRI_image_dataloader = torch.utils.data.DataLoader(MRI_image_dataset_train, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0daeac5a",
      "metadata": {
        "id": "0daeac5a",
        "outputId": "4ffa6dfc-5248-46e0-9feb-ba9e1da0f200"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Size([8, 3, 256, 256])"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a = iter(MRI_image_dataloader)\n",
        "x = a.next()\n",
        "x.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c6f7fe00",
      "metadata": {
        "id": "c6f7fe00"
      },
      "outputs": [],
      "source": [
        "# def Dict_ToggleButtons(dic, description=''):\n",
        "#     w = widgets.ToggleButtons(\n",
        "#         options=dic.keys(),\n",
        "#         description=description,\n",
        "#         disabled=False,\n",
        "#         button_style='', # 'success', 'info', 'warning', 'danger' or ''\n",
        "#     #     tooltips=['Description of slow', 'Description of regular', 'Description of fast'],\n",
        "#     #     icons=['check'] * 3\n",
        "#     )\n",
        "    \n",
        "#     return w"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e4215208",
      "metadata": {
        "id": "e4215208"
      },
      "outputs": [],
      "source": [
        "# #@title [Colab Form] 測試不同的dataloader \n",
        "# #@markdown 可能的選擇包括 'dataloader_train', 'dataloader_test', 'CT_dataloader_test'\n",
        "\n",
        "# def foo(): \n",
        "    \n",
        "#     dic = dict(zip(\n",
        "#     ['dataloader_train', 'dataloader_test', 'CT_dataloader_train', 'CT_dataloader_test'],\n",
        "#     [dataloader_train, dataloader_test, CT_dataloader_train, CT_dataloader_test]\n",
        "#     ))\n",
        "    \n",
        "#     w = Dict_ToggleButtons(dic)\n",
        "#     # display(w)\n",
        "#     dataloader = dic[w.value]\n",
        "\n",
        "#     dataloader = iter(dataloader) \n",
        "#     image, mask = dataloader.next()\n",
        "# #     print(len(mask))\n",
        "# #     print(len(image))\n",
        "#     for idx, data in enumerate(zip(image, mask), 1):\n",
        "#         x,y = data\n",
        "#         print(x.shape, y.shape)\n",
        "#         show_image(x[0].numpy(), y.squeeze(0).numpy())\n",
        "\n",
        "#         showPics = 3 #@param {type:\"integer\"}\n",
        "#         if idx==showPics: \n",
        "#             break\n",
        "# foo()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "55d8609d",
      "metadata": {
        "id": "55d8609d"
      },
      "source": [
        "## 進行最佳化器與模型設計\n",
        "- 最佳化器(預設使用BECLossWithLogit)\n",
        "- 模型(並將模型放到torch指定的運算空間: [cpu, cuda])\n",
        "\n",
        ">二元分類，只有一個channel的輸出的場合。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f1b6c176",
      "metadata": {
        "id": "f1b6c176"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        " \n",
        "class DiceLoss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(DiceLoss, self).__init__()\n",
        " \n",
        "    def\tforward(self, input, target):\n",
        "        N = target.size(0)\n",
        "        smooth = 1\n",
        " \n",
        "        input_flat = input.view(N, -1)\n",
        "        target_flat = target.view(N, -1)\n",
        " \n",
        "        intersection = input_flat * target_flat\n",
        " \n",
        "        loss = (2 * intersection.sum(1) + smooth) / (input_flat.sum(1) + target_flat.sum(1) + smooth)\n",
        "        loss = 1 - loss.sum() / N\n",
        " \n",
        "        return loss\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2717ceea",
      "metadata": {
        "id": "2717ceea"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "\n",
        "class FocalLoss(nn.Module):\n",
        "    r\"\"\"\n",
        "        This criterion is a implemenation of Focal Loss, which is proposed in \n",
        "        Focal Loss for Dense Object Detection.\n",
        "\n",
        "            Loss(x, class) = - \\alpha (1-softmax(x)[class])^gamma \\log(softmax(x)[class])\n",
        "\n",
        "        The losses are averaged across observations for each minibatch.\n",
        "\n",
        "        Args:\n",
        "            alpha(1D Tensor, Variable) : the scalar factor for this criterion\n",
        "            gamma(float, double) : gamma > 0; reduces the relative loss for well-classiﬁed examples (p > .5), \n",
        "                                   putting more focus on hard, misclassiﬁed examples\n",
        "            size_average(bool): By default, the losses are averaged over observations for each minibatch.\n",
        "                                However, if the field size_average is set to False, the losses are\n",
        "                                instead summed for each minibatch.\n",
        "\n",
        "\n",
        "    \"\"\"\n",
        "    def __init__(self, class_num, alpha=None, gamma=2, size_average=True):\n",
        "        super(FocalLoss, self).__init__()\n",
        "        if alpha is None:\n",
        "            self.alpha = Variable(torch.ones(class_num, 1))\n",
        "        else:\n",
        "            if isinstance(alpha, Variable):\n",
        "                self.alpha = alpha\n",
        "            else:\n",
        "                self.alpha = Variable(alpha)\n",
        "        self.gamma = gamma\n",
        "        self.class_num = class_num\n",
        "        self.size_average = size_average\n",
        "\n",
        "    def forward(self, inputs, targets):\n",
        "        N = inputs.size(0)\n",
        "        C = inputs.size(1)\n",
        "        P = F.softmax(inputs)\n",
        "\n",
        "        class_mask = inputs.data.new(N, C).fill_(0)\n",
        "        class_mask = Variable(class_mask)\n",
        "        ids = targets.view(-1, 1)\n",
        "        class_mask.scatter_(1, ids.data, 1.)\n",
        "        #print(class_mask)\n",
        "\n",
        "\n",
        "        if inputs.is_cuda and not self.alpha.is_cuda:\n",
        "            self.alpha = self.alpha.cuda()\n",
        "        alpha = self.alpha[ids.data.view(-1)]\n",
        "\n",
        "        probs = (P*class_mask).sum(1).view(-1,1)\n",
        "\n",
        "        log_p = probs.log()\n",
        "        #print('probs size= {}'.format(probs.size()))\n",
        "        #print(probs)\n",
        "\n",
        "        batch_loss = -alpha*(torch.pow((1-probs), self.gamma))*log_p \n",
        "        #print('-----bacth_loss------')\n",
        "        #print(batch_loss)\n",
        "\n",
        "\n",
        "        if self.size_average:\n",
        "            loss = batch_loss.mean()\n",
        "        else:\n",
        "            loss = batch_loss.sum()\n",
        "        return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "29fc0111",
      "metadata": {
        "id": "29fc0111"
      },
      "outputs": [],
      "source": [
        "class BCEWithRCFLoss(nn.Module):\n",
        "    '''\n",
        "    red: https://github.com/meteorshowers/RCF-pytorch/blob/master/functions.py\n",
        "    '''\n",
        "    def __init__(self):\n",
        "        super(BCEWithRCFLoss, self).__init__()\n",
        "  \n",
        "    def\tforward(self, input, target):\n",
        "        target = target.long()\n",
        "        num_positive = torch.sum((mask==1)).float()\n",
        "        num_negative = torch.sum((mask==0)).float()\n",
        "\n",
        "        mask[mask == 1] = 1.0 * num_negative / (num_positive + num_negative)\n",
        "        mask[mask == 0] = 1.1 * num_positive / (num_positive + num_negative)\n",
        "        cost = torch.nn.functional.binary_cross_entropy_with_logits(\n",
        "                input.float(),target.float(), weight=mask, reduce=False)\n",
        "        return torch.sum(cost)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "697b5eb4",
      "metadata": {
        "id": "697b5eb4"
      },
      "source": [
        "## 訓練模型 DANN+Unet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b5c711dd",
      "metadata": {
        "id": "b5c711dd"
      },
      "outputs": [],
      "source": [
        "# from unet import convBlock, Encoder, UpSampleConvs, Decoder\n",
        "# from unet import Encoder\n",
        "# from unet import Decoder\n",
        "from torch.nn import functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3437f8dd",
      "metadata": {
        "id": "3437f8dd"
      },
      "outputs": [],
      "source": [
        "## 加入instance normalization\n",
        "class convBlock(nn.Module):\n",
        "    def __init__(self, in_ch, out_ch, padding = 'same', kernel_size=3):\n",
        "        super().__init__()\n",
        "        kernel_size = kernel_size\n",
        "        pad_size = lambda kernel_size:(kernel_size-1)//2\n",
        "        if padding=='same':\n",
        "            self.padding = pad_size(kernel_size)\n",
        "        else:\n",
        "            self.padding = padding\n",
        "        \n",
        "        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size, padding=self.padding, bias=False)\n",
        "        self.relu  = nn.ReLU()\n",
        "        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size, padding=self.padding, bias=False)\n",
        "        self.INorm = torch.nn.InstanceNorm2d(out_ch, affine=True)\n",
        "        self.BN = nn.BatchNorm2d(out_ch)\n",
        "        \n",
        "    def forward(self, x):\n",
        "#         x = self.INorm(self.conv1(x))\n",
        "        x = self.BN(self.conv1(x))\n",
        "        x = self.relu(x)\n",
        "#         x = self.INorm(self.conv2(x))\n",
        "        x = self.BN(self.conv2(x))\n",
        "        x = self.relu(x)\n",
        "        return x\n",
        "    \n",
        "class Encoder(nn.Module):\n",
        "    def __init__(self, chs=(3,32,64,128,256,512), padding='same'):\n",
        "        super().__init__()\n",
        "        self.FPN_enc_ftrs = nn.ModuleList([convBlock(chs[i], chs[i+1], padding) for i in range(len(chs)-1)])\n",
        "        self.pool = torch.max_pool2d\n",
        "        \n",
        "    def forward(self, x):\n",
        "        features = []\n",
        "        \n",
        "        for block in self.FPN_enc_ftrs:\n",
        "            x = block(x)\n",
        "            features.append(x)\n",
        "            x = self.pool(x, kernel_size=2)\n",
        "        return features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd52e714",
      "metadata": {
        "id": "dd52e714"
      },
      "outputs": [],
      "source": [
        "import torchvision\n",
        "class UpSampleConvs(nn.Module):\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super().__init__()\n",
        "        self.conv = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False)\n",
        "        self.relu  = nn.ReLU()\n",
        "        self.upSample = nn.Upsample(scale_factor=2)\n",
        "        self.INorm = torch.nn.InstanceNorm2d(out_ch)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.upSample(x)\n",
        "        x = self.conv(x)\n",
        "        x = self.relu(x)\n",
        "        x = self.INorm(x)\n",
        "        return x\n",
        "\n",
        "class Decoder(nn.Module):\n",
        "    def __init__(self, chs=(512, 256, 128, 64, 32), padding='same', multi_level = 0):\n",
        "        super().__init__()\n",
        "\n",
        "        self.chs = chs\n",
        "        self.padding = padding\n",
        "        self.upconvs = nn.ModuleList([UpSampleConvs(chs[i], chs[i+1]) for i in range(len(chs)-1)]) # 上採樣後卷積\n",
        "        self.FPN_dec_ftrs = nn.ModuleList([convBlock(chs[i], chs[i+1], padding=padding) for i in range(len(chs)-1)]) \n",
        "        self.multi_level = multi_level # int like\n",
        "  \n",
        "        \n",
        "    def forward(self, x, encoder_features):\n",
        "#         print(len(encoder_features))\n",
        "        if self.multi_level == 0:\n",
        "            multi_level_Ftrs = []\n",
        "        elif self.multi_level == -1:\n",
        "            multi_level_Ftrs = [0]*(len(encoder_features)+1)\n",
        "            multi_level_Ftrs[0] = x\n",
        "        else:\n",
        "            multi_level_Ftrs = [0]*(self.multi_level+1)\n",
        "            multi_level_Ftrs[0] = x\n",
        "            \n",
        "        for i in range(len(self.chs)-1):\n",
        "            enc_ftrs = encoder_features[i]\n",
        "\n",
        "            x = self.upconvs[i](x)\n",
        "            enc_ftrs = self.crop(encoder_features[i], x)\n",
        "            x = torch.cat([x, enc_ftrs], dim=1)\n",
        "            x = self.FPN_dec_ftrs[i](x)\n",
        "            if self.multi_level == 0:\n",
        "                pass\n",
        "            elif i<len(multi_level_Ftrs)-1:\n",
        "                multi_level_Ftrs[i+1] = x\n",
        "\n",
        "        return x, multi_level_Ftrs\n",
        "    \n",
        "    def crop(self, enc_ftrs, x):\n",
        "\n",
        "        _, _, H, W = x.shape\n",
        "        enc_ftrs   = torchvision.transforms.CenterCrop([H, W])(enc_ftrs)\n",
        "        return enc_ftrs"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ea6b3c4f",
      "metadata": {
        "id": "ea6b3c4f"
      },
      "source": [
        "### 調整unet ch數"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0219f4d1",
      "metadata": {
        "id": "0219f4d1"
      },
      "outputs": [],
      "source": [
        "SCALE=4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "39505765",
      "metadata": {
        "id": "39505765"
      },
      "outputs": [],
      "source": [
        "class FeatureExtractor(nn.Module):\n",
        "    def __init__(self, enc_chs=(3,64,128,256,512,1024), padding='same'):\n",
        "        super().__init__()\n",
        "        tmp = enc_chs[0]\n",
        "        enc_chs = [i//SCALE for i in enc_chs]\n",
        "        enc_chs[0]=tmp\n",
        "        enc_chs = tuple(enc_chs)\n",
        "        self.BN0       = nn.BatchNorm2d(enc_chs[0])\n",
        "        self.encoder     = Encoder(enc_chs, padding=padding)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.BN0(x)\n",
        "        enc_ftrs = self.encoder(x)\n",
        "        return enc_ftrs\n",
        "\n",
        "class LabelPredictor(nn.Module):\n",
        "    def __init__(self, out_sz, num_class=1, retain_dim=False, padding='same',\n",
        "                         activation=None, dec_chs=(1024, 512, 256, 128, 64), multi_level=0):\n",
        "        super().__init__()\n",
        "        dec_chs = tuple(i//SCALE for i in dec_chs)\n",
        "        self.head  = nn.Conv2d(dec_chs[-1], num_class, 1)\n",
        "        self.retain_dim  = retain_dim\n",
        "        self.out_sz = out_sz\n",
        "        self.activation = activation\n",
        "        self.decoder     = Decoder(dec_chs, multi_level=multi_level, padding=padding)\n",
        "        self.multi_level = multi_level\n",
        "        \n",
        "    def forward(self, enc_ftrs):\n",
        "        \n",
        "        out, multi_level_Ftrs  = self.decoder(enc_ftrs[::-1][0], enc_ftrs[::-1][1:]) # 把不同尺度的所有featuremap都輸入decoder，我們在decoder需要做featuremap的拼接\n",
        "        out = self.head(out)\n",
        "        if self.activation:\n",
        "            out = self.activation(out)\n",
        "        if self.retain_dim:\n",
        "            out = F.interpolate(out, self.out_sz)\n",
        "            \n",
        "        if self.multi_level == 0:\n",
        "            return out, []\n",
        "        else:  \n",
        "            return out, multi_level_Ftrs\n",
        "\n",
        "\n",
        "class conv(nn.Module):\n",
        "    def __init__(self,ch_in,ch_out):\n",
        "        super().__init__()\n",
        "        self.cell=nn.Sequential(\n",
        "            nn.Conv2d(ch_in, ch_out, 3, 1, 1),\n",
        "            nn.BatchNorm2d(ch_out),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(2),\n",
        "            nn.Dropout(0.2)\n",
        "        )\n",
        "    def forward(self,x):\n",
        "        return self.cell(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2c4fc772",
      "metadata": {
        "id": "2c4fc772"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "兩個domain的image要送進去，\n",
        "每個domain的shape為: B, in_channels, H, W\n",
        "output的shape是: B, 1, H, W\n",
        "'''\n",
        "class Discriminator(nn.Module):\n",
        "    def __init__(self, in_channels=3, level=4, flatten=False):\n",
        "        super(Discriminator, self).__init__()\n",
        "\n",
        "        def discriminator_block(in_filters, out_filters, normalization=True):\n",
        "            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n",
        "            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n",
        "            if normalization:\n",
        "                layers.append(nn.InstanceNorm2d(out_filters))\n",
        "            layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
        "            return layers\n",
        "        \n",
        "        chs = [16*2**i for i in range(level)]\n",
        "        layers = [*discriminator_block(in_channels, 16, normalization=False)]\n",
        "        for i in range(level-1):\n",
        "            layers += discriminator_block(chs[i], chs[i+1])\n",
        "        \n",
        "        if flatten:\n",
        "            self.model = nn.Sequential(*layers,\n",
        "                                       nn.AdaptiveAvgPool2d((1,1)),\n",
        "                                        nn.Flatten(),\n",
        "                                       nn.Linear(16*2**(level-1), 1),\n",
        "#                                        nn.Linear(16*2**(level-1), 128),\n",
        "#                                        nn.Linear(128, 1), \n",
        "#                                        nn.Sigmoid(),\n",
        "                                      )\n",
        "        else:\n",
        "            self.model = nn.Sequential(*layers)\n",
        "        \n",
        "\n",
        "    def forward(self, img_A, img_B=False):\n",
        "        # Concatenate image and condition image by channels to produce input\n",
        "#         img_input = torch.cat((img_A, img_B), 1d\n",
        "        if isinstance(img_B, torch.Tensor):\n",
        "            img_input = torch.cat((img_A, img_B), dim=1)\n",
        "        else:\n",
        "            img_input = img_A\n",
        "        return self.model(img_input)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "43e5e130",
      "metadata": {
        "id": "43e5e130"
      },
      "outputs": [],
      "source": [
        "class Unet(nn.Module):\n",
        "    def __init__(self, out_sz, out_channels=3, activation=None, multi_level=0):\n",
        "        super().__init__()\n",
        "        self.FeatureExtractor = FeatureExtractor()\n",
        "        LP = LabelPredictor(out_sz=out_sz, activation=activation, multi_level=0)\n",
        "        LP.head = nn.Conv2d(64//SCALE, out_channels, 1)\n",
        "        self.LabelPredictor = LP\n",
        "    def forward(self, x):\n",
        "        x = self.FeatureExtractor(x)\n",
        "        x, _ = self.LabelPredictor(x)\n",
        "        return x, _"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b296e134",
      "metadata": {
        "id": "b296e134"
      },
      "outputs": [],
      "source": [
        "class Generator(nn.Module):\n",
        "    def __init__(self, out_sz, out_channels=3, activation=None, multi_level=0):\n",
        "        super().__init__()\n",
        "        self.FeatureExtractor = FeatureExtractor(enc_chs=(3*2,64,128,256))\n",
        "        LP = LabelPredictor(out_sz=out_sz, dec_chs=(256, 128, 64),\n",
        "                            activation=activation, multi_level=multi_level)\n",
        "        LP.head = nn.Conv2d(64//SCALE, out_channels, 1)\n",
        "        self.LabelPredictor = LP\n",
        "        \n",
        "    def forward(self, x, domain_label):\n",
        "        x = torch.cat([x, domain_label], dim=1)\n",
        "\n",
        "        x = self.FeatureExtractor(x)\n",
        "        x, _ = self.LabelPredictor(x)\n",
        "        \n",
        "        return x, _\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "974c330f",
      "metadata": {
        "id": "974c330f"
      },
      "outputs": [],
      "source": [
        "# class Generator(nn.Module):\n",
        "#     '''\n",
        "#     netG_CT2MRI = Generator(out_sz=(HEIGHT, WIDTH), out_channels=3,\n",
        "#                              activation=nn.Sigmoid(), multi_level=2,\n",
        "#                             load_path=f'{save_root}share_model_MRI.bin').to(device)\n",
        "\n",
        "#     netG_MRI2CT = Generator(out_sz=(HEIGHT, WIDTH), out_channels=3,\n",
        "#                              activation=nn.Sigmoid(), multi_level=2,\n",
        "#                             load_path=f'{save_root}share_model_MRI.bin').to(device)\n",
        "#     '''\n",
        "#     def __init__(self, out_sz, out_channels=3, activation=None, load_path=None):\n",
        "#         super().__init__()\n",
        "#         self.unet =  Unet(out_sz=(HEIGHT, WIDTH), out_channels=1, activation=None, multi_level=0).to(device)\n",
        "#         if load_path:\n",
        "#             print(f'load model from: {load_path}')\n",
        "#             self.unet.load_state_dict(torch.load(load_path, map_location=device)) \n",
        "#         self.head = nn.Conv2d(1,3,1)\n",
        "#         self.activation = activation\n",
        "#     def forward(self, x):\n",
        "#         x, _ = self.unet(x)\n",
        "#         x = self.head(x)\n",
        "#         x = self.activation(x)\n",
        "#         return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee4fbafe",
      "metadata": {
        "id": "ee4fbafe"
      },
      "outputs": [],
      "source": [
        "# class focal_loss(nn.Module):    \n",
        "#     def __init__(self, alpha=0.25, gamma=2, num_classes = 3, size_average=True):\n",
        "#         \"\"\"\n",
        "#         focal_loss损失函数, -α(1-yi)**γ *ce_loss(xi,yi)      \n",
        "#         步骤详细的实现了 focal_loss损失函数.\n",
        "#         :param alpha:   阿尔法α,类别权重.      当α是列表时,为各类别权重,当α为常数时,类别权重为[α, 1-α, 1-α, ....],常用于 目标检测算法中抑制背景类 , retainnet中设置为0.25\n",
        "#         :param gamma:   伽马γ,难易样本调节参数. retainnet中设置为2\n",
        "#         :param num_classes:     类别数量\n",
        "#         :param size_average:    损失计算方式,默认取均值\n",
        "#         \"\"\"\n",
        "\n",
        "#         super(focal_loss,self).__init__()\n",
        "#         self.size_average = size_average\n",
        "#         if isinstance(alpha,list):\n",
        "#             assert len(alpha)==num_classes   # α可以以list方式输入,size:[num_classes] 用于对不同类别精细地赋予权重\n",
        "#             print(\"Focal_loss alpha = {}, 将对每一类权重进行精细化赋值\".format(alpha))\n",
        "#             self.alpha = torch.Tensor(alpha)\n",
        "#         else:\n",
        "#             assert alpha<1   #如果α为一个常数,则降低第一类的影响,在目标检测中为第一类\n",
        "#             print(\" --- Focal_loss alpha = {} ,将对背景类进行衰减,请在目标检测任务中使用 --- \".format(alpha))\n",
        "#             self.alpha = torch.zeros(num_classes)\n",
        "#             self.alpha[0] += alpha\n",
        "#             self.alpha[1:] += (1-alpha) # α 最终为 [ α, 1-α, 1-α, 1-α, 1-α, ...] size:[num_classes]\n",
        "#         self.gamma = gamma\n",
        "\n",
        "#     def forward(self, preds, labels):\n",
        "#         \"\"\"\n",
        "#         focal_loss损失计算        \n",
        "#         :param preds:   预测类别. size:[B,N,C] or [B,C]    分别对应与检测与分类任务, B 批次, N检测框数, C类别数        \n",
        "#         :param labels:  实际类别. size:[B,N] or [B]        \n",
        "#         :return:\n",
        "#         \"\"\"        \n",
        "#         # assert preds.dim()==2 and labels.dim()==1        \n",
        "#         preds = preds.view(-1,preds.size(-1))        \n",
        "#         self.alpha = self.alpha.to(preds.device)        \n",
        "#         preds_softmax = F.softmax(preds, dim=1) # 这里并没有直接使用log_softmax, 因为后面会用到softmax的结果(当然你也可以使用log_softmax,然后进行exp操作)        \n",
        "#         preds_logsoft = torch.log(preds_softmax)\n",
        "#         preds_softmax = preds_softmax.gather(1,labels.view(-1,1))   # 这部分实现nll_loss ( crossempty = log_softmax + nll )        \n",
        "#         preds_logsoft = preds_logsoft.gather(1,labels.view(-1,1))        \n",
        "#         self.alpha = self.alpha.gather(0,labels.view(-1))        \n",
        "#         loss = -torch.mul(torch.pow((1-preds_softmax), self.gamma), preds_logsoft)  # torch.pow((1-preds_softmax), self.gamma) 为focal loss中 (1-pt)**γ\n",
        "#         loss = torch.mul(self.alpha, loss.t())        \n",
        "#         if self.size_average:        \n",
        "#             loss = loss.mean()        \n",
        "#         else:            \n",
        "#             loss = loss.sum()        \n",
        "#         return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cf75debc",
      "metadata": {
        "id": "cf75debc"
      },
      "outputs": [],
      "source": [
        "# #PyTorch\n",
        "# ALPHA = 0.8\n",
        "# GAMMA = 2\n",
        "\n",
        "# class FocalLoss(nn.Module):\n",
        "#     def __init__(self, weight=None, size_average=True):\n",
        "#         super(FocalLoss, self).__init__()\n",
        "\n",
        "#     def forward(self, inputs, targets, alpha=ALPHA, gamma=GAMMA, smooth=1):\n",
        "        \n",
        "#         #comment out if your model contains a sigmoid or equivalent activation layer\n",
        "#         inputs = F.sigmoid(inputs)       \n",
        "        \n",
        "#         #flatten label and prediction tensors\n",
        "#         inputs = inputs.view(-1)\n",
        "#         targets = targets.view(-1)\n",
        "        \n",
        "#         #first compute binary cross-entropy \n",
        "#         BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')\n",
        "#         BCE_EXP = torch.exp(-BCE)\n",
        "#         focal_loss = alpha * (1-BCE_EXP)**gamma * BCE\n",
        "                       \n",
        "#         return focal_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7b0a3246",
      "metadata": {
        "id": "7b0a3246"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "https://github.com/clcarwin/focal_loss_pytorch/blob/e11e75bad957aecf641db6998a1016204722c1bb/focalloss.py\n",
        "'''\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "\n",
        "class FocalLoss(nn.Module):\n",
        "    def __init__(self, gamma=0, alpha=None, size_average=True):\n",
        "        super(FocalLoss, self).__init__()\n",
        "        self.gamma = gamma\n",
        "        self.alpha = alpha\n",
        "        if isinstance(alpha,(float,int)): self.alpha = torch.Tensor([alpha,1-alpha])\n",
        "        if isinstance(alpha,list): self.alpha = torch.Tensor(alpha)\n",
        "        self.size_average = size_average\n",
        "\n",
        "    def forward(self, inputs, target):\n",
        "        if inputs.dim()>2:\n",
        "#             input = input.view(input.size(0),input.size(1),-1)  # N,C,H,W => N,C,H*W\n",
        "#             input = input.transpose(1,2)    # N,C,H*W => N,H*W,C\n",
        "#             input = input.contiguous().view(-1,input.size(2))   # N,H*W,C => N*H*W,C\n",
        "            C = inputs.shape[1] # num class\n",
        "            inputs = inputs.transpose(1,-1)\n",
        "            inputs = inputs.reshape(-1, C)\n",
        "        \n",
        "        target = target.to(torch.int64)\n",
        "        # flatten all pixel\n",
        "        target = target.view(-1,1) # input should be a tensor (N, 1, H, W), 1 for 1 ch, class shound be [1,C], dtype=Long\n",
        "        \n",
        "        \n",
        "        logpt = F.log_softmax(inputs) # log(softmax(x))\n",
        "        logpt = logpt.gather(1, target) # explain by list, logpt = [logpt[i, target[i]] for i in range(len(target))]\n",
        "        logpt = logpt.view(-1)\n",
        "        pt = Variable(logpt.data.exp()) # reverse the log operation\n",
        "\n",
        "        if self.alpha is not None:\n",
        "            if self.alpha.type()!=inputs.data.type():\n",
        "                self.alpha = self.alpha.type_as(inputs.data) \n",
        "            at = self.alpha.gather(0,target.data.view(-1))\n",
        "            logpt = logpt * Variable(at)\n",
        "\n",
        "        loss = -1 * (1-pt)**self.gamma * logpt\n",
        "        if self.size_average: return loss.mean()\n",
        "        else: return loss.sum()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7c237dab",
      "metadata": {
        "id": "7c237dab",
        "outputId": "d50c94b3-29cd-4af1-9ad1-94b6a59bac3f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "細節手刻: tensor(0.1944)\n",
            "使用Focal Loss class forward: tensor(0.1944)\n"
          ]
        }
      ],
      "source": [
        "# gather example\n",
        "'para'\n",
        "size_average= True\n",
        "gamma = 0.1 \n",
        "alpha = torch.Tensor([0.1, 0.5, 0.3])\n",
        "num_class = 2\n",
        "\n",
        "'input'\n",
        "sample_x    = torch.zeros(2, num_class, WIDTH, HEIGHT)\n",
        "x = sample_x.clone()\n",
        "y =  torch.rand(2, 1, WIDTH, HEIGHT)*num_class\n",
        "target = y.clone()\n",
        "target = target.long()\n",
        "\n",
        "'operation'\n",
        "C = x.shape[1]\n",
        "x = x.transpose(1,-1)\n",
        "x = x.reshape(-1, C)\n",
        "logpt = F.log_softmax(x)\n",
        "# print(logpt)\n",
        "\n",
        "target = target.view(-1,1)\n",
        "\n",
        "logpt = logpt.gather(1, target)\n",
        "\n",
        "''' # gather opersation explain\n",
        "will be like this in python:\n",
        "logpt = [logpt[i][idx] for i, idx in enumerate(target)]\n",
        "\n",
        "# example\n",
        "src = list(logpt)\n",
        "index = list(target)\n",
        "src = [src[i][idx] for i, idx in enumerate(index)]\n",
        "print(src)\n",
        "'''\n",
        "\n",
        "logpt = logpt.view(-1) # gather後，攤平\n",
        "\n",
        "pt = Variable(logpt.data.exp())\n",
        "\n",
        "alpha = torch.Tensor([0.1, 0.5, 0.3]) # 對類別做weight\n",
        "alpha = alpha.type_as(x.data) \n",
        "at = alpha.gather(0,target.data.view(-1)) # 對所有樣本產生weight，weight基於target(mask annotation)的index\n",
        "# print(at) # 已經被攤平了\n",
        "# print(logpt)\n",
        "logpt = logpt * Variable(at)\n",
        "# print(logpt)\n",
        "\n",
        "\n",
        "loss = -1 * (1-pt)**gamma * logpt\n",
        "if size_average: \n",
        "    loss = loss.mean()\n",
        "else: \n",
        "    loss = loss.sum()\n",
        "print('細節手刻:',loss)\n",
        "\n",
        "\n",
        "\n",
        "x = sample_x.clone()\n",
        "target = y.long().clone()\n",
        "\n",
        "a = FocalLoss(gamma=gamma, alpha=alpha)\n",
        "print('使用Focal Loss class forward:' ,a(x, target))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb82c829",
      "metadata": {
        "id": "eb82c829"
      },
      "outputs": [],
      "source": [
        "# def foo():\n",
        "    \n",
        "#     x    = torch.randn(2, 3, WIDTH, HEIGHT)\n",
        "    \n",
        "# #     a = FeatureExtractor()\n",
        "# #     features = a(x)\n",
        "# #     print('FeatureExtractor\\n', [i.shape for i in features])\n",
        "    \n",
        "# #     a = LabelPredictor(out_sz=(WIDTH, HEIGHT), activation=nn.Sigmoid(), multi_level=0)\n",
        "# #     y_pred, multi_ftrs = a(features)\n",
        "# #     print('LabelPredictor >>>1\\n', y_pred.shape, [i.shape for i in multi_ftrs])\n",
        "    \n",
        "# #     a = LabelPredictor(out_sz=(WIDTH, HEIGHT), activation=nn.Sigmoid(), multi_level=4)\n",
        "# #     y_pred, multi_ftrs = a(features)\n",
        "# #     print('LabelPredictor >>>2\\n', y_pred.shape, [i.shape for i in multi_ftrs])\n",
        "    \n",
        "# #     a  = Discriminator(multi_ftrs[0].shape[1], level=2, flatten=True)\n",
        "# #     logit = a(multi_ftrs[0])\n",
        "# #     print('Discriminator for FeatureExtractor >>>1\\n', multi_ftrs[0].shape[1], logit.shape)\n",
        "    \n",
        "# #     a  = Discriminator(multi_ftrs[1].shape[1], level=4, flatten=True)\n",
        "# #     logit = a(multi_ftrs[1])\n",
        "# #     print('Discriminator for FeatureExtractor >>>2\\n', multi_ftrs[1].shape[1], logit.shape)\n",
        "    \n",
        "# #     a  = Discriminator(multi_ftrs[-1].shape[1], level=4, flatten=True)\n",
        "# #     logit = a(multi_ftrs[-1])\n",
        "# #     print('Discriminator for FeatureExtractor >>>3\\n', multi_ftrs[-1].shape[1], logit.shape)\n",
        "    \n",
        "# #     a = Unet(out_sz=(HEIGHT, WIDTH), out_channels=3, activation=nn.Sigmoid()) \n",
        "# #     pred, _ = a(x)\n",
        "# #     print('Unet\\n', pred.shape)\n",
        "    \n",
        "# #     a  = Discriminator(pred.shape[1], level=2, flatten=True) # input: (2,1,256,256)\n",
        "# #     logit = a(pred)\n",
        "# #      print(logit)\n",
        "# #     print('Discriminator for Unet\\n', logit.shape)\n",
        "    \n",
        "# #    a = Generator(out_sz=(HEIGHT, WIDTH), out_channels=3, activation=nn.Sigmoid(), multi_level=2) \n",
        "# ##    source_domain_tensor = torch.ones(source_data.data.size()).to(device)\n",
        "# #    target_domain_tensor = torch.zeros(source_data.data.size()).to(device)\n",
        "# #    pred, _ = a(x, source_domain_tensor)\n",
        "# #    print(pred.shape)\n",
        "# #    print([i.shape for i in _])\n",
        "# #     print(y_pred.shape)\n",
        "# #     print( F.softmax(y_pred, dim=0).shape)\n",
        "# #     F.softmax(pred_MRI)\n",
        "# #     torch.zeros(source_data.data.size()).to(device)\n",
        "#     class_criterion = FocalLoss(class_num=2)\n",
        "#     loss = class_criterion(x, torch.zeros(x.data.size()))\n",
        "#     print(loss)\n",
        "# foo()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7787e11c",
      "metadata": {
        "id": "7787e11c"
      },
      "outputs": [],
      "source": [
        "# mem_params = sum([param.nelement()*param.element_size() for param in a.parameters()])\n",
        "# mem_bufs = sum([buf.nelement()*buf.element_size() for buf in a.buffers()])\n",
        "# mem = mem_params + mem_bufs\n",
        "# mem/1024/1024 #bytes/ kb/ mb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1f3feeb4",
      "metadata": {
        "id": "1f3feeb4"
      },
      "outputs": [],
      "source": [
        "save_root = 'saved/0413-5/'\n",
        "# class_criterion = DiceLoss()\n",
        "class_criterion = FocalLoss()\n",
        "domain_criterion = nn.BCEWithLogitsLoss()\n",
        "consist_criterion = nn.L1Loss()\n",
        "out_channels = 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b2e282f0",
      "metadata": {
        "id": "b2e282f0"
      },
      "outputs": [],
      "source": [
        "model_D_image_source = Discriminator(in_channels=3, level=3, flatten=True).to(device)\n",
        "model_D_image_target = Discriminator(in_channels=3, level=3, flatten=True).to(device)\n",
        "model_MRI =  Unet(out_sz=(HEIGHT, WIDTH), out_channels=out_channels, activation=nn.Sigmoid(), multi_level=0).to(device)\n",
        "\n",
        "\n",
        "discOptim = optim.Adam([\n",
        "                {'params': model_D_image_source.parameters()},\n",
        "                {'params': model_D_image_target.parameters()},\n",
        "            ], lr=1e-4)  # default lr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1e3184d5",
      "metadata": {
        "id": "1e3184d5"
      },
      "outputs": [],
      "source": [
        "def train_disc(backward=True):\n",
        "    \n",
        "    for model_D in [model_D_image_source, model_D_image_target]:\n",
        "        for param in model_D.parameters():\n",
        "            param.requires_grad = True\n",
        "    \n",
        "    target_domain_label = 0\n",
        "    source_domain_label = 1\n",
        "    \n",
        "    '''\n",
        "    train domain classifier\n",
        "    ---\n",
        "        make domain predction\n",
        "    '''\n",
        "    MRI_pred = model_MRI(target_data)\n",
        "    \n",
        "    target_domain_label = 0\n",
        "    source_domain_label = 1\n",
        "\n",
        "    'prediction discrimator'\n",
        "    domain_logit = model_D_image_source(source_data)\n",
        "    disc_loss = domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device))\n",
        "    \n",
        "    'image discrimator'\n",
        "    domain_logit = model_D_image_target(target_data)\n",
        "    disc_loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "    if backward:\n",
        "        disc_loss.backward()\n",
        "        discOptim.step()\n",
        "        discOptim.zero_grad()\n",
        "    return disc_loss/2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0abd2ac6",
      "metadata": {
        "id": "0abd2ac6"
      },
      "outputs": [],
      "source": [
        "def train_share_weight_unet(backward=True):# model_MRI/ model_CT\n",
        "\n",
        "    for model in [model_MRI]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "            \n",
        "    tmpOptim = optim.Adam([{'params': model_MRI.parameters()},\n",
        "                          ], lr=1e-1)\n",
        "    \n",
        "    target_domain_label = 0\n",
        "    '''\n",
        "    compute model_MRI shared weight\n",
        "    '''\n",
        "    mixed_data = torch.cat([source_data, target_data], dim=0)\n",
        "    mixed_label = torch.cat([source_label, target_label], dim=0)\n",
        "    \n",
        "    y_pred_mixed, _ = model_MRI(mixed_data) \n",
        "    loss = class_criterion(y_pred_mixed, mixed_label)\n",
        "    \n",
        "    domain_logit = model_D_image_target(target_data)\n",
        "    disc_loss = domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "    \n",
        "    loss = loss - 0.1*disc_loss\n",
        "    \n",
        "    if backward:\n",
        "        loss.backward()\n",
        "        tmpOptim.step()\n",
        "        tmpOptim.zero_grad()\n",
        "        \n",
        "        del mixed_data, y_pred_mixed, mixed_label\n",
        "        torch.cuda.empty_cache()\n",
        "    \n",
        "    return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c8f0b5f6",
      "metadata": {
        "scrolled": true,
        "id": "c8f0b5f6",
        "outputId": "685f4ac4-75cb-4007-94e2-1a71ba32398d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "domain clf loss 0.6918910145759583\n"
          ]
        },
        {
          "ename": "KeyboardInterrupt",
          "evalue": "",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m/tmp/ipykernel_552091/3532760646.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     22\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'domain clf loss'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m         \u001b[0mshare_weight_value\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mtrain_share_weight_unet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0;32mdel\u001b[0m \u001b[0msource_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msource_label\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_label\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/tmp/ipykernel_552091/3862085133.py\u001b[0m in \u001b[0;36mtrain_share_weight_unet\u001b[0;34m(backward)\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mbackward\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0mtmpOptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m         \u001b[0mtmpOptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/opt/conda/lib/python3.7/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    253\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    254\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 255\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    256\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    257\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    147\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m    148\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 149\u001b[0;31m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    150\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    151\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "source_dataloader, target_dataloader = CT_dataloader_train, dataloader_train\n",
        "test_dataloader = dataloader_test\n",
        "\n",
        "EPOCHS = 100\n",
        "min_target_loss_value = 100\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "\n",
        "    share_weight_value = 0.0\n",
        "    testing_loss_value = 0.0\n",
        "    for i, ((source_data, source_label), (target_data, target_label)) in enumerate(zip(source_dataloader, target_dataloader)):\n",
        "        source_data = source_data.to(device)\n",
        "        source_label = source_label.to(device)\n",
        "        target_data = target_data.to(device)\n",
        "        target_label = target_label.to(device)\n",
        "    \n",
        "        \n",
        "        print(i, end='\\r')\n",
        "        if epoch==0 and i==0:\n",
        "            for loop in range(2):\n",
        "                loss = train_disc()\n",
        "            print('domain clf loss', loss.item())\n",
        "        \n",
        "        share_weight_value += train_share_weight_unet().item()\n",
        "        \n",
        "        del source_data, source_label, target_data, target_label\n",
        "        torch.cuda.empty_cache()\n",
        "        \n",
        "    for i2, (target_data, target_label) in enumerate(test_dataloader):\n",
        "        target_data = target_data.to(device)\n",
        "        target_label = target_label.to(device)\n",
        "        \n",
        "        y_pred, _ = model_MRI(target_data)\n",
        "        loss = class_criterion(y_pred, target_label)\n",
        "        testing_loss_value += loss.item()\n",
        "        \n",
        "        del y_pred\n",
        "        torch.cuda.empty_cache()\n",
        "        \n",
        "    print(f'epoch: {epoch}, testing_loss_value: {testing_loss_value/(i2+1)}')\n",
        "    testing_loss_value /= (i2+1)\n",
        "    if testing_loss_value < 0.5:\n",
        "        torch.save(model_MRI.state_dict(), f'{save_root}share_model_MRI.bin')\n",
        "        break\n",
        "    torch.save(model_MRI.state_dict(), f'{save_root}share_model_MRI.bin')\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4043d96d",
      "metadata": {
        "id": "4043d96d"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b2746dca",
      "metadata": {
        "id": "b2746dca"
      },
      "source": [
        "### GAN disc loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2736674b",
      "metadata": {
        "id": "2736674b"
      },
      "outputs": [],
      "source": [
        "model_MRI =  Unet(out_sz=(HEIGHT, WIDTH), out_channels=out_channels, activation=nn.Sigmoid(), multi_level=0).to(device)\n",
        "model_CT = Unet(out_sz=(HEIGHT, WIDTH), out_channels=out_channels, activation=nn.Sigmoid(), multi_level=0).to(device)\n",
        "\n",
        "model_MRI.load_state_dict(torch.load(f'{save_root}share_model_MRI.bin', map_location=device)) \n",
        "model_CT.load_state_dict(torch.load(f'{save_root}share_model_MRI.bin', map_location=device)) \n",
        "\n",
        "model_D_pred_source = Discriminator(in_channels=out_channels, level=3, flatten=True).to(device)\n",
        "model_D_pred_target = Discriminator(in_channels=out_channels, level=3, flatten=True).to(device)\n",
        "model_D_image_source = Discriminator(in_channels=3, level=3, flatten=True).to(device)\n",
        "model_D_image_target = Discriminator(in_channels=3, level=3, flatten=True).to(device)\n",
        "\n",
        "\n",
        "netG_CT2MRI = Generator(out_sz=(HEIGHT, WIDTH), out_channels=3, activation=None, multi_level=0).to(device) # stacked conv + linear  \n",
        "netG_MRI2CT = Generator(out_sz=(HEIGHT, WIDTH), out_channels=3, activation=None, multi_level=0).to(device)\n",
        "\n",
        "discOptim = optim.Adam([\n",
        "                {'params': model_D_pred_source.parameters()},\n",
        "                {'params': model_D_pred_target.parameters()},\n",
        "                {'params': model_D_image_source.parameters()},\n",
        "                {'params': model_D_image_target.parameters()},\n",
        "            ], lr=1e-4)  # default lr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c8a77d3c",
      "metadata": {
        "id": "c8a77d3c"
      },
      "outputs": [],
      "source": [
        "def train_disc(backward=True):\n",
        "    for model_D in [model_D_pred_source, model_D_pred_target, model_D_image_source, model_D_image_target]:\n",
        "        for param in model_D.parameters():\n",
        "            param.requires_grad = True\n",
        "    \n",
        "#     for model in [netG_CT2MRI, netG_MRI2CT]:\n",
        "#         for param in model.parameters():\n",
        "#             param.requires_grad = False\n",
        "    with torch.no_grad():\n",
        "        source_domain_tensor = torch.ones(source_data.data.size()).to(device)\n",
        "        target_domain_tensor = torch.zeros(source_data.data.size()).to(device)\n",
        "        \n",
        "        fake_MRI, _ = netG_CT2MRI(source_data, target_domain_tensor)\n",
        "        fake_CT, _ = netG_MRI2CT(target_data, source_domain_tensor)\n",
        "\n",
        "    with torch.no_grad():\n",
        "        CT_pred, _ = model_CT(source_data)\n",
        "        MRI_pred, _ = model_MRI(target_data)\n",
        "        fake_CT_pred, _ = model_CT(fake_CT)\n",
        "        fake_MRI_pred, _ = model_MRI(fake_MRI)\n",
        "    \n",
        "    '''\n",
        "    prediction discrimator\n",
        "    '''    \n",
        "    domain_logit = model_D_pred_source(CT_pred.detach())\n",
        "    loss = domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device))\n",
        "    \n",
        "    domain_logit = model_D_pred_source(fake_CT_pred.detach())\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device))\n",
        "\n",
        "    domain_logit = model_D_pred_target(MRI_pred.detach())\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "    \n",
        "    domain_logit = model_D_pred_target(fake_MRI_pred.detach())\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "    \n",
        "    '''\n",
        "    image discrimator\n",
        "    '''\n",
        "    domain_logit = model_D_image_source(source_data)\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device))\n",
        "    \n",
        "    domain_logit = model_D_image_source(fake_CT.detach())\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device))\n",
        "                             \n",
        "    domain_logit = model_D_image_target(target_data)\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "                             \n",
        "    domain_logit = model_D_image_target(fake_MRI.detach())\n",
        "    loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device))\n",
        "\n",
        "    \n",
        "    if backward:\n",
        "        loss.backward()\n",
        "        discOptim.step()\n",
        "        discOptim.zero_grad()\n",
        "        \n",
        "        del CT_pred, MRI_pred, fake_CT_pred, fake_MRI_pred, fake_MRI, fake_CT, source_domain_tensor, target_domain_tensor\n",
        "        torch.cuda.empty_cache()\n",
        "    loss /= 8\n",
        "    return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "02ee9e25",
      "metadata": {
        "id": "02ee9e25"
      },
      "outputs": [],
      "source": [
        "source_dataloader, target_dataloader = CT_dataloader_train, dataloader_train\n",
        "test_dataloader = dataloader_test\n",
        "\n",
        "EPOCHS = 3\n",
        "min_target_loss_value = 100\n",
        "source_domain_label = 1\n",
        "target_domain_label = 0\n",
        "\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "    disc_loss_value = 0.0\n",
        "    for i, ((source_data, source_label), (target_data, target_label)) in enumerate(zip(source_dataloader, target_dataloader)):\n",
        "        source_data = source_data.to(device)\n",
        "        target_data = target_data.to(device)\n",
        "    \n",
        "        print(i, end='\\r')\n",
        "        disc_loss_value += train_disc().item()\n",
        "        \n",
        "        del source_data, target_data\n",
        "        torch.cuda.empty_cache()\n",
        "    \n",
        "    disc_loss_value /= (i+1)     \n",
        "    print(f'disc_loss_value:{disc_loss_value}')\n",
        "    if disc_loss_value <= 0.6:\n",
        "        break\n",
        "    \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a1de6482",
      "metadata": {
        "id": "a1de6482"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e02dc98",
      "metadata": {
        "id": "9e02dc98"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ec48068a",
      "metadata": {
        "id": "ec48068a"
      },
      "source": [
        "### GAN gen loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ef8a5e5b",
      "metadata": {
        "id": "ef8a5e5b"
      },
      "outputs": [],
      "source": [
        "def train_gen(backward=True):\n",
        "        # part 1\n",
        "    for model in [model_CT, model_MRI]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "            \n",
        "    for model_D in [model_D_pred_source, model_D_pred_target, model_D_image_source, model_D_image_target]:\n",
        "        for param in model_D.parameters():\n",
        "            param.requires_grad = False\n",
        "            \n",
        "    for model in [netG_CT2MRI, netG_MRI2CT]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "            \n",
        "        # part2 \n",
        "    source_domain_tensor = torch.ones(source_data.data.size()).to(device)\n",
        "    target_domain_tensor = torch.zeros(target_data.data.size()).to(device)\n",
        "\n",
        "    fake_MRI, _ = netG_CT2MRI(source_data, target_domain_tensor)\n",
        "    fake_CT, _ = netG_MRI2CT(target_data, source_domain_tensor)\n",
        "    \n",
        "    'y pred'\n",
        "    CT_pred, _ = model_CT(source_data)\n",
        "    MRI_pred, _ = model_MRI(target_data)\n",
        "    fake_CT_pred, _ = model_CT(fake_CT)\n",
        "    fake_MRI_pred, _ = model_MRI(fake_MRI)\n",
        "    \n",
        "    # part 3\n",
        "    '''\n",
        "    discrimator loss (domain_loss)\n",
        "    '''\n",
        "    'image'\n",
        "    domain_logit = model_D_image_source(fake_CT) \n",
        "    domain_loss = domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device)) # 給相反的標記\n",
        "    \n",
        "    domain_logit = model_D_image_target(fake_MRI)\n",
        "    domain_loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device)) # 給相反的標記\n",
        "    'pred'\n",
        "    domain_logit = model_D_pred_source(fake_CT_pred)\n",
        "    domain_loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(target_domain_label).to(device)) # 給相反的標記\n",
        "\n",
        "    domain_logit = model_D_pred_target(fake_MRI_pred)\n",
        "    domain_loss += domain_criterion(domain_logit, \n",
        "                           torch.zeros(domain_logit.data.size()).fill_(source_domain_label).to(device)) # 給相反的標記\n",
        "\n",
        "    # part 4\n",
        "    '''\n",
        "    cycle loss (consistency loss)\n",
        "    '''\n",
        "    rec_CT, _ = netG_MRI2CT(fake_MRI, source_domain_tensor)\n",
        "    rec_MRI, _ = netG_CT2MRI(fake_CT, target_domain_tensor)\n",
        "    \n",
        "    csis_loss = consist_criterion(rec_CT, source_data)\n",
        "    csis_loss += consist_criterion(rec_MRI, target_data)\n",
        "    \n",
        "    \n",
        "    \n",
        "    loss = domain_loss + csis_loss\n",
        "    if backward:\n",
        "        loss.backward()\n",
        "        unetOptim.step()\n",
        "        unetOptim.zero_grad()\n",
        "        \n",
        "    del CT_pred, MRI_pred, fake_CT_pred, fake_MRI_pred, source_domain_tensor, target_domain_tensor\n",
        "    del fake_MRI, fake_CT, rec_CT, rec_MRI\n",
        "    torch.cuda.empty_cache()\n",
        "    \n",
        "    return domain_loss/4, csis_loss/2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3a625759",
      "metadata": {
        "id": "3a625759"
      },
      "outputs": [],
      "source": [
        "unetOptim = optim.Adam([\n",
        "                {'params': model_MRI.parameters(), 'lr': 1e-1},\n",
        "                {'params': model_CT.parameters()}, # 讓MRI有小一點的lr，ct用default lr\n",
        "                {'params': netG_CT2MRI.parameters()},\n",
        "                {'params': netG_MRI2CT.parameters()},\n",
        "            ], lr=1e-1)  # default lr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c7dc610c",
      "metadata": {
        "id": "c7dc610c"
      },
      "outputs": [],
      "source": [
        "source_dataloader, target_dataloader = CT_dataloader_train, dataloader_train\n",
        "test_dataloader = dataloader_test\n",
        "\n",
        "EPOCHS = 100\n",
        "min_target_loss_value = 100\n",
        "source_domain_label = 1\n",
        "target_domain_label = 0\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "    domain_loss_value = 0.0\n",
        "    csis_loss_value = 0.0\n",
        "    for i, ((source_data, source_label), (target_data, target_label)) in enumerate(zip(source_dataloader, target_dataloader)):\n",
        "        source_data = source_data.to(device)\n",
        "        target_data = target_data.to(device)\n",
        "\n",
        "        print(i, end='\\r')\n",
        "\n",
        "        a, b = train_gen()\n",
        "        domain_loss_value += a.item()\n",
        "        csis_loss_value += b.item()\n",
        "\n",
        "        del source_data, target_data, a, b\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "    domain_loss_value /= (i+1)   \n",
        "    csis_loss_value /= (i+1)   \n",
        "\n",
        "    print(f'epoch: {epoch}, domain_loss_value:{domain_loss_value}, csis_loss_value: {csis_loss_value}')\n",
        "    torch.save(netG_MRI2CT.state_dict(), f'{save_root}netG_MRI2CT.bin')\n",
        "    torch.save(netG_CT2MRI.state_dict(), f'{save_root}netG_CT2MRI.bin')\n",
        "    if domain_loss_value <= 0.4:\n",
        "        break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "66d72888",
      "metadata": {
        "id": "66d72888"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b696cc53",
      "metadata": {
        "id": "b696cc53",
        "outputId": "4b8ea7d8-72b1-40d4-c9e9-284f3cf01905"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 51,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "netG_CT2MRI.load_state_dict(torch.load(f'{save_root}netG_CT2MRI.bin')) \n",
        "netG_MRI2CT.load_state_dict(torch.load(f'{save_root}netG_MRI2CT.bin')) "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "648ecebf",
      "metadata": {
        "id": "648ecebf"
      },
      "source": [
        "### nonlabel_unet (model_MRI)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2ebde9c9",
      "metadata": {
        "id": "2ebde9c9"
      },
      "outputs": [],
      "source": [
        "def train_nonlabel_unet(backward=True):\n",
        "    # part1 \n",
        "    for model in [model_CT, model_MRI]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "            \n",
        "    for model_D in [model_D_pred_source, model_D_pred_target, model_D_image_source, model_D_image_target]:\n",
        "        for param in model_D.parameters():\n",
        "            param.requires_grad = False\n",
        "            \n",
        "    for model in [netG_CT2MRI, netG_MRI2CT]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "            \n",
        "    # part 2 \n",
        "    source_domain_tensor = torch.ones(source_data.data.size()).to(device)\n",
        "    target_domain_tensor = torch.zeros(MRI_image.data.size()).to(device)\n",
        "    \n",
        "    fake_MRI, _ = netG_CT2MRI(source_data, target_domain_tensor)\n",
        "    fake_CT, _ = netG_MRI2CT(MRI_image, source_domain_tensor)\n",
        "#     fake_CT_MMRI, _ = netG_MRI2CT(MRI_image, source_domain_tensor)\n",
        "    \n",
        "    'y pred'\n",
        "    CT_pred, _ = model_CT(source_data)\n",
        "    MRI_pred, _ = model_MRI(MRI_image)\n",
        "#     MRI_pred_MMRI, _ = model_MRI(MRI_image)\n",
        "    \n",
        "    fake_CT_pred, _ = model_CT(fake_CT)\n",
        "    fake_MRI_pred, _ = model_MRI(fake_MRI)\n",
        "#     fake_CT_MMRI_pred, _ = model_CT(fake_CT_MMRI)\n",
        "    \n",
        "    \n",
        "    \n",
        "    # part 3 \n",
        "    '''\n",
        "    class loss\n",
        "    '''\n",
        "    class_loss = class_criterion(CT_pred, source_label)\n",
        "    class_loss += class_criterion(fake_MRI_pred, source_label)\n",
        "    \n",
        "    # part 4\n",
        "    '''\n",
        "    non-label consistency loss\n",
        "    '''\n",
        "    csis_loss = consist_criterion(MRI_pred, fake_CT_pred)\n",
        "        \n",
        "    # part 5\n",
        "    '''\n",
        "    cycle loss (consistency loss)\n",
        "    '''\n",
        "    rec_CT, _ = netG_MRI2CT(fake_MRI, source_domain_tensor)\n",
        "    rec_MRI, _ = netG_CT2MRI(fake_CT, target_domain_tensor)\n",
        "#     rec_MRI_MMRI, _ = netG_CT2MRI(fake_CT_MMRI, target_domain_tensor)\n",
        "    \n",
        "    cycle_loss = consist_criterion(rec_CT, source_data)\n",
        "    cycle_loss += consist_criterion(rec_MRI, target_data)\n",
        "#     cycle_loss += consist_criterion(rec_MRI_MMRI, MRI_image)\n",
        "    \n",
        "    loss = class_loss + csis_loss + cycle_loss\n",
        "    \n",
        "    \n",
        "    if backward:\n",
        "        loss.backward()\n",
        "        unetOptim.step()\n",
        "        unetOptim.zero_grad()\n",
        "        \n",
        "    del source_domain_tensor, target_domain_tensor\n",
        "    del CT_pred, fake_CT, fake_CT_pred, rec_CT\n",
        "    del MRI_pred, fake_MRI, fake_MRI_pred,  rec_MRI, \n",
        "#     del MRI_pred_MMRI, fake_CT_MMRI_pred, fake_CT_MMRI, rec_MRI_MMRI\n",
        "    torch.cuda.empty_cache()\n",
        "    \n",
        "    return class_loss/2, csis_loss, cycle_loss/3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4517242a",
      "metadata": {
        "id": "4517242a"
      },
      "outputs": [],
      "source": [
        "def train_label_unet(backward=True):\n",
        "    # part1 \n",
        "    for model in [model_CT, model_MRI]:\n",
        "        for param in model.parameters():\n",
        "            param.requires_grad = True\n",
        "    \n",
        "    # part2 \n",
        "    source_domain_tensor = torch.ones(source_data.data.size()).to(device)\n",
        "    target_domain_tensor = torch.zeros(source_data.data.size()).to(device)\n",
        "    \n",
        "    fake_MRI, _ = netG_CT2MRI(source_data, target_domain_tensor)\n",
        "    fake_CT, _ = netG_MRI2CT(target_data, source_domain_tensor)\n",
        "    \n",
        "    'y pred'\n",
        "    CT_pred, _ = model_CT(source_data)\n",
        "    MRI_pred, _ = model_MRI(target_data)\n",
        "    fake_CT_pred, _ = model_CT(fake_CT)\n",
        "    fake_MRI_pred, _ = model_MRI(fake_MRI)\n",
        "    \n",
        "    # part3 \n",
        "    '''\n",
        "    class loss\n",
        "    '''\n",
        "    class_loss = class_criterion(CT_pred, source_label)\n",
        "    class_loss += class_criterion(fake_MRI_pred, source_label)\n",
        "    class_loss += class_criterion(MRI_pred, target_label)\n",
        "    class_loss += class_criterion(fake_CT_pred, target_label)\n",
        "    \n",
        "    '''\n",
        "    label consistency loss\n",
        "    '''\n",
        "    csis_loss = consist_criterion(MRI_pred, fake_CT_pred)\n",
        "    csis_loss += consist_criterion(CT_pred, fake_MRI_pred)\n",
        "        \n",
        "    '''\n",
        "    cycle loss (consistency loss)\n",
        "    '''\n",
        "    rec_CT, _ = netG_MRI2CT(fake_MRI, source_domain_tensor)\n",
        "    rec_MRI, _ = netG_CT2MRI(fake_CT, target_domain_tensor)\n",
        "    \n",
        "    cycle_loss = consist_criterion(rec_CT, source_data)\n",
        "    cycle_loss += consist_criterion(rec_MRI, target_data)\n",
        "    \n",
        "    loss = class_loss + csis_loss + cycle_loss\n",
        "    \n",
        "    \n",
        "    if backward:\n",
        "        loss.backward()\n",
        "        unetOptim.step()\n",
        "        unetOptim.zero_grad()\n",
        "        \n",
        "    del CT_pred, MRI_pred, fake_MRI, fake_CT, fake_CT_pred, fake_MRI_pred, source_domain_tensor, target_domain_tensor\n",
        "    del rec_CT, rec_MRI\n",
        "    torch.cuda.empty_cache()\n",
        "    \n",
        "    return class_loss/4, csis_loss/2, cycle_loss/2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8f04a5d4",
      "metadata": {
        "id": "8f04a5d4",
        "outputId": "841af7f1-de9d-4432-a7b5-366ca5fb5378"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch: 0, class_loss_value:0.07108763605356216, csis_loss_value: 0.006933868862688541, cycle_loss_value:0.022816015407443047\n",
            ">>>> testing_loss_value: 0.6814392805099487\n",
            "epoch: 1, class_loss_value:0.05375444516539574, csis_loss_value: 0.005334228277206421, cycle_loss_value:0.02531367726624012\n",
            ">>>> testing_loss_value: 0.6600275784730911\n",
            "epoch: 2, class_loss_value:0.06673753261566162, csis_loss_value: 0.005733566824346781, cycle_loss_value:0.029614800587296486\n",
            ">>>> testing_loss_value: 0.7427933812141418\n",
            "epoch: 3, class_loss_value:0.06925664097070694, csis_loss_value: 0.004644803237169981, cycle_loss_value:0.025160755962133408\n",
            ">>>> testing_loss_value: 0.6136516779661179\n",
            "epoch: 4, class_loss_value:0.06514846533536911, csis_loss_value: 0.0063852895982563496, cycle_loss_value:0.028559641912579536\n",
            ">>>> testing_loss_value: 0.712475374341011\n",
            "epoch: 5, class_loss_value:0.07294768840074539, csis_loss_value: 0.006324455142021179, cycle_loss_value:0.03126144036650658\n",
            ">>>> testing_loss_value: 0.7089353501796722\n",
            "epoch: 6, class_loss_value:0.0721866562962532, csis_loss_value: 0.006731495261192322, cycle_loss_value:0.03278528153896332\n",
            ">>>> testing_loss_value: 0.744776040315628\n",
            "4\r"
          ]
        },
        {
          "ename": "KeyboardInterrupt",
          "evalue": "",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m/tmp/ipykernel_440999/3755202357.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;31m#         print(MRI_image.shape)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'\\r'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 32\u001b[0;31m         \u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_label_unet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     33\u001b[0m         \u001b[0mclass_loss_value\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m         \u001b[0mcsis_loss_value\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/tmp/ipykernel_440999/1268052096.py\u001b[0m in \u001b[0;36mtrain_label_unet\u001b[0;34m(backward)\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     47\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mbackward\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 48\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     49\u001b[0m         \u001b[0munetOptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m         \u001b[0munetOptim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/opt/conda/lib/python3.7/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    253\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    254\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 255\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    256\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    257\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    147\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m    148\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 149\u001b[0;31m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    150\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    151\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "source_dataloader, target_dataloader = CT_dataloader_train, dataloader_train\n",
        "test_dataloader = dataloader_test\n",
        "\n",
        "\n",
        "unetOptim = optim.Adam([\n",
        "                {'params': model_MRI.parameters(),},\n",
        "                {'params': model_CT.parameters()}, # 讓MRI有小一點的lr，ct用default lr\n",
        "                {'params': netG_CT2MRI.parameters()},\n",
        "                {'params': netG_MRI2CT.parameters()},\n",
        "            ], lr=1e-1)  # default lr\n",
        "\n",
        "EPOCHS = 150\n",
        "min_target_loss_value = 100\n",
        "source_domain_label = 1\n",
        "target_domain_label = 0\n",
        "\n",
        "for epoch in range(EPOCHS):  \n",
        "    class_loss_value = 0.0\n",
        "    csis_loss_value = 0.0\n",
        "    cycle_loss_value = 0.0\n",
        "    testing_loss_value = 0.0\n",
        "    for i, ((source_data, source_label), (target_data, target_label), (MRI_image)) in enumerate(zip(source_dataloader,\n",
        "                                                                                       target_dataloader, MRI_image_dataloader)):\n",
        "        source_data = source_data.to(device)\n",
        "        source_label = source_label.to(device)\n",
        "        target_data = target_data.to(device)\n",
        "        target_label = target_label.to(device)\n",
        "        MRI_image = MRI_image.to(device)\n",
        "        \n",
        "#         print(MRI_image.shape)\n",
        "        print(i, end='\\r')\n",
        "        a,b,c = train_label_unet()\n",
        "        class_loss_value += a\n",
        "        csis_loss_value += b\n",
        "        cycle_loss_value += c\n",
        "        \n",
        "        a, b, c = train_nonlabel_unet()\n",
        "        class_loss_value += a\n",
        "        csis_loss_value += b\n",
        "        cycle_loss_value += c\n",
        "        if epoch%5==0:\n",
        "            train_disc()\n",
        "        train_gen()\n",
        "        \n",
        "        class_loss_value /= 2\n",
        "        csis_loss_value /= 2\n",
        "        cycle_loss_value /= 2\n",
        "        \n",
        "        del source_data, source_label, target_data, a, b, c\n",
        "        torch.cuda.empty_cache()\n",
        "        \n",
        "    class_loss_value /= (i+1)   \n",
        "    csis_loss_value /= (i+1)   \n",
        "    cycle_loss_value /= (i+1)\n",
        "    \n",
        "    for i2, (target_data, target_label) in enumerate(test_dataloader):\n",
        "        target_data = target_data.to(device)\n",
        "        target_label = target_label.to(device)\n",
        "        \n",
        "        y_pred, _ = model_MRI(target_data)\n",
        "        loss = class_criterion(y_pred, target_label)\n",
        "        testing_loss_value += loss.item()\n",
        "        \n",
        "        del target_data, target_label\n",
        "        torch.cuda.empty_cache()\n",
        "    testing_loss_value /= (i2+1)\n",
        "\n",
        "\n",
        "\n",
        "    \n",
        "    print(f'epoch: {epoch}, class_loss_value:{class_loss_value}, csis_loss_value: {csis_loss_value}, cycle_loss_value:{cycle_loss_value}')\n",
        "    print(f'>>>> testing_loss_value: {testing_loss_value}')\n",
        "    if testing_loss_value < 0.9*min_target_loss_value:\n",
        "        min_target_loss_value = testing_loss_value\n",
        "        torch.save(netG_MRI2CT.state_dict(), f'{save_root}best_netG_MRI2CT.bin')\n",
        "        torch.save(netG_CT2MRI.state_dict(), f'{save_root}best_netG_CT2MRI.bin')\n",
        "        torch.save(model_CT.state_dict(), f'{save_root}best_model_CT.bin')\n",
        "        torch.save(model_MRI.state_dict(), f'{save_root}best_model_MRI.bin')\n",
        "    else:\n",
        "\n",
        "        torch.save(netG_MRI2CT.state_dict(), f'{save_root}netG_MRI2CT.bin')\n",
        "        torch.save(netG_CT2MRI.state_dict(), f'{save_root}netG_CT2MRI.bin')\n",
        "        torch.save(model_CT.state_dict(), f'{save_root}model_CT.bin')\n",
        "        torch.save(model_MRI.state_dict(), f'{save_root}model_MRI.bin')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6d85958b",
      "metadata": {
        "id": "6d85958b"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d21a618",
      "metadata": {
        "id": "9d21a618"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a3d69694",
      "metadata": {
        "id": "a3d69694"
      },
      "source": [
        "## 測試模型"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f6aad708",
      "metadata": {
        "id": "f6aad708"
      },
      "source": [
        "### load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e65ab920",
      "metadata": {
        "id": "e65ab920"
      },
      "outputs": [],
      "source": [
        "root = '/home/jovyan/DA/DATA/TA_data/CHAOS_AIAdatasets/2_Domain_Adaptation_dataset/testset'\n",
        "\n",
        "dic = {}\n",
        "for a,b,c in os.walk(root, topdown=True):\n",
        "    if len(c)>0: # 當前目錄內包含檔案\n",
        "        if not a.__contains__('OutPhase'):\n",
        "            dic[a] = c\n",
        "dic.keys()\n",
        "dataset = {}\n",
        "# lis = ['CT', 'MRI']\n",
        "lis = ['DICOM_anon', 'Ground']\n",
        "for task in lis:\n",
        "    class_lis = []\n",
        "    for sub_folder in dic.keys():\n",
        "        if task in sub_folder.split('/'):\n",
        "            class_lis+=[sub_folder+'/'+filename for filename in dic[sub_folder]]\n",
        "    dataset[task] = class_lis\n",
        "    \n",
        "\n",
        "dataset['MRI_T2SPIR_test'] = sorted([i for i in dataset['DICOM_anon'] if 'T2SPIR' in i])\n",
        "dataset['MRI_T2SPIR_mask'] = sorted([i for i in dataset['Ground'] if 'T2SPIR' in i])\n",
        "assert len(dataset['MRI_T2SPIR_test'])==len(dataset['MRI_T2SPIR_mask'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb25a5c7",
      "metadata": {
        "id": "eb25a5c7"
      },
      "outputs": [],
      "source": [
        "MRI_test = list(zip(dataset['MRI_T2SPIR_test'], dataset['MRI_T2SPIR_mask']))\n",
        "\n",
        "# 建議同時間只有8個(256,256)的sample進行計算 (Total = BATCH_SIZE*MULTIPLE_BATCH)\n",
        "dataset_test = CTMRI_ImageDataset(MRI_test, dtype='MRI', transform=target_transform) # **如果要正式使用要記得把這裡換成X_test\n",
        "dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "91f75efc",
      "metadata": {
        "id": "91f75efc"
      },
      "outputs": [],
      "source": [
        "a = iter(dataloader_test)\n",
        "bx, by = a.next()\n",
        "print(bx.shape, by.shape)\n",
        "del a, bx, by"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c213273e",
      "metadata": {
        "id": "c213273e"
      },
      "outputs": [],
      "source": [
        "model_MRI =  Unet(out_sz=(HEIGHT, WIDTH), out_channels=out_channels, activation=nn.Sigmoid(), multi_level=0).to(device)\n",
        "model_MRI.load_state_dict(torch.load(f'{save_root}model_MRI.bin')) \n",
        "class_criterion = DiceLoss()\n",
        "# class_criterion = FocalLoss()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6a288939",
      "metadata": {
        "id": "6a288939"
      },
      "outputs": [],
      "source": [
        "#@title [Colab Form] 測試不同的dataloader \n",
        "#@markdown 可能的選擇包括 'dataloader_train', 'dataloader_test', 'CT_dataloader_test'\n",
        "# print(w.value)\n",
        "# test_dataloader = dic[w.value]\n",
        "import pandas as pd\n",
        "# test_dataloader = dataloader_test\n",
        "\n",
        "\n",
        "testing_loss = 0.0\n",
        "# 建立3個block\n",
        "\n",
        "with torch.no_grad(): # 避免torch計算gradient產生記憶體負擔\n",
        "    for i, data in enumerate(dataloader_test, 1): \n",
        "        torch.cuda.empty_cache()\n",
        "        image, mask = data\n",
        "\n",
        "        image = image.to(device)\n",
        "        mask = mask.to(device)\n",
        "        \n",
        "        outputs, _ = model_MRI(image)\n",
        "\n",
        "        if True:\n",
        "            thres = 0.1\n",
        "            outputs[outputs<thres] = 0 \n",
        "            outputs[outputs!=0] = 1\n",
        "            outputs = outputs.long()\n",
        "\n",
        "#             outputs = mask # ground truth diceloss=0\n",
        "#         a = outputs\n",
        "        loss = class_criterion(outputs, mask)\n",
        "        testing_loss += loss\n",
        "    loss =  testing_loss/len(dataloader_test)\n",
        "    \n",
        "del outputs\n",
        "        \n",
        "loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a528bc60",
      "metadata": {
        "id": "a528bc60"
      },
      "outputs": [],
      "source": [
        "def mask2rle(img):\n",
        "    \"\"\"\n",
        "    img: numpy array, 1 - mask, 0 - background\n",
        "    Returns run length as string formatted\n",
        "    \"\"\"\n",
        "    pixels = img.flatten()\n",
        "    pixels = np.concatenate([[0], pixels, [0]])\n",
        "    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n",
        "    runs[1::2] -= runs[::2]\n",
        "    return ' '.join(str(x) for x in runs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d05c8a5f",
      "metadata": {
        "id": "d05c8a5f"
      },
      "outputs": [],
      "source": [
        "submission = []\n",
        "for file_list, dataloader in zip([MRI_test], [dataloader_test]):\n",
        "    test_list = [['-'.join([str(i[0].split('/')[idx]) for idx in [-4,-3,-1]])] for i in file_list]\n",
        "#     len(CT_test_list)\n",
        "\n",
        "    dataloader = iter(dataloader)\n",
        "    print(len(file_list))\n",
        "    i = 0\n",
        "    while 1:\n",
        "        try:\n",
        "            image, mask = dataloader.next()\n",
        "            image = image.to(device)\n",
        "#             mask = mask.to(device)\n",
        "            \n",
        "            outputs, _ = model_MRI(image)\n",
        "        \n",
        "\n",
        "            outputs[outputs<0.1] = 0.\n",
        "            outputs[outputs!=0] = 1.\n",
        "            outputs = outputs.detach().cpu()\n",
        "\n",
        "            for out in outputs:\n",
        "                test_list[i].append(mask2rle(out))\n",
        "\n",
        "                i += 1\n",
        "        except StopIteration:\n",
        "            print(i)\n",
        "            print('complete')\n",
        "            break\n",
        "    submission+=test_list\n",
        "    assert i==len(test_list)\n",
        "    \n",
        "pd.DataFrame(submission, columns=['filename', 'rle']).to_csv('submission.csv', index=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4d2a69ff",
      "metadata": {
        "id": "4d2a69ff"
      },
      "outputs": [],
      "source": [
        "# #@title [Colab Form] 測試不同的dataloader \n",
        "# #@markdown 可能的選擇包括 'dataloader_test', 'CT_dataloader_test'\n",
        "\n",
        "# test_dataloader = CT_dataloader_test\n",
        "\n",
        "\n",
        "# def test():\n",
        "#     with torch.no_grad(): # 避免torch計算gradient產生記憶體負擔\n",
        "#         for i, data in enumerate(test_dataloader, 1): \n",
        "#             torch.cuda.empty_cache()\n",
        "#             image, mask = data\n",
        "#             print(len(image))\n",
        "\n",
        "#             image = image.to(device)\n",
        "#             mask = mask.to(device)\n",
        "\n",
        "#             outputs, _ = model_MRI(image)\n",
        "# #             outputs = label_predictor(enc_ftrs)\n",
        "# #             print(outputs)\n",
        "#             # outputs = model(image)\n",
        "#             loss = class_criterion(outputs, mask)\n",
        "#             print(loss)\n",
        "\n",
        "#             '''\n",
        "#             設定閥值輸出:\n",
        "#             其實不設定也是可以輸出，不設定的情況會用灰階輸出，細節會更好一點\n",
        "#             '''\n",
        "#             if True: # 可以改成false，採灰階輸出， True的話會通過閾值輸出\n",
        "#                 # a = sum([i.max() for i in outputs])/len(outputs)\n",
        "#                 # b = sum([i.std() for i in outputs])/len(outputs)\n",
        "#                 # \n",
        "#                 # thres = a-1*b\n",
        "#                 thres = 0.1\n",
        "#                 outputs[outputs>=thres] = 1 \n",
        "#                 outputs[outputs<thres] = 0\n",
        "#                 outputs = outputs.long()\n",
        "\n",
        "#             print(image[:,0,:,:].shape)\n",
        "#             print(mask.shape)\n",
        "#             print(outputs.shape)\n",
        "\n",
        "#             img_process = lambda image:image[:,0,:,:].detach().cpu().numpy()\n",
        "#             mask_process = lambda mask:mask.squeeze(1).detach().cpu().numpy()\n",
        "\n",
        "#             for x, m, p in [*zip(img_process(image), mask_process(mask), mask_process(outputs))]:\n",
        "#                   # show_image_mask(x, m, p)\n",
        "#                 show_image(x,m,p)\n",
        "\n",
        "#             del image, mask, outputs, data, loss # 把記憶體清掉\n",
        "#             torch.cuda.empty_cache()\n",
        "#             break\n",
        "  \n",
        "    \n",
        "# tmp = test()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e2213eb1",
      "metadata": {
        "id": "e2213eb1"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "daTest_0416_selfsup.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}