write comments and improve measure computation

mu_map/eval/measures.py

 import numpy as np
+import pandas as pd
 
+from mu_map.dataset.default import MuMapDataset
+from mu_map.models.unet import UNet
 
-def mse(prediction: np.array, target: np.array):
+
+def mse(prediction: np.array, target: np.array) -> float:
+    """
+    Compute the mean squared error (MSE) between a prediction and
+    a target array.
+
+    Parameters
+    ----------
+    prediction: np.ndarray
+    target: np.ndarray
+    """
     se = (prediction - target) ** 2
     mse = se.sum() / se.size
     return mse
 
 
-def nmae(prediction: np.array, target: np.array, vmax:float=None, vmin:float=None):
-    if vmax is None:
-        vmax = target.max()
-    if vmin is None:
-        vmin = target.min()
+def nmae(
+    prediction: np.array, target: np.array, vmax: float = None, vmin: float = None
+):
+    """
+    Compute the normalized mean absolute error (NMAE) between a prediction
+    and a target array.
+
+    Parameters
+    ----------
+    prediction: np.ndarray
+    target: np.ndarray
+    vmax: float, optional
+        maximum value for normalization, defaults to the maximal value in the target
+    vmin: float, optional
+        minimum value for normalization, defaults to the minimal value in the target
+    """
+    vmax = target.max() if vmax is None else vmax
+    vmin = target.min() if vmin is None else vmin
 
     ae = np.absolute(prediction - target)
     mae = ae.sum() / ae.size
@@ -19,16 +45,51 @@ def nmae(prediction: np.array, target: np.array, vmax:float=None, vmin:float=Non
     return nmae
 
 
+def compute_measures(dataset: MuMapDataset, model: UNet) -> pd.DataFrame:
+    """
+    Compute measures (MSE, NMAE) for all images in a dataset.
+
+    Parameters
+    ----------
+    dataset: MuMapDataset
+        the dataset containing the reconstructions and mu maps for which the scores are computed
+    model: UNet
+        the UNet model which is used to predict mu maps from reconstructions
+
+    Returns
+    -------
+    pd.DataFrame
+        a dataframe containing containing the measures for each image in the dataset
+    """
+    measures = {"NMAE": nmae, "MSE": mse}
+    values = pd.DataFrame(dict(map(lambda x: (x, []), measures.keys())))
+    for i, (recon, mu_map) in enumerate(dataset):
+        _id = dataset.table.iloc[i]["id"]
+        print(
+            f"Process input {str(i):>{len(str(len(dataset)))}}/{len(dataset)}", end="\r"
+        )
+        prediction = model(recon.unsqueeze(dim=0).to(device))
+
+        prediction = prediction.squeeze().cpu().numpy()
+        mu_map = mu_map.squeeze().cpu().numpy()
+
+        row = dict(
+            map(lambda item: (item[0], [item[1](prediction, mu_map)]), measures.items())
+        )
+        row["id"] = _id
+        row = pd.DataFrame(row)
+        values = pd.concat((values, row), ignore_index=True)
+    print(f" " * 100, end="\r")
+    return values
+
+
 if __name__ == "__main__":
     import argparse
 
-    import pandas as pd
     import torch
 
-    from mu_map.dataset.default import MuMapDataset
     from mu_map.dataset.normalization import norm_by_str, norm_choices
     from mu_map.dataset.transform import SequenceTransform, PadCropTranform
-    from mu_map.models.unet import UNet
 
     parser = argparse.ArgumentParser(
         description="Compute, print and store measures for a given model",
@@ -37,7 +98,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--device",
         type=str,
-        default="cpu",
+        default="cuda" if torch.cuda.is_available() else "cpu",
         choices=["cpu", "cuda"],
         help="the device on which the model is evaluated (cpu or cuda)",
     )
@@ -84,7 +145,7 @@ if __name__ == "__main__":
     torch.set_grad_enabled(False)
 
     device = torch.device(args.device)
-    model = UNet()
+    model = UNet(features=[32, 64, 128, 256, 512])
     model.load_state_dict(torch.load(args.weights, map_location=device))
     model = model.to(device).eval()
 
@@ -101,28 +162,16 @@ if __name__ == "__main__":
         scatter_correction=args.scatter_corrected,
     )
 
-    measures = {"NMAE": nmae, "MSE": mse}
-    values = pd.DataFrame(dict(map(lambda x: (x, []), measures.keys())))
-    for i, (recon, mu_map) in enumerate(dataset):
-        print(
-            f"Process input {str(i):>{len(str(len(dataset)))}}/{len(dataset)}", end="\r"
-        )
-        prediction = model(recon.unsqueeze(dim=0).to(device))
-
-        prediction = prediction.squeeze().cpu().numpy()
-        mu_map = mu_map.squeeze().cpu().numpy()
-
-        row = pd.DataFrame(dict(
-            map(lambda item: (item[0], [item[1](prediction, mu_map)]), measures.items())
-        ))
-        values = pd.concat((values, row), ignore_index=True)
-    print(f" " * 100, end="\r")
+    values = compute_measures(dataset, model)
 
     if args.out:
         values.to_csv(args.out, index=False)
 
     print("Scores:")
     for measure_name, measure_values in values.items():
+        if measure_name == "id":
+            continue
+
         mean = measure_values.mean()
         std = np.std(measure_values)
         print(f" - {measure_name:>6}: {mean:.6f}±{std:.6f}")