prepare_pseudo_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import SimpleITK as sitk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "felix2paot_label_mapping = {\n",
    "    1: 8,    # Aorta\n",
    "    2: 12,    # Adrenal Gland\n",
    "    4: 33,    # CBD\n",
    "    5: 25,    # Celiac abdominal aorta\n",
    "    6: 18,    # Colon\n",
    "    7: 14,    # Duodenum\n",
    "    8: 4,    # Gall Bladder\n",
    "    9: 9,    # IVC / postcava\n",
    "    10: 3,    # L Kidney\n",
    "    11: 2,    # R Kidney\n",
    "    12: 6,    # Liver\n",
    "    13: 11,    # Pancreas\n",
    "    14: 34,    # pancreatic duct\n",
    "    15: 35,    # SMA\n",
    "    16: 19,    # small bowel / intestine\n",
    "    18: 1,    # spleen\n",
    "    19: 7,    # stomach\n",
    "    20: 10,    # veins\n",
    "    21: 36,    # Kidney_LtRV\n",
    "    22: 37,    # Kidney_RtRV\n",
    "    24: 38,    # CBD stent\n",
    "}\n",
    "\n",
    "paot2felix_label_mapping = {v: k for k, v in felix2paot_label_mapping.items()}\n",
    "paot2felix_label_mapping[13] = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'FELIX0449_ARTERIAL.nii.gz'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fnames = open('dataset/dataset_list/felix_small_step2_train.txt').readlines()\n",
    "fnames = [x.split()[1] for x in fnames]\n",
    "fnames = [x.split('/')[-1] for x in fnames]\n",
    "fnames[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = 'output/felix_onehot_step1/step2_pseudo_label'\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "for fname in fnames:\n",
    "    label = sitk.ReadImage(os.path.join('Input Image Dir', fname))\n",
    "    pred = sitk.ReadImage(os.path.join('output/felix_onehot_step1/test_epoch_100/predict', fname.split('.')[0], fname.split('.')[0] + '_pred.nii.gz'))\n",
    "    label_np = sitk.GetArrayFromImage(label)\n",
    "    pred_np = sitk.GetArrayFromImage(pred)\n",
    "    pseudo = np.zeros_like(label_np)    # felix style\n",
    "    label_cls = list(np.unique(label_np))    # felix style\n",
    "    label_cls.remove(0)\n",
    "    for paot_cls in [1, 2, 3, 4, 6, 9, 11, 12, 13]:    # pretrained classes\n",
    "        felix_cls = paot2felix_label_mapping[paot_cls]\n",
    "        pseudo[pred_np == paot_cls] = felix_cls\n",
    "        if felix_cls in label_cls:\n",
    "            label_cls.remove(felix_cls)\n",
    "    for icls in label_cls:\n",
    "        pseudo[label_np == icls] = icls\n",
    "    pseudo = sitk.GetImageFromArray(pseudo)\n",
    "    pseudo.SetOrigin(label.GetOrigin())\n",
    "    pseudo.SetSpacing(label.GetSpacing())\n",
    "    pseudo.SetDirection(label.GetDirection())\n",
    "    sitk.WriteImage(pseudo, os.path.join(save_dir, fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'FELIX1260_ARTERIAL.nii.gz'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fnames = open('dataset/dataset_list/felix_small_step3_train.txt').readlines()\n",
    "fnames = [x.split()[1] for x in fnames]\n",
    "fnames = [x.split('/')[-1] for x in fnames]\n",
    "fnames[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = 'output/felix_v3_step2_run2/step3_pseudo_label'\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "for fname in fnames:\n",
    "    label = sitk.ReadImage(os.path.join('Input Image Dir', fname))\n",
    "    pred = sitk.ReadImage(os.path.join('output/felix_v3_step2_run2/test_epoch_100/predict', fname.split('.')[0], fname.split('.')[0] + '_pred.nii.gz'))\n",
    "    label_np = sitk.GetArrayFromImage(label)\n",
    "    pred_np = sitk.GetArrayFromImage(pred)\n",
    "    pseudo = np.zeros_like(label_np)    # felix style\n",
    "    label_cls = list(np.unique(label_np))    # felix style\n",
    "    label_cls.remove(0)\n",
    "    for paot_cls in [1, 2, 3, 4, 6, 9, 11, 12, 13, 7, 14, 18, 19]:    # pretrained classes\n",
    "        felix_cls = paot2felix_label_mapping[paot_cls]\n",
    "        pseudo[pred_np == paot_cls] = felix_cls\n",
    "        if felix_cls in label_cls:\n",
    "            label_cls.remove(felix_cls)\n",
    "    for icls in label_cls:\n",
    "        pseudo[label_np == icls] = icls\n",
    "    pseudo = sitk.GetImageFromArray(pseudo)\n",
    "    pseudo.SetOrigin(label.GetOrigin())\n",
    "    pseudo.SetSpacing(label.GetSpacing())\n",
    "    pseudo.SetDirection(label.GetDirection())\n",
    "    sitk.WriteImage(pseudo, os.path.join(save_dir, fname))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BTCV -> LiTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "lits2paot_label_mapping = {\n",
    "    1: 6,    # Liver\n",
    "    2: 27,    # Liver tumor\n",
    "}\n",
    "\n",
    "paot2lits_label_mapping = {v: k for k, v in lits2paot_label_mapping.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "liver_1\n",
      "liver_10\n",
      "liver_100\n",
      "liver_101\n",
      "liver_102\n",
      "liver_103\n",
      "liver_104\n",
      "liver_106\n",
      "liver_107\n",
      "liver_108\n",
      "liver_110\n",
      "liver_112\n",
      "liver_113\n",
      "liver_114\n",
      "liver_115\n",
      "liver_116\n",
      "liver_117\n",
      "liver_119\n",
      "liver_12\n",
      "liver_120\n",
      "liver_122\n",
      "liver_123\n",
      "liver_125\n",
      "liver_126\n",
      "liver_129\n",
      "liver_13\n",
      "liver_14\n",
      "liver_15\n",
      "liver_17\n",
      "liver_18\n",
      "liver_19\n",
      "liver_2\n",
      "liver_22\n",
      "liver_23\n",
      "liver_24\n",
      "liver_25\n",
      "liver_26\n",
      "liver_27\n",
      "liver_29\n",
      "liver_3\n",
      "liver_30\n",
      "liver_31\n",
      "liver_32\n",
      "liver_33\n",
      "liver_34\n",
      "liver_35\n",
      "liver_36\n",
      "liver_39\n",
      "liver_4\n",
      "liver_40\n",
      "liver_42\n",
      "liver_44\n",
      "liver_45\n",
      "liver_46\n",
      "liver_49\n",
      "liver_5\n",
      "liver_50\n",
      "liver_51\n",
      "liver_52\n",
      "liver_54\n",
      "liver_55\n",
      "liver_56\n",
      "liver_57\n",
      "liver_59\n",
      "liver_6\n",
      "liver_63\n",
      "liver_64\n",
      "liver_65\n",
      "liver_66\n",
      "liver_67\n",
      "liver_68\n",
      "liver_69\n",
      "liver_7\n",
      "liver_71\n",
      "liver_72\n",
      "liver_74\n",
      "liver_75\n",
      "liver_77\n",
      "liver_78\n",
      "liver_8\n",
      "liver_80\n",
      "liver_81\n",
      "liver_84\n",
      "liver_86\n",
      "liver_87\n",
      "liver_88\n",
      "liver_89\n",
      "liver_9\n",
      "liver_90\n",
      "liver_92\n",
      "liver_93\n",
      "liver_94\n",
      "liver_95\n",
      "liver_97\n",
      "liver_99\n"
     ]
    }
   ],
   "source": [
    "fnames = open('./dataset/dataset_list/lits_train.txt').readlines()\n",
    "fnames = [fname.split('/')[-1].split('.')[0] for fname in fnames]\n",
    "os.makedirs('output/btcv_v3_cls-all/pseudo_label_lits', exist_ok=True)\n",
    "for fname in fnames:\n",
    "    print(fname)\n",
    "    label = sitk.ReadImage(os.path.join('LiTS label dir', fname))\n",
    "    pred = sitk.ReadImage(os.path.join('output/btcv_v3_cls-all/test_epoch_100/predict', fname, fname + '_pred.nii.gz'))\n",
    "    label_np = sitk.GetArrayFromImage(label)\n",
    "    pred_np = sitk.GetArrayFromImage(pred)\n",
    "    pseudo = np.zeros_like(label_np)\n",
    "    label_cls = list(np.unique(label_np))    # lits style\n",
    "    label_cls.remove(0)\n",
    "    for icls in [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14]:    # no liver\n",
    "        pseudo[pred_np == icls] = icls\n",
    "        # if icls in label_cls:\n",
    "        #     label_cls.remove(icls)\n",
    "    for lits_cls in label_cls:\n",
    "        pseudo[label_np == lits_cls] = lits2paot_label_mapping[lits_cls]\n",
    "    pseudo = sitk.GetImageFromArray(pseudo)\n",
    "    pseudo.CopyInformation(label)\n",
    "    sitk.WriteImage(pseudo, os.path.join('output/btcv_v3_cls-all/pseudo_label_lits/', fname + '.nii.gz'))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = open('./dataset/dataset_list/lits_train.txt').readlines()\n",
    "fnames = [fname.split('/')[-1].split('.')[0] for fname in fnames]\n",
    "os.makedirs('output/btcv_swin-ce_cls-all/pseudo_label_lits', exist_ok=True)\n",
    "for fname in fnames:\n",
    "    label = sitk.ReadImage(os.path.join('LiTS label dir', fname))\n",
    "    pred = sitk.ReadImage(os.path.join('output/btcv_swin-ce_cls-all/test_epoch_100/predict', fname, fname + '_pred.nii.gz'))\n",
    "    label_np = sitk.GetArrayFromImage(label)\n",
    "    pred_np = sitk.GetArrayFromImage(pred)\n",
    "    pseudo = np.zeros_like(label_np)\n",
    "    label_cls = list(np.unique(label_np))    # lits style\n",
    "    label_cls.remove(0)\n",
    "    for icls in [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14]:\n",
    "        pseudo[pred_np == icls] = icls\n",
    "        # if icls in label_cls:\n",
    "        #     label_cls.remove(icls)\n",
    "    for lits_cls in label_cls:\n",
    "        pseudo[label_np == lits_cls] = lits2paot_label_mapping[lits_cls]\n",
    "    pseudo = sitk.GetImageFromArray(pseudo)\n",
    "    pseudo.CopyInformation(label)\n",
    "    sitk.WriteImage(pseudo, os.path.join('output/btcv_swin-ce_cls-all/pseudo_label_lits/', fname + '.nii.gz'))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = open('./dataset/dataset_list/lits_train.txt').readlines()\n",
    "fnames = [fname.split('/')[-1].split('.')[0] for fname in fnames]\n",
    "os.makedirs('output/btcv_onehot_cls-all/pseudo_label_lits', exist_ok=True)\n",
    "for fname in fnames:\n",
    "    print(fname)\n",
    "    label = sitk.ReadImage(os.path.join('LiTS label dir', fname))\n",
    "    pred = sitk.ReadImage(os.path.join('output/btcv_onehot_cls-all/test_epoch_100/predict', fname, fname + '_pred.nii.gz'))\n",
    "    label_np = sitk.GetArrayFromImage(label)\n",
    "    pred_np = sitk.GetArrayFromImage(pred)\n",
    "    pseudo = np.zeros_like(label_np)\n",
    "    label_cls = list(np.unique(label_np))    # lits style\n",
    "    label_cls.remove(0)\n",
    "    for icls in [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14]:    # no liver\n",
    "        pseudo[pred_np == icls] = icls\n",
    "        # if icls in label_cls:\n",
    "        #     label_cls.remove(icls)\n",
    "    for lits_cls in label_cls:\n",
    "        pseudo[label_np == lits_cls] = lits2paot_label_mapping[lits_cls]\n",
    "    pseudo = sitk.GetImageFromArray(pseudo)\n",
    "    pseudo.CopyInformation(label)\n",
    "    sitk.WriteImage(pseudo, os.path.join('output/btcv_onehot_cls-all/pseudo_label_lits/', fname + '.nii.gz'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch1.11",
   "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.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "be14455941fa8fdac530bc913020d2951320e551ab8e5b1887df6d07763caad0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}