From f64fd2c1c592828b3ae91c43601480f90a4af911 Mon Sep 17 00:00:00 2001
From: David Meunier <david_meunier_79@hotmail.fr>
Date: Wed, 25 Sep 2024 13:29:53 +0200
Subject: [PATCH] Bug nopriors (#264)

* if "use_priors" in params["segment_atropos_pipe"].keys():

* post_pat

* post_pat
---
 macapype/pipelines/full_pipelines.py | 19 +++++++-----
 macapype/pipelines/segment.py        | 44 +++++++++++++++-------------
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/macapype/pipelines/full_pipelines.py b/macapype/pipelines/full_pipelines.py
index 090f8d4d6..ef5613910 100644
--- a/macapype/pipelines/full_pipelines.py
+++ b/macapype/pipelines/full_pipelines.py
@@ -638,14 +638,17 @@ def create_brain_segment_from_mask_pipe(
         segment_atropos_pipe = create_segment_atropos_seg_pipe(
             params=parse_key(params, "segment_atropos_pipe"))
 
-        if "register_NMT_pipe" in params:
-            brain_segment_pipe.connect(
-                register_NMT_pipe, 'outputnode.native_template_seg',
-                segment_atropos_pipe, "inputnode.seg_file")
-        elif "reg" in params:
-            brain_segment_pipe.connect(
-                register_seg_to_nat, 'out_file',
-                segment_atropos_pipe, "inputnode.seg_file")
+        # linking priors if "use_priors" in params
+        if "use_priors" in params["segment_atropos_pipe"].keys():
+
+            if "register_NMT_pipe" in params:
+                brain_segment_pipe.connect(
+                    register_NMT_pipe, 'outputnode.native_template_seg',
+                    segment_atropos_pipe, "inputnode.seg_file")
+            elif "reg" in params:
+                brain_segment_pipe.connect(
+                    register_seg_to_nat, 'out_file',
+                    segment_atropos_pipe, "inputnode.seg_file")
     else:
         print("#### create_segment_atropos_pipe (3 tissues) ")
 
diff --git a/macapype/pipelines/segment.py b/macapype/pipelines/segment.py
index b2aa106fd..652cbf868 100644
--- a/macapype/pipelines/segment.py
+++ b/macapype/pipelines/segment.py
@@ -172,6 +172,12 @@ def create_segment_atropos_seg_pipe(params={}, name="segment_atropos_pipe"):
 
     # merging probseg
     for tissue, index_tissue in tissue_dict.items():
+
+        if "use_priors" in params.keys():
+            post_pat = "SegmentationPosteriors{:02d}"
+        else:
+            post_pat = "SegmentationPosteriors{}"
+
         if isinstance(index_tissue, list):
 
             # Merging as file list
@@ -182,7 +188,7 @@ def create_segment_atropos_seg_pipe(params={}, name="segment_atropos_pipe"):
                 segment_pipe.connect(
                     seg_at,
                     ('segmented_files', get_pattern,
-                     "SegmentationPosteriors{:02d}".format(
+                     post_pat.format(
                          int(sub_index_tissue))),
                     merge_list, 'in' + str(index+1))
 
@@ -207,7 +213,7 @@ def create_segment_atropos_seg_pipe(params={}, name="segment_atropos_pipe"):
             segment_pipe.connect(
                 seg_at,
                 ('segmented_files', get_pattern,
-                 "SegmentationPosteriors{:02d}".format(int(index_tissue))),
+                 post_pat.format(int(index_tissue))),
                 outputnode, 'prob_'+tissue)
 
     return segment_pipe
@@ -444,6 +450,12 @@ def create_segment_atropos_pipe(params={}, name="segment_atropos_pipe"):
 
     # merging probseg
     for tissue, index_tissue in tissue_dict.items():
+
+        if "use_priors" in params.keys():
+            post_pat = "SegmentationPosteriors{:02d}"
+        else:
+            post_pat = "SegmentationPosteriors{}"
+
         if isinstance(index_tissue, list):
 
             # Merging as file list
@@ -453,9 +465,9 @@ def create_segment_atropos_pipe(params={}, name="segment_atropos_pipe"):
             for index, sub_index_tissue in enumerate(index_tissue):
                 segment_pipe.connect(
                     seg_at,
-                    ('segmented_files', get_pattern,
-                     "SegmentationPosteriors{:02d}".format(
-                         int(sub_index_tissue))),
+                    ('segmented_files',
+                     get_pattern,
+                     post_pat.format(int(sub_index_tissue))),
                     merge_list, 'in' + str(index+1))
 
             # Merging files in the same nifti
@@ -475,22 +487,12 @@ def create_segment_atropos_pipe(params={}, name="segment_atropos_pipe"):
 
         else:
 
-            if "use_priors" in params.keys():
-
-                # prob output
-                segment_pipe.connect(
-                    seg_at,
-                    ('segmented_files', get_pattern,
-                     "SegmentationPosteriors{:02d}".format(int(index_tissue))),
-                    outputnode, 'prob_'+tissue)
-            else:
-
-                # prob output
-                segment_pipe.connect(
-                    seg_at,
-                    ('segmented_files', get_pattern,
-                     "SegmentationPosteriors{}".format(int(index_tissue))),
-                    outputnode, 'prob_'+tissue)
+            # prob output
+            segment_pipe.connect(
+                seg_at,
+                ('segmented_files', get_pattern,
+                    post_pat.format(int(index_tissue))),
+                outputnode, 'prob_'+tissue)
 
     return segment_pipe