Added possibility to save results

btk/__init__.py

@@ -17,3 +17,4 @@
 from . import plot_utils
 from . import sampling_functions
 from . import survey
+from . import utils

btk/draw_blends.py

@@ -148,6 +148,7 @@ def __init__(
         shifts=None,
         indexes=None,
         channels_last=False,
+        save_path=None,
     ):
         """Initializes the DrawBlendsGenerator class.
 
@@ -170,6 +171,8 @@ def __init__(
             channels_last (bool): Whether to return images as numpy arrays with the channel
                                 (band) dimension as the last dimension or before the pixels
                                 dimensions (default).
+            save_path (str): Path to save the results, ending by the file name root. If left
+                            ton None, results will not be saved.
         """
         self.blend_generator = BlendGenerator(
             catalog, sampling_function, batch_size, shifts, indexes, verbose
@@ -191,8 +194,8 @@ def __init__(
 
         self.add_noise = add_noise
         self.verbose = verbose
-
         self.channels_last = channels_last
+        self.save_path = save_path
 
     def __iter__(self):
         """Returns iterable which is the object itself."""
@@ -274,6 +277,15 @@ def __next__(self):
                 blend_images[s.name][i] = batch_results[i][0]
                 isolated_images[s.name][i] = batch_results[i][1]
                 batch_blend_cat[s.name].append(batch_results[i][2])
+
+            if self.save_path is not None:
+                np.save(f"{self.save_path}_{s.name}_blended", blend_images[s.name])
+                np.save(f"{self.save_path}_{s.name}_isolated", isolated_images[s.name])
+                for i in range(len(batch_results)):
+                    batch_blend_cat[s.name][i].write(
+                        f"{self.save_path}_{s.name}_blend_info_{i}", format="ascii", overwrite=True
+                    )
+
         if len(self.surveys) > 1:
             output = {
                 "blend_images": blend_images,
@@ -291,6 +303,7 @@ def __next__(self):
                 "psf": psfs[survey_name],
                 "wcs": wcss[survey_name],
             }
+
         return output
 
     def render_mini_batch(self, blend_list, psf, wcs, survey, extra_data=None):

btk/measure.py

@@ -151,6 +151,7 @@ def __init__(
         cpus=1,
         verbose=False,
         measure_kwargs: dict = None,
+        save_path=None,
     ):
         """Initialize measurement generator.
 
@@ -163,6 +164,8 @@ def __init__(
             verbose (bool): Whether to print information about measurement.
             measure_kwargs (dict): Dictionary containing keyword arguments to be passed
                 in to each of the `measure_functions`.
+            save_path (str): Path to save the results, ending by the file name root. If left
+                            ton None, results will not be saved.
         """
         # setup and verify measure_functions.
         if callable(measure_functions):
@@ -183,6 +186,7 @@ def __init__(
         self.batch_size = self.draw_blend_generator.batch_size
         self.channels_last = self.draw_blend_generator.channels_last
         self.verbose = verbose
+        self.save_path = save_path
 
         # initialize measure_kwargs dictionary.
         self.measure_kwargs = {} if measure_kwargs is None else measure_kwargs
@@ -268,11 +272,21 @@ def __next__(self):
         measure_results = {}
         for i, f in enumerate(self.measure_functions):
             measure_dic = {}
-            for key in ["catalog", "deblended_images", "segmentation"]:
+            for key in self.measure_params:
                 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 self.save_path is not None:
+                for key in ["segmentation", "deblended_images"]:
+                    if key in measure_dic.keys():
+                        np.save(f"{self.save_path}_{f.__name__}_{key}", measure_dic[key])
+                for j, cat in enumerate(measure_dic["catalog"]):
+                    cat.write(
+                        f"{self.save_path}_{f.__name__}_detection_catalog_{j}",
+                        format="ascii",
+                        overwrite=True,
+                    )
 
         return blend_output, measure_results

btk/metrics.py

@@ -509,6 +509,7 @@ def compute_metrics(  # noqa: C901
     meas_band_num=0,
     target_meas={},
     channels_last=False,
+    save_path=None,
 ):
     """Computes all requested metrics given information in a single batch from measure_generator.
 
@@ -542,6 +543,8 @@ def compute_metrics(  # noqa: C901
                              be returned for both isolated and deblended images to compare.
         channels_last (bool) : Indicates whether the images should be channels first (NCHW)
                           or channels last (NHWC).
+        save_path (str): Path to save the results, ending by the file name root. If left
+                            ton None, results will not be saved.
 
     Returns:
         results (dict) : Contains all the computed metrics. Entries are :
@@ -626,6 +629,10 @@ def compute_metrics(  # noqa: C901
                 for k in reconstruction_keys:
                     row[k] = results["reconstruction"][k][i][j]
             results["galaxy_summary"].add_row(row[0])
+    if save_path is not None:
+        for key in use_metrics:
+            np.save(f"{save_path}_{key}_metric", results[key])
+        results["galaxy_summary"].write(f"{save_path}_galaxy_summary", format="ascii")
 
     return results
 
@@ -640,6 +647,7 @@ def __init__(
         meas_band_num=0,
         target_meas={},
         noise_threshold_factor=3,
+        save_path=None,
     ):
         """Initialize metrics generator.
 
@@ -658,12 +666,15 @@ def __init__(
                 applied when getting segmentations from true images. A value of 3 would
                 correspond to a threshold of 3 sigmas (with sigma the standard deviation of
                 the noise)
+            save_path (str): Path to save the results, ending by the file name root. If left
+                            ton None, results will not be saved.
         """
         self.measure_generator: MeasureGenerator = measure_generator
         self.use_metrics = use_metrics
         self.meas_band_num = meas_band_num
         self.target_meas = target_meas
         self.noise_threshold_factor = noise_threshold_factor
+        self.save_path = save_path
 
     def __next__(self):
         """Returns metric results calculated on one batch."""
@@ -695,6 +706,7 @@ def __next__(self):
                 self.meas_band_num,
                 target_meas,
                 channels_last=self.measure_generator.channels_last,
+                save_path=self.save_path + "_" + meas_func if self.save_path is not None else None,
             )
             metrics_results[meas_func] = metrics_results_f
 

btk/utils.py

@@ -0,0 +1,124 @@
+"""Contains utility functions, including functions for loading saved results."""
+import numpy as np
+from astropy.table import Table
+
+
+def load_blend_results(path, survey):
+    """Load results exported from a DrawBlendsGenerator.
+
+    Args;
+        path (str): Path to the files. Should be the same as the save_path
+                    which was provided to the DrawBlendsGenerator to save
+                    the files.
+        survey (str): Name of the survey for which you want to load the files.
+
+    Returns:
+        Dictionnary containing the blend images, the isolated images and the
+        informations about the blends.
+    """
+    blend_images = np.load(f"{path}_{survey}_blended.npy", allow_pickle=True)
+    isolated_images = np.load(f"{path}_{survey}_isolated.npy", allow_pickle=True)
+    blend_list = []
+    for i in range(blend_images.shape[0]):
+        blend_list.append(Table.read(f"{path}_{survey}_blend_info_{i}", format="ascii"))
+    return {
+        "blend_images": blend_images,
+        "isolated_images": isolated_images,
+        "blend_list": blend_list,
+    }
+
+
+def load_measure_results(path, measure_name, n_batch):
+    """Load results exported from a MeasureGenerator.
+
+    Args:
+        path (str): Path to the files. Should be the same as the save_path
+                    which was provided to the MeasureGenerator to save
+                    the files.
+        measure_name (str): Name of the measure function for which you
+                    want to load the files
+        n_batch (int): Number of blends in the batch you want to load
+
+    Returns:
+        Dictionnary containing the detection catalogs, the segmentations
+        and the deblended images.
+    """
+    measure_results = {}
+    for key in ["segmentation", "deblended_images"]:
+        try:
+            measure_results[key] = np.load(f"{path}_{measure_name}_{key}.npy", allow_pickle=True)
+        except FileNotFoundError:
+            print(f"No {key} found.")
+    catalog = []
+    for j in range(n_batch):
+        catalog.append(
+            Table.read(
+                f"{path}_{measure_name}_detection_catalog_{j}",
+                format="ascii",
+            )
+        )
+    measure_results["catalog"] = catalog
+    return measure_results
+
+
+def load_metrics_results(path, measure_name):
+    """Load results exported from a MetricsGenerator.
+
+    Args:
+        path (str): Path to the files. Should be the same as the save_path
+                    which was provided to the MetricsGenerator to save
+                    the files.
+        measure_name (str): Name of the measure function for which you
+                    want to load the files
+
+    Returns:
+        Dictionnary containing the detection catalogs, the segmentations
+        and the deblended images.
+    """
+    metrics_results = {}
+    for key in ["detection", "segmentation", "reconstruction"]:
+        try:
+            metrics_results[key] = np.load(
+                f"{path}_{measure_name}_{key}_metric.npy", allow_pickle=True
+            )
+        except FileNotFoundError:
+            print(f"No {key} metrics found.")
+
+    metrics_results["galaxy_summary"] = Table.read(
+        f"{path}_{measure_name}_galaxy_summary",
+        format="ascii",
+    )
+    return metrics_results
+
+
+def load_all_results(path, surveys, measure_names, n_batch):
+    """Load results exported from a MetricsGenerator.
+
+    Args:
+        path (str): Path to the files. Should be the same as the save_path
+                    which was provided to the MetricsGenerator to save
+                    the files.
+        surveys (list): Names of the surveys for which you want to load
+                        the files
+        measure_names (list): Names of the measure functions for which you
+                    want to load the files
+        n_batch (int): Number of blends in the batch you want to load
+
+    Returns:
+        The three dictionnaries corresponding to the results.
+    """
+    blend_results = {}
+    for key in ["blend_images", "isolated_images", "blend_list"]:
+        blend_results[key] = {}
+    measure_results = {}
+    metrics_results = {}
+    for s in surveys:
+        blend_results_temp = load_blend_results(path, s)
+        for key in ["blend_images", "isolated_images", "blend_list"]:
+            blend_results[key][s] = blend_results_temp[key]
+
+    for meas in measure_names:
+        measure_results[meas] = load_measure_results(path, meas, n_batch)
+        metrics_results[meas] = load_metrics_results(path, meas)
+
+    return blend_results, measure_results, metrics_results

tests/test_save.py

@@ -0,0 +1,57 @@
+import os
+import shutil
+
+import numpy as np
+
+import btk
+from btk.survey import Rubin
+
+
+def test_save():
+    try:
+        shutil.rmtree("./output")
+    except Exception:
+        pass
+    os.mkdir("./output")
+    catalog_name = "data/sample_input_catalog.fits"
+    stamp_size = 24.0
+    batch_size = 8
+    catalog = btk.catalog.CatsimCatalog.from_file(catalog_name)
+    sampling_function = btk.sampling_functions.DefaultSampling(stamp_size=stamp_size)
+    draw_blend_generator = btk.draw_blends.CatsimGenerator(
+        catalog,
+        sampling_function,
+        [Rubin],
+        batch_size=batch_size,
+        stamp_size=stamp_size,
+        save_path="./output/test",
+    )
+    meas_generator = btk.measure.MeasureGenerator(
+        btk.measure.sep_measure, draw_blend_generator, save_path="./output/test"
+    )
+    metrics_generator = btk.metrics.MetricsGenerator(
+        meas_generator,
+        use_metrics=("detection", "segmentation", "reconstruction"),
+        target_meas={"ellipticity": btk.metrics.meas_ksb_ellipticity},
+        save_path="./output/test",
+    )
+    blend_results, measure_results, metrics_results = next(metrics_generator)
+    blend_results2, measure_results2, metrics_results2 = btk.utils.load_all_results(
+        "./output/test", ["LSST"], ["sep_measure"], batch_size
+    )
+    np.testing.assert_array_equal(
+        blend_results["blend_images"], blend_results2["blend_images"]["LSST"]
+    )
+    np.testing.assert_array_equal(
+        measure_results["sep_measure"]["segmentation"][0],
+        measure_results2["sep_measure"]["segmentation"][0],
+    )
+    np.testing.assert_array_equal(
+        measure_results["sep_measure"]["deblended_images"][0],
+        measure_results2["sep_measure"]["deblended_images"][0],
+    )
+    np.testing.assert_array_equal(
+        metrics_results["sep_measure"]["galaxy_summary"]["distance_closest_galaxy"],
+        metrics_results2["sep_measure"]["galaxy_summary"]["distance_closest_galaxy"],
+    )
+    shutil.rmtree("./output")