rework metrics

btk/draw_blends.py

@@ -194,7 +194,7 @@ def __init__(
 
         if dim_order not in ("NCHW", "NHWC"):
             raise ValueError("dim_order must be either 'NCHW' or 'NHWC'.")
-        self.dim_order = (0, 1, 2) if dim_order == "NCHW" else (1, 2, 0)
+        self.dim_order = dim_order
 
     def __iter__(self):
         """Returns iterable which is the object itself."""
@@ -245,9 +245,17 @@ def __next__(self):
             batch_blend_cat[s.name] = []
             batch_obs_cond[s.name] = []
             image_shape = (
-                len(s.filters),
-                pix_stamp_size,
-                pix_stamp_size,
+                (
+                    len(s.filters),
+                    pix_stamp_size,
+                    pix_stamp_size,
+                )
+                if self.dim_order == "NCHW"
+                else (
+                    pix_stamp_size,
+                    pix_stamp_size,
+                    len(s.filters),
+                )
             )
             blend_images[s.name] = np.zeros((self.batch_size, *image_shape))
             isolated_images[s.name] = np.zeros((self.batch_size, self.max_number, *image_shape))
@@ -355,7 +363,7 @@ def render_mini_batch(self, blend_list, psf, wcs, survey, extra_data=None):
                 iso_image_multi[:, b, :, :] = single_band_output[1]
 
             # transpose if requested.
-            dim_order = np.array(self.dim_order)
+            dim_order = np.array((0, 1, 2) if self.dim_order == "NCHW" else (1, 2, 0))
             blend_image_multi = blend_image_multi.transpose(dim_order)
             iso_image_multi = iso_image_multi.transpose(0, *(dim_order + 1))
 

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
@@ -89,7 +89,7 @@ def sep_measure(batch, idx):
 
     image = batch["blend_images"][idx]
     stamp_size = image.shape[-2]  # true for both 'NCHW' or 'NHWC' formats.
-    coadd = np.mean(image, axis=0)
+    coadd = np.mean(image, axis=np.argmin(image.shape))  # Smallest dimension is the channels
     bkg = sep.Background(coadd)
     # Here the 1.5 value corresponds to a 1.5 sigma threshold for detection against noise.
     catalog, segmentation = sep.extract(coadd, 1.5, err=bkg.globalrms, segmentation_map=True)
@@ -99,7 +99,11 @@ def sep_measure(batch, idx):
     for i in range(n_objects):
         seg_i = segmentation == i + 1
         segmentation_exp[i] = seg_i
-        deblended_images[i] = image * seg_i[np.newaxis, ...]
+        seg_i_reshaped = np.zeros((np.min(image.shape), stamp_size, stamp_size))
+        for j in range(np.min(image.shape)):
+            seg_i_reshaped[j] = seg_i
+        seg_i_reshaped = np.moveaxis(seg_i_reshaped, 0, np.argmin(image.shape))
+        deblended_images[i] = image * seg_i_reshaped
 
     t = astropy.table.Table()
     t["x_peak"] = catalog["x"]
@@ -203,11 +207,13 @@ def run_batch(self, batch, index):
                             f"The output '{key}' of at least one of your measurement"
                             f"functions is not a numpy array."
                         )
-                    if not out[key].shape[-2:] == batch["blend_images"].shape[-2:]:
-                        raise ValueError(
-                            f"The shapes of the blended images in your {key} don't"
-                            f"match for at least one your measurement functions."
-                        )
+                    if key == "deblended_images":
+                        if not out[key].shape[-3:] == batch["blend_images"].shape[-3:]:
+                            raise ValueError(
+                                f"The shapes of the blended images in your {key} don't "
+                                f"match for at least one your measurement functions."
+                                f"{out[key].shape[-3:]} vs {batch['blend_images'].shape[-3:]}"
+                            )
 
             out = {k: out.get(k, None) for k in self.measure_params}
             output.append(out)
@@ -217,19 +223,32 @@ def __next__(self):
         """Return measurement results on a single batch from the draw_blend_generator.
 
         Returns:
-            draw_blend_generator output from `__next__` method.
-            measurement output: 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.
+            draw_blend_generator output from its `__next__` method.
+            measurement_results (dict): Dictionary with keys being the name of each
+                `measure_function` passed in. Each value is a dictionary containing keys
+                `catalog`, `deblended_images`, and `segmentation` storing the values returned by
+                the corresponding measure_function` for one 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 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
+        if len(self.measure_functions) == 1:
+            measure_results = measure_results[self.measure_functions[0].__name__]
+
         return blend_output, measure_results