rework metrics

btk/measure.py

@@ -25,7 +25,7 @@
                                     `(n_objects, stamp_size, stamp_size, n_bands)` depending on
                                     convention. The order of this array should correspond to the
                                     order in the returned `catalog`. Where `n_objects` is the
-                                    number of detected objects
+                                    number of detected objects by the algorithm.
     - segmentation (np.ndarray): Array of booleans with shape `(n_objects,stamp_size,stamp_size)`
                                  The pixels set to True in the i-th channel correspond to the i-th
                                  object. The order should correspond to the order in the returned
@@ -218,18 +218,27 @@ def __next__(self):
 
         Returns:
             draw_blend_generator output from `__next__` method.
-            measurement output: List of length `batch_size`, where each element is a list of
+            measurement_results: List of length `batch_size`, where each element is a list of
                                 `len(measure_functions)` corresponding to the measurements made by
                                 each function on each element of the batch.
         """
         blend_output = next(self.draw_blend_generator)
         input_args = ((blend_output, i) for i in range(self.batch_size))
-        measure_results = multiprocess(
+        measure_output = multiprocess(
             self.run_batch,
             input_args,
             self.cpus,
             self.verbose,
         )
         if self.verbose:
             print("Measurement performed on batch")
+        measure_results = {}
+        for i, f in enumerate(self.measure_functions):
+            measure_dic = {}
+            for key in ["catalog", "deblended_images", "segmentation"]:
+                if key in measure_output[0][i] and measure_output[0][i][key] is not None:
+                    measure_dic[key] = [
+                        measure_output[j][i][key] for j in range(len(measure_output))
+                    ]
+            measure_results[f.__name__] = measure_dic
         return blend_output, measure_results