Merge pull request #487 from frheault/add_merge_option

Add options to merge groups of labels

scripts/scil_combine_labels.py

@@ -48,18 +48,19 @@ def _build_arg_parser():
     o_ids = p.add_mutually_exclusive_group()
     o_ids.add_argument('--out_labels_ids', type=int, nargs='+', default=[],
                        help='Give a list of labels indices for output images.')
-
     o_ids.add_argument('--unique', action='store_true',
                        help='Output id with unique labels,'
                             ' excluding first background value.')
-
     o_ids.add_argument('--group_in_m', action='store_true',
                        help='Add (x*1000000) to each volume labels,'
                             ' where x is the input volume order number.')
 
     p.add_argument('--background', type=int, default=0,
                    help='Background id, excluded from output [%(default)s],\n'
                         ' the value is used as output background value.')
+    p.add_argument('--merge_groups', action='store_true',
+                   help='Each group for the -v option will be merge as a'
+                        'single labels.')
     add_overwrite_arg(p)
     return p
 
@@ -74,7 +75,7 @@ def main():
     used_indices_all = False
     for v_args in args.volume_ids:
         if len(v_args) < 2:
-            logging.error("No indices was given for a given volume")
+            parser.error("No indices was given for a given volume.")
 
         image_files.append(v_args[0])
         if "all" in v_args:
@@ -84,7 +85,11 @@ def main():
             indices_per_volume.append(np.asarray(v_args[1:], dtype=int))
 
     if used_indices_all and args.out_labels_ids:
-        logging.error("'all' indices cannot be used with 'out_labels_ids'")
+        parser.error("'all' indices cannot be used with --out_labels_ids.")
+
+    if args.merge_groups and (args.group_in_m or args.unique):
+        parser.error("Cannot use --unique and --group_in_m with "
+                     "--merge_groups.")
 
     # Check inputs / output
     assert_inputs_exist(parser, image_files)
@@ -110,12 +115,17 @@ def main():
         id_list = np.asarray(id_list)
         new_ids = id_list[~np.in1d(id_list, args.background)]
         filtered_ids_per_vol.append(new_ids)
+
     # Prepare output indices
     if args.out_labels_ids:
         out_labels = args.out_labels_ids
-        if len(out_labels) != len(np.hstack(indices_per_volume)):
-            logging.error("--out_labels_ids, requires the same amount"
-                          " of total given input indices")
+        if not args.merge_groups \
+                and len(out_labels) != len(np.hstack(indices_per_volume)):
+            parser.error("--out_labels_ids, requires the same amount"
+                         " of total given input indices.")
+        elif len(out_labels) != len(args.volume_ids):
+            parser.error("--out_labels_ids, requires the same amount"
+                         " of total given groups (to merge).")
     elif args.unique:
         stack = np.hstack(filtered_ids_per_vol)
         ids = np.arange(len(stack) + 1)
@@ -127,7 +137,10 @@ def main():
             m_list.append(prefix + np.asarray(filtered_ids_per_vol[i]))
         out_labels = np.hstack(m_list)
     else:
-        out_labels = np.hstack(filtered_ids_per_vol)
+        if args.merge_groups:
+            out_labels = np.arange(len(args.volume_ids))+1
+        else:
+            out_labels = np.hstack(filtered_ids_per_vol)
 
     if len(np.unique(out_labels)) != len(out_labels):
         logging.error("The same output label number was used "
@@ -138,14 +151,22 @@ def main():
     resulting_labels = (np.ones_like(data_list[0], dtype=np.uint16)
                         * args.background)
     for i in range(len(image_files)):
-        # Add Given labels for each volume
+        # Add given labels for each volume
         for index in filtered_ids_per_vol[i]:
-            mask = data_list[i] == index
-            resulting_labels[mask] = out_labels[current_id]
-            current_id += 1
-
-            if np.count_nonzero(mask) == 0:
-                logging.warning("Label {} was not in the volume".format(index))
+            if args.merge_groups:
+                for j, curr_volume_ids in enumerate(args.volume_ids):
+                    if str(index) in curr_volume_ids[1:]:
+                        where_at = j
+                mask = data_list[i] == index
+                resulting_labels[mask] = out_labels[where_at]
+            else:
+                mask = data_list[i] == index
+                resulting_labels[mask] = out_labels[current_id]
+                current_id += 1
+
+                if np.count_nonzero(mask) == 0:
+                    logging.warning(
+                        "Label {} was not in the volume".format(index))
 
     # Save final combined volume
     nib.save(nib.Nifti1Image(resulting_labels, first_img.affine,

scripts/tests/test_combine_labels.py

@@ -4,7 +4,7 @@
 import os
 import tempfile
 
-from scilpy.io.fetcher import fetch_data, get_home, get_testing_files_dict 
+from scilpy.io.fetcher import fetch_data, get_home, get_testing_files_dict
 
 # If they already exist, this only takes 5 seconds (check md5sum)
 fetch_data(get_testing_files_dict(), keys=['atlas.zip'])
@@ -27,3 +27,16 @@ def test_execution_atlas(script_runner):
                             '253', '254', '1022', '1024', '2022', '2024',
                             '-v', in_brainstem, '16')
     assert ret.success
+
+
+def test_execution_atlas_merge(script_runner):
+    os.chdir(os.path.expanduser(tmp_dir.name))
+    in_atlas_1 = os.path.join(get_home(), 'atlas',
+                              'atlas_freesurfer_v2.nii.gz')
+    in_brainstem = os.path.join(get_home(), 'atlas', 'brainstem.nii.gz')
+    ret = script_runner.run('scil_combine_labels.py',
+                            'atlas_freesurfer_v2_merge_brainstem.nii.gz',
+                            '-v', in_atlas_1, '8', '47', '251', '252',
+                            '253', '254', '1022', '1024', '2022', '2024',
+                            '-v', in_brainstem, '16', '--merge_groups')
+    assert ret.success