get_0_180_radios: add return_indices

nabu/resources/dataset_analyzer.py

@@ -407,7 +407,7 @@ def analyze_dataset(dataset_path, processes_file=None, extra_options=None, logge
 
 
 
-def get_0_180_radios(dataset_info, angles=None):
+def get_0_180_radios(dataset_info, angles=None, return_indices=False):
     """
     Get the radios at 0 degres and 180 degrees.
 
@@ -417,6 +417,14 @@ def get_0_180_radios(dataset_info, angles=None):
         Data structure with the dataset information
     angles: array, optional
         Array with the rotation angles. If provided, it overwrites the information from 'dataset_info', if any.
+    return_indices: bool, optional
+        Whether to return radios indices along with the radios array.
+
+    Returns
+    -------
+    res: array or tuple
+        If return_indices is True, return a tuple (radios, indices).
+        Otherwise, return an array with the radios.
     """
 
     if angles is None:
@@ -436,7 +444,10 @@ def get_0_180_radios(dataset_info, angles=None):
     for i in range(n_radios):
         radio_idx = radios_indices[i]
         radios[i] = get_data(dataset_info.projections[radio_idx]).astype("f")
-    return radios
+    if return_indices:
+        return radios, radios_indices
+    else:
+        return radios
 
 
 def get_angles(dataset_info, logger=None):