registration_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json \n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from preprocess import *\n",
    "import SimpleITK as sitk\n",
    "from collections import OrderedDict\n",
    "import os\n",
    "from register_images import *\n",
    "import sys\n",
    "sys.path.insert(0, '../')\n",
    "from parse_registration_json import ParserRegistrationJson\n",
    "from parse_study_dict import ParserStudyDict\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'do_affine'\n",
      "'do_deformable'\n",
      "'do_reconstruction'\n",
      "'fast_execution'\n",
      "'use_imaging_constraints'\n",
      "Reading aaa0069 Study Json C:\\Users\\weishao\\Desktop\\Code_ProsRegNet\\jsonData\\reg_aaa0069.json\n"
     ]
    }
   ],
   "source": [
    "####### INPUTS\n",
    "json_path = \"jsonData/TCIA_FUSION.json\"\n",
    "preprocess_moving = True\n",
    "preprocess_fixed = True\n",
    "run_registration = True\n",
    "extension = 'nii.gz'\n",
    "timings = {}\n",
    "\n",
    "try:\n",
    "    with open('coord.txt') as f:\n",
    "        coord = json.load(f)   \n",
    "        \n",
    "except:\n",
    "    coord = {}\n",
    "\n",
    "############### START REGISTRATION HERE\n",
    "\n",
    "json_obj = ParserRegistrationJson(json_path)\n",
    "\n",
    "studies = json_obj.studies\n",
    "toProcess = json_obj.ToProcess\n",
    "outputPath = json_obj.output_path\n",
    "#cases = toProcess.keys()\n",
    "\n",
    "if not os.path.isdir(outputPath):\n",
    "    os.mkdir(outputPath) \n",
    "\n",
    "###### PREPROCESSING DESTINATIONS ######################################\n",
    "preprocess_moving_dest = outputPath + 'preprocess\\\\hist\\\\'\n",
    "preprocess_fixed_dest = outputPath + 'preprocess\\\\mri\\\\'\n",
    "\n",
    "if not os.path.isdir(outputPath + 'preprocess\\\\'):\n",
    "    os.mkdir(outputPath + 'preprocess\\\\')\n",
    "\n",
    "if not os.path.isdir(preprocess_moving_dest):\n",
    "    os.mkdir(preprocess_moving_dest)\n",
    "    \n",
    "if not os.path.isdir(preprocess_fixed_dest):\n",
    "    os.mkdir(preprocess_fixed_dest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx Processing aaa0069 xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\n",
      "Preprocessing moving sid: aaa0069 ...\n",
      "Finished preprocessing aaa0069\n",
      "Preprocessing fixed case: aaa0069 ...\n",
      "Finished processing fixed mha aaa0069\n",
      ".............................. Begin deep learning registration for aaa0069..............................\n",
      "Creating CNN model...\n",
      "Loading trained model weights...\n",
      "Registration done in  0.156(min)\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# start doing preprocessing on each case and register\n",
    "for s in json_obj.studies:\n",
    "    if json_obj.ToProcess:\n",
    "        if not (s in json_obj.ToProcess):\n",
    "            print(\"Skipping\", s)\n",
    "            continue\n",
    "\n",
    "    print(\"x\"*30, \"Processing\", s,\"x\"*30)\n",
    "    studyDict = json_obj.studies[s] \n",
    "\n",
    "\n",
    "    studyParser = ParserStudyDict(studyDict)\n",
    "    \n",
    "    sid = studyParser.id\n",
    "    fixed_img_mha = studyParser.fixed_filename\n",
    "    fixed_seg = studyParser.fixed_segmentation_filename\n",
    "    moving_dict = studyParser.ReadMovingImage()\n",
    "\n",
    "    ###### PREPROCESSING HISTOLOGY HERE #############################################################\n",
    "    if preprocess_moving == True: \n",
    "        print('Preprocessing moving sid:', sid, '...')\n",
    "        preprocess_hist(moving_dict, preprocess_moving_dest, sid)\n",
    "        print('Finished preprocessing', sid)\n",
    "\n",
    "    ###### PREPROCESSING MRI HERE #############################################################\n",
    "    if preprocess_fixed == True:\n",
    "        print (\"Preprocessing fixed case:\", sid, '...')\n",
    "\n",
    "        coord = preprocess_mri(fixed_img_mha, fixed_seg, preprocess_fixed_dest, coord, sid)\n",
    "\n",
    "        print(\"Finished processing fixed mha\", sid)\n",
    "\n",
    "        with open('coord.txt', 'w') as json_file: \n",
    "            json.dump(coord, json_file)\n",
    "    ##### ALIGNMENT HERE ########################################################################\n",
    "    if run_registration == True: \n",
    "        \n",
    "        ######## LOAD MODELS\n",
    "        print('.'*30, 'Begin deep learning registration for ' + sid + '.'*30)\n",
    "\n",
    "        try:\n",
    "            model_cache\n",
    "        except NameError:\n",
    "            feature_extraction_cnn = 'resnet101'\n",
    "\n",
    "            model_aff_path = 'C:/Users/weishao/Desktop/Code_ProsRegNet/trained_models/best_CombinedLoss_affine_resnet101.pth.tar'\n",
    "            model_tps_path = 'C:/Users/weishao/Desktop/Code_ProsRegNet/trained_models/best_CombinedLoss_tps_resnet101.pth.tar'\n",
    "\n",
    "            model_cache = load_models(feature_extraction_cnn, model_aff_path, model_tps_path, do_deformable=True)\n",
    "        \n",
    "        start = time.time()\n",
    "        output3D_cache = register(preprocess_moving_dest + sid + '\\\\', preprocess_fixed_dest + sid + '\\\\', coord, model_cache, sid)\n",
    "        out3Dhist_highRes, out3Dmri_highRes, out3Dcancer_highRes, out3D_region00, out3D_region10, out3D_region09, out3Dmri_mask = output3D_cache\n",
    "        end = time.time()\n",
    "        print(\"Registration done in {:6.3f}(min)\".format((end-start)/60.0))\n",
    "        imMri = sitk.ReadImage(fixed_img_mha)\n",
    "        mriOrigin = imMri[:,:,coord[sid]['slice'][0]:coord[sid]['slice'][-1]].GetOrigin()\n",
    "        mriSpace  = imMri.GetSpacing()\n",
    "        mriDirection = imMri.GetDirection()\n",
    "\n",
    "        imSpatialInfo = (mriOrigin, mriSpace, mriDirection)\n",
    "\n",
    "        # write output hist 3D volume to .nii.gz format\n",
    "        fn_moving_highRes = '_moved_highres_rgb.'\n",
    "        output_results_high_res(preprocess_moving_dest + sid + '\\\\',preprocess_fixed_dest,outputPath, out3Dhist_highRes, sid, fn_moving_highRes, imSpatialInfo, coord, imMri, extension = \"nii.gz\")\n",
    "\n",
    "        #write output mri 3D volume to .nii.gz format\n",
    "        fn_fixed_highRes = '_fixed_image.'\n",
    "        output_results(outputPath, out3Dmri_highRes, sid, fn_fixed_highRes, imSpatialInfo, extension = \"nii.gz\")\n",
    "\n",
    "        #write output cancer outline 3D volume to .nii.gz format\n",
    "        fn_cancer_highRes = '_moved_highres_region01_label.'\n",
    "        output_results_high_res(preprocess_moving_dest + sid + '\\\\',preprocess_fixed_dest,outputPath, out3Dcancer_highRes, sid, fn_cancer_highRes, imSpatialInfo, coord, imMri, extension = \"nii.gz\")\n",
    "        \n",
    "        #write region00\n",
    "        fn_region00 = '_moved_highres_region00_label.'\n",
    "        output_results_high_res(preprocess_moving_dest + sid + '\\\\',preprocess_fixed_dest,outputPath, out3D_region00, sid, fn_region00, imSpatialInfo, coord, imMri, extension = \"nii.gz\")\n",
    "        \n",
    "        #write region10\n",
    "        fn_region00 = '_moved_highres_region10_label.'\n",
    "        output_results_high_res(preprocess_moving_dest + sid + '\\\\',preprocess_fixed_dest,outputPath, out3D_region10, sid, fn_region00, imSpatialInfo, coord, imMri, extension = \"nii.gz\")\n",
    "        \n",
    "        #write region09\n",
    "        fn_region00 = '_moved_highres_region09_label.'\n",
    "        output_results_high_res(preprocess_moving_dest + sid + '\\\\',preprocess_fixed_dest,outputPath, out3D_region09, sid, fn_region00, imSpatialInfo, coord, imMri, extension = \"nii.gz\")\n",
    "        \n",
    "        #write mriMask\n",
    "        fn_mriMask = '_fixed_mask_label.'\n",
    "        output_results(outputPath, out3Dmri_mask, sid, fn_mriMask, imSpatialInfo, extension = \"nii.gz\")\n",
    "\n",
    "        timings[s] = (end-start)/60.0\n",
    "        print('Done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json.dump(timings, open(\"timings.txt\",'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}