ENH: Use of DataSinker node to save final results

Medical-Image-Analysis-Laboratory · Jun 5, 2020 · b130576 · b130576
1 parent c4fbabf
commit b130576
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 35 deletions.
diff --git a/notebooks/srr.ipynb b/notebooks/srr.ipynb
@@ -26,7 +26,7 @@
     "import shutil\n",
     "\n",
     "# Imports from nipype\n",
-    "from nipype.interfaces.io import BIDSDataGrabber,DataGrabber\n",
+    "from nipype.interfaces.io import BIDSDataGrabber,DataGrabber, DataSink\n",
     "from nipype.pipeline import Node, Workflow\n",
     "\n",
     "# Import the implemented interface from pymialsrtk\n",
@@ -76,10 +76,8 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "raw",
    "metadata": {},
-   "outputs": [],
    "source": [
     "# \n",
     "# Copy result files to BIDS\n",
@@ -207,30 +205,31 @@
    "outputs": [],
    "source": [
     "## Node linkage\n",
-    "def create_workflow(bids_dir, subject, p_stacksOrder, session=None):\n",
+    "def create_workflow(bids_dir, process_dir, subject, p_stacksOrder, session=None, deltatTV = 0.01, lambdaTV = 0.75, primal_dual_loops=10):\n",
     "#     wf_base_dir = os.path.join(\"{}\".format(output_dir),\"superres-mri\",\"sub-{}\".format(subject),\"nipype\")\n",
-    "    output_dir = os.path.join(\"{}\".format(bids_dir),\"derivatives\",\"superres-mri\")\n",
     "    \n",
     "    if session is None:\n",
-    "        wf_base_dir = os.path.join(output_dir, subject)\n",
+    "        wf_base_dir = os.path.join(process_dir, subject)\n",
+    "        process_dir = os.path.join(process_dir, subject)\n",
     "    else:\n",
-    "        wf_base_dir = os.path.join(output_dir, subject, session)\n",
+    "        wf_base_dir = os.path.join(process_dir, subject, session)\n",
+    "        process_dir = os.path.join(process_dir, subject, session)\n",
     "\n",
-    "    if not os.path.exists(output_dir):\n",
-    "        os.makedirs(output_dir)\n",
-    "    print(\"Ouput directory: {}\".format(wf_base_dir))\n",
+    "    if not os.path.exists(process_dir):\n",
+    "        os.makedirs(process_dir)\n",
+    "    print(\"Process directory: {}\".format(wf_base_dir))\n",
     "\n",
     "    wf = Workflow(name=\"srr_nipype\",base_dir=wf_base_dir)\n",
     "    srr_nipype_dir = os.path.join(wf.base_dir, wf.name )\n",
     "    \n",
     "    \n",
     "    # Initialization\n",
-    "    if os.path.isfile(os.path.join(output_dir,\"pypeline_\"+subject+\".log\")):\n",
-    "        os.unlink(os.path.join(output_dir,\"pypeline_\"+subject+\".log\"))\n",
-    "#         open(os.path.join(output_dir,\"pypeline.log\"), 'a').close()\n",
+    "    if os.path.isfile(os.path.join(process_dir,\"pypeline_\"+subject+\".log\")):\n",
+    "        os.unlink(os.path.join(process_dir,\"pypeline_\"+subject+\".log\"))\n",
+    "#         open(os.path.join(process_dir,\"pypeline.log\"), 'a').close()\n",
     "        \n",
     "\n",
-    "    config.update_config({'logging': {'log_directory': os.path.join(output_dir), 'log_to_file': True},\n",
+    "    config.update_config({'logging': {'log_directory': os.path.join(process_dir), 'log_to_file': True},\n",
     "                          'execution': {\n",
     "                              'remove_unnecessary_outputs': False,\n",
     "                              'stop_on_first_crash': True,\n",
@@ -337,15 +336,14 @@
     "    if session != None:\n",
     "        sub_ses = ''.join([sub_ses, '_', session])\n",
     "    srtkImageReconstruction.inputs.sub_ses = sub_ses\n",
-    "\n",
     "    \n",
     "    srtkTVSuperResolution = Node(interface=reconstruction.MialsrtkTVSuperResolution(), name='srtkTVSuperResolution')  \n",
     "    srtkTVSuperResolution.inputs.bids_dir = bids_dir\n",
     "    srtkTVSuperResolution.inputs.stacksOrder = p_stacksOrder\n",
     "    srtkTVSuperResolution.inputs.sub_ses = sub_ses\n",
-    "    srtkTVSuperResolution.inputs.in_loop = 10\n",
-    "    srtkTVSuperResolution.inputs.in_deltat = 0.01\n",
-    "    srtkTVSuperResolution.inputs.in_lambda = 0.75\n",
+    "    srtkTVSuperResolution.inputs.in_loop = primal_dual_loops\n",
+    "    srtkTVSuperResolution.inputs.in_deltat = deltatTV\n",
+    "    srtkTVSuperResolution.inputs.in_lambda = lambdaTV\n",
     "    \n",
     "    \n",
     "\n",
@@ -360,6 +358,10 @@
     "    srtkMaskImage02 = Node(interface=preprocess.MialsrtkMaskImage(), name='srtkMaskImage02')\n",
     "    srtkMaskImage02.inputs.bids_dir = bids_dir\n",
     "    \n",
+    "    datasink = Node(DataSink(), name='sinker')\n",
+    "    output_dir = os.path.join(\"{}\".format(bids_dir),\"derivatives\",\"mialsrtk-py\")\n",
+    "    datasink.inputs.base_directory = output_dir\n",
+    "    \n",
     "    #\n",
     "    ## Nodes ready - Linking now\n",
     "    \n",
@@ -403,10 +405,10 @@
     "    wf.connect(srtkIntensityStandardization02_nlm, \"output_images\", srtkMaskImage01, \"input_images\")\n",
     "    wf.connect(dg, \"masks\", srtkMaskImage01, \"input_masks\")\n",
     "    \n",
+    "    \n",
     "    wf.connect(srtkMaskImage01, \"output_images\", srtkImageReconstruction, \"input_images\")\n",
     "    wf.connect(dg, \"masks\", srtkImageReconstruction, \"input_masks\")\n",
     "    \n",
-    "    \n",
     "    wf.connect(srtkIntensityStandardization02, \"output_images\", srtkTVSuperResolution, \"input_images\")\n",
     "    wf.connect(srtkImageReconstruction, \"output_transforms\", srtkTVSuperResolution, \"input_transforms\")\n",
     "    wf.connect(dg, \"masks\", srtkTVSuperResolution, \"input_masks\")\n",
@@ -425,31 +427,93 @@
     "    wf.connect(srtkRefineHRMaskByIntersection, \"output_SRmask\", srtkMaskImage02, \"in_mask\")\n",
     "    \n",
     "    \n",
+    "    \n",
+    "    #\n",
+    "    ### - Saving files\n",
+    "    \n",
+    "    \n",
+    "    substitutions = []\n",
+    "    for stack in stacksOrder:\n",
+    "    \n",
+    "        print( sub_ses+'_run-'+str(stack)+'_T2w_nlm_uni_bcorr_histnorm.nii.gz', '    --->     ',sub_ses+'_run-'+str(stack)+'_T2w_preproc.nii.gz')\n",
+    "        substitutions.append( ( sub_ses+'_run-'+str(stack)+'_T2w_nlm_uni_bcorr_histnorm.nii.gz', sub_ses+'_run-'+str(stack)+'_T2w_preproc.nii.gz') )\n",
+    "        \n",
+    "        print( sub_ses+'_run-'+str(stack)+'_T2w_nlm_uni_bcorr_histnorm_transform_'+str(len(stacksOrder))+'V.txt', '    --->     ', sub_ses+'_run-'+str(stack)+'_T2w_from-origin_to-SDI_mode-image_xfm.txt')\n",
+    "        substitutions.append( ( sub_ses+'_run-'+str(stack)+'_T2w_nlm_uni_bcorr_histnorm_transform_'+str(len(stacksOrder))+'V.txt', sub_ses+'_run-'+str(stack)+'_T2w_from-origin_to-SDI_mode-image_xfm.txt') )\n",
+    "        \n",
+    "        print( sub_ses+'_run-'+str(stack)+'_T2w_uni_bcorr_histnorm_LRmask.nii.gz', '    --->     ', sub_ses+'_run-'+str(stack)+'_T2w_desc-LRmask.nii.gz')\n",
+    "        substitutions.append( ( sub_ses+'_run-'+str(stack)+'_T2w_uni_bcorr_histnorm_LRmask.nii.gz', sub_ses+'_run-'+str(stack)+'_T2w_desc-LRmask.nii.gz') )\n",
+    "\n",
+    "        \n",
+    "    print( 'SDI_'+sub_ses+'_'+str(len(stacksOrder))+'V_rad1.nii.gz', '    --->     ', sub_ses+'_rec-SDI_T2w.nii.gz')\n",
+    "    substitutions.append( ( 'SDI_'+sub_ses+'_'+str(len(stacksOrder))+'V_rad1.nii.gz', sub_ses+'_rec-SDI_T2w.nii.gz') )\n",
+    "\n",
+    "    print( 'SRTV_'+sub_ses+'_'+str(len(stacksOrder))+'V_rad1_gbcorr.nii.gz', '    --->     ', sub_ses+'_rec-SR_T2w.nii.gz')\n",
+    "    substitutions.append( ( 'SRTV_'+sub_ses+'_'+str(len(stacksOrder))+'V_rad1_gbcorr.nii.gz', sub_ses+'_rec-SR_T2w.nii.gz') )\n",
+    "    \n",
+    "\n",
+    "    print( sub_ses+'_T2w_uni_bcorr_histnorm_srMask.nii.gz', '    --->     ', sub_ses+'_rec-SR_T2w_desc-brain_mask.nii.gz')\n",
+    "    substitutions.append( ( sub_ses+'_T2w_uni_bcorr_histnorm_srMask.nii.gz', sub_ses+'_rec-SR_T2w_desc-SRmask.nii.gz') )\n",
+    "\n",
+    "    \n",
+    "        \n",
+    "    datasink.inputs.substitutions = substitutions\n",
+    "    \n",
+    "    wf.connect(srtkMaskImage01, \"output_images\", datasink, 'preproc')\n",
+    "    wf.connect(srtkImageReconstruction, \"output_transforms\", datasink, 'xfm')\n",
+    "    wf.connect(srtkRefineHRMaskByIntersection, \"output_LRmasks\", datasink, 'postproc')\n",
+    "    \n",
+    "    wf.connect(srtkImageReconstruction, \"output_sdi\", datasink, 'anat')\n",
+    "    wf.connect(srtkN4BiasFieldCorrection, \"output_image\", datasink, 'anat.@SR')\n",
+    "    wf.connect(srtkRefineHRMaskByIntersection, \"output_SRmask\", datasink, 'postproc.@SRmask')\n",
+    "    \n",
+    "    \n",
+    "    # JSON file SRTV\n",
+    "    output_dict = {}\n",
+    "\n",
+    "    output_dict[\"Description\"] = \"Isotropic high-resolution image reconstructed using the Total-Variation Super-Resolution algorithm provided by MIALSRTK\"\n",
+    "    # output_dict[\"Sources\"] = sources\n",
+    "    output_dict[\"Input sources run order\"] = stacksOrder\n",
+    "    output_dict[\"CustomMetaData\"] = {}\n",
+    "    output_dict[\"CustomMetaData\"][\"Number of scans used\"] = str(len(p_stacksOrder))\n",
+    "    output_dict[\"CustomMetaData\"][\"TV regularization weight lambda\"] = lambdaTV\n",
+    "    output_dict[\"CustomMetaData\"][\"Optimization time step\"] = deltatTV\n",
+    "    output_dict[\"CustomMetaData\"][\"Primal/dual loops\"] = primal_dual_loops\n",
+    "\n",
+    "    output_json = os.path.join(output_dir, 'anat', ''.join([subject, '_rec-SR.json']))\n",
+    "    with open(output_json, 'w+', encoding='utf8') as outfile:\n",
+    "        json.dump(output_dict, outfile, indent=4)\n",
     "        \n",
     "    return wf"
    ]
   },
+  {
+   "cell_type": "raw",
+   "metadata": {},
+   "source": [
+    "%%bash\n",
+    "\n",
+    "\n",
+    "python srr.py /fetaldata /fetaldata/output participant --participant_label \"HK01\" \"ctrl0022\""
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# m_wf = create_workflow(bids_dir, subject, p_stacksOrder=stacsksOrder)\n",
-    "m_wf = create_workflow(bids_dir, subject, p_stacksOrder=stacksOrder, session = session)\n",
+    "m_wf = create_workflow(bids_dir, process_dir='/fetaldata/derivatives/tmp_proc', subject=subject, p_stacksOrder=stacksOrder, session = session)\n",
     "m_wf.write_graph()\n",
-    "aa = m_wf.run()\n",
-    "\n",
-    "\n",
-    "copy_results_to_bids(p_bids_dir= bids_dir, \n",
-    "                     nipype_dir=os.path.join(m_wf.base_dir, m_wf.name), \n",
-    "                     p_wf_base_dir=m_wf.base_dir, \n",
-    "                     p_stacksOrder=stacksOrder, \n",
-    "                     subject=subject, \n",
-    "                     session=session)"
+    "aa = m_wf.run()\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

diff --git a/pymialsrtk/interfaces/postprocess.py b/pymialsrtk/interfaces/postprocess.py
@@ -96,6 +96,8 @@ def _run_interface(self, runtime):
             cmd += ['-O', out_file]
 
         _, name, ext = split_filename(os.path.abspath(self.inputs.input_images[0]))
+        run_id = (name.split('run-')[1]).split('_')[0]
+        name = name.replace('_run-'+run_id+'_', '_')
         out_file = os.path.join(os.getcwd().replace(self.inputs.bids_dir,'/fetaldata'), ''.join((name, self.inputs.out_srmask_postfix, ext)))
 
         cmd += ['-r', self.inputs.input_sr]
@@ -116,7 +118,10 @@ def _run_interface(self, runtime):
     def _list_outputs(self):
         outputs = self._outputs().get()
         _, name, ext = split_filename(os.path.abspath(self.inputs.input_images[0]))
+        run_id = (name.split('run-')[1]).split('_')[0]
+        name = name.replace('_run-'+run_id+'_', '_')
         outputs['output_SRmask'] = os.path.join(os.getcwd().replace(self.inputs.bids_dir,'/fetaldata'), ''.join((name, self.inputs.out_srmask_postfix, ext)))
+
         outputs['output_LRmasks'] = glob(os.path.abspath(''.join(["*", self.inputs.out_LRmask_postfix, ext])))
         return outputs
 

diff --git a/pymialsrtk/interfaces/preprocess.py b/pymialsrtk/interfaces/preprocess.py
@@ -518,7 +518,7 @@ class MialsrtkMaskImageInputSpec(BaseInterfaceInputSpec):
     out_im_postfix = traits.Str("", usedefault=True)
 
 class MialsrtkMaskImageOutputSpec(TraitedSpec):
-    out_im_file = File(desc='Bias field corrected image')
+    out_im_file = File(desc='Masked image')
 
 
 class MialsrtkMaskImage(BaseInterface):