Fix figure issues

README.md

@@ -216,8 +216,8 @@ Options:
                                   per run).
   -similarity, --vis_sim          Visualise train-val model similarity matrix.
   -va, --vis_all                  Visualise all above.
-  -vf, --vis_format TEXT          The format of saved images. Options: png ,
-                                  pdf
+  -vf, --vis_format TEXT          The format of saved images. Options: png,
+                                  pdf, svg
   -fev, --freq_eval INTEGER       Frequency at which to evaluate test set.
   -fs, --freq_sta INTEGER         Frequency at which to save state
   -fac, --freq_acc INTEGER        Frequency at which to visualise confusion
@@ -246,8 +246,6 @@ Options:
                                   and the maximum to one
   -res, --rescale INTEGER         Rescale images to given value (tuple, one
                                   value per dim).
-  -tb, --tensorboard              Log metrics and figures to tensorboard
-                                  during training
   -st, --strategy TEXT            Define the strategy for distributed
                                   training. Options are: 'ddp', 'deepspeed' or
                                   'fsdp

avae/config.py

@@ -177,9 +177,18 @@ class AffinityConfig(pydantic.BaseModel):
     #### Visualisation parameters
     # setup params
     vis_format: None | str = pydantic.Field(
-        "png", description="The format of saved images. Options: png , pdf"
+        "png",
+        description="The format of saved images. Options: png, pdf, svg "
+        "(vector formats stay sharp at any size and are fastest for "
+        "large matrices)",
+    )
+    vis_print: bool = pydantic.Field(
+        False,
+        description="Render matrix visualisations (confusion, affinity, "
+        "similarity) at print quality (300dpi) instead of the faster "
+        "default resolution. The dpi setting does not apply to vector "
+        "vis_format (pdf, svg)",
     )
-    tensorboard: bool = pydantic.Field(False, description="Use tensorboard")
     vis_his: bool = pydantic.Field(
         False, description="Visualise class distribution histogram"
     )
@@ -328,6 +337,8 @@ def load_config_params(
         # if no config file is provided, start from default and update with command line arguments
         params = {}
 
+    command_line_params = set()
+
     # check for command line input values and overwrite config file values
     # we're using sys args because click has populated defaults from confing
     if sys_args is not None:
@@ -337,6 +348,7 @@ def load_config_params(
                 continue
             name = arg[2:].split("=")[0]
             if name in local_args.keys():
+                command_line_params.add(name)
                 # overwrite config file value with command line argument value
                 # but only if its on system args (click has defaults)
                 if name in params.keys():
@@ -372,11 +384,16 @@ def load_config_params(
         if type(val) == pathlib.Path:
             # turn relative paths to absolute
             params[key] = str(val.absolute())
-        if 'vis' in key and params['vis_all'] is not None:
+        if (
+            'vis' in key
+            and params['vis_all'] is not None
+            and key not in command_line_params
+        ):
             # set visualisation to vis_all if it is not set, except for vis_z_n_int and vis_pose_class
             if key in [
                 'vis_all',
                 'vis_format',
+                'vis_print',
                 'vis_z_n_int',
                 'vis_pose_class',
             ]:
@@ -385,7 +402,11 @@ def load_config_params(
             logging.warning(
                 f"Visualisation parameter 'vis_all' is overriding {key} to {params['vis_all']}"
             )
-        if 'freq' in key and params['freq_all'] is not None:
+        if (
+            'freq' in key
+            and params['freq_all'] is not None
+            and key not in command_line_params
+        ):
             # set frequency to freq_all if it is not set
             if key in ['freq_all']:
                 continue

avae/data.py

@@ -34,6 +34,7 @@ def load_data(
     vis_his: bool = False,
     vis_aff: bool = False,
     vis_format: str = "png",
+    vis_print: bool = False,
 ) -> tuple[torch.utils.data.DataLoader, int]:
     ...
 
@@ -57,6 +58,7 @@ def load_data(
     vis_his: bool = False,
     vis_aff: bool = False,
     vis_format: str = "png",
+    vis_print: bool = False,
 ) -> tuple[
     torch.utils.data.DataLoader,
     torch.utils.data.DataLoader,
@@ -85,13 +87,17 @@ def load_data(
     vis_his: bool = False,
     vis_aff: bool = False,
     vis_format: str = "png",
-) -> tuple[
-    torch.utils.data.DataLoader,
-    torch.utils.data.DataLoader,
-    torch.utils.data.DataLoader,
-    pd.DataFrame,
-    int,
-] | tuple[torch.utils.data.DataLoader, int]:
+    vis_print: bool = False,
+) -> (
+    tuple[
+        torch.utils.data.DataLoader,
+        torch.utils.data.DataLoader,
+        torch.utils.data.DataLoader,
+        pd.DataFrame,
+        int,
+    ]
+    | tuple[torch.utils.data.DataLoader, int]
+):
     """This function a wrapper around the DiskDataLoader class from the caked library. It loads data from a given path, selects a subset of classes if requested, splits it into train / val and test in batch sets, and loads an affinity matrix.
     The function is overloaded to return different types of data depending on the value of the eval parameter. If eval is True, the function returns only the test data and the dimension of the data. If eval is False, the function returns train, validation, and test data, the affinity matrix, and the dimension of the data.
 
@@ -140,7 +146,7 @@ def load_data(
 
     logging.info("\n")
     logging.info("############################################### DATA")
-    logging.info(f"Loading data...")
+    logging.info("Loading data...\n")
 
     # read the class list, if not provided all classes in the dataset will be used as default
     if classes is not None:
@@ -188,6 +194,7 @@ def load_data(
                 classes_list,
                 vis_aff=vis_aff and fabric.global_rank == 0,
                 vis_format=vis_format,
+                vis_print=vis_print,
             )
 
         # assign the affinity matrix to the dataset (small modification from the caked DiskDataset, which only returns data and labels, and we need the affinity matrix indexes for training).
@@ -202,11 +209,19 @@ def load_data(
 
         # Plot the complete splits once, before Fabric shards the loaders.
         if vis_his and fabric.global_rank == 0:
-            train_y = list(sum([y[1] for _, y in enumerate(trains)], ()))
-            val_y = list(sum([y[1] for _, y in enumerate(vals)], ()))
-            plot_classes_distribution(train_y, "train", vis_format=vis_format)
+            train_y = _labels(trains.dataset)
+            val_y = _labels(vals.dataset)
+            plot_classes_distribution(
+                train_y,
+                "train",
+                vis_format=vis_format,
+                vis_print=vis_print,
+            )
             plot_classes_distribution(
-                val_y, "validation", vis_format=vis_format
+                val_y,
+                "validation",
+                vis_format=vis_format,
+                vis_print=vis_print,
             )
 
         train_size = len(trains.dataset)
@@ -251,9 +266,12 @@ def load_data(
         tests = test_loader.get_loader(batch_size=batch)
 
         if vis_his and fabric.global_rank == 0:
-            eval_y = list(sum([y[1] for _, y in enumerate(tests)], ()))
+            eval_y = _labels(test_loader.dataset)
             plot_classes_distribution(
-                eval_y, "evaluation", vis_format=vis_format
+                eval_y,
+                "evaluation",
+                vis_format=vis_format,
+                vis_print=vis_print,
             )
 
         tests = fabric.setup_dataloaders(tests)
@@ -278,6 +296,7 @@ def get_affinity_matrix(
     classes: list = [],
     vis_aff: bool = False,
     vis_format: str = "png",
+    vis_print: bool = False,
 ) -> pd.DataFrame:
     """Loads affinity matrix from a given path, subsets it given selected classes and returns it as a pandas DataFrame.
 
@@ -309,12 +328,38 @@ def get_affinity_matrix(
                 )
             )
         if vis_aff:
-            plot_affinity_matrix(
-                lookup=affinity,
-                all_classes=affinity.columns.tolist(),
-                selected_classes=classes,
-                vis_format=vis_format,
+            plot_path = f"plots/affinity_matrix.{vis_format}"
+            # selected classes are highlighted in the plot, so they are part
+            # of what makes an existing plot up to date
+            classes_path = "plots/affinity_matrix_classes.txt"
+            selected = "\n".join(
+                [
+                    str(pathlib.Path(affinity_path).resolve()),
+                    f"vis_print={vis_print}",
+                    *map(str, classes),
+                ]
             )
+            if (
+                os.path.exists(plot_path)
+                and os.path.getmtime(plot_path)
+                > os.path.getmtime(affinity_path)
+                and os.path.exists(classes_path)
+                and pathlib.Path(classes_path).read_text() == selected
+            ):
+                logging.info(
+                    "Affinity matrix plot is up to date, skipping: "
+                    f"{plot_path}"
+                )
+            else:
+                plot_affinity_matrix(
+                    lookup=affinity,
+                    all_classes=affinity.columns.tolist(),
+                    selected_classes=classes,
+                    vis_format=vis_format,
+                    vis_print=vis_print,
+                )
+                os.makedirs("plots", exist_ok=True)
+                pathlib.Path(classes_path).write_text(selected)
 
         # subset affinity matrix with only the relevant classes
         index = [affinity.columns.get_loc(f"{columns}") for columns in classes]
@@ -326,8 +371,22 @@ def get_affinity_matrix(
         return None
 
 
-class AffinityDiskDataset(caked.dataloader.DiskDataset):
+def _labels(dataset, indices=None) -> list:
+    """Get labels from file names (no data loading) for a dataset with `paths`
+    or a torch Subset of one. `indices` selects positions within `dataset`."""
+    if isinstance(dataset, torch.utils.data.Subset):
+        indices = (
+            list(dataset.indices)
+            if indices is None
+            else [dataset.indices[i] for i in indices]
+        )
+        dataset = dataset.dataset
+    if indices is None:
+        indices = range(len(dataset.paths))
+    return [pathlib.Path(dataset.paths[i]).name.split("_")[0] for i in indices]
 
+
+class AffinityDiskDataset(caked.dataloader.DiskDataset):
     """Modified version of the caked DiskDataset to include the affinity matrix and data metadata that is needed for the
     affinity pipeline"""
 
@@ -357,7 +416,7 @@ def __getitem__(self, index):
         # get file basename
         filename = pathlib.Path(self.paths[index]).name
         # ground truth
-        y = pathlib.Path(filename).name.split("_")[0]
+        y = _labels(self, [index])[0]
 
         # similarity column / vector
         if self.affinity is not None:

avae/decoders/differentiable.py

@@ -250,7 +250,7 @@ def configure_renderer(
         `shape` of the output.
         """
         self._shape = shape
-        self._default_axis = default_axis.as_tensor()
+        self._default_axis = default_axis.as_tensor().to(device)
         self._splatter = GaussianSplatRenderer(
             shape,
             device=device,

avae/evaluate.py

@@ -192,6 +192,7 @@ def evaluate(params):
             data_dim,
             mode="evl",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
 
     # visualise latent disentanglement
@@ -204,6 +205,7 @@ def evaluate(params):
             poses=p_test,
             mode="_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
 
     # visualise pose disentanglement
@@ -216,6 +218,7 @@ def evaluate(params):
             device,
             mode="_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
 
     if pose_dims != 0 and params.vis_pose_class:
@@ -229,6 +232,7 @@ def evaluate(params):
             device,
             mode="_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
     # visualise interpolations
     if params.vis_int:
@@ -241,6 +245,7 @@ def evaluate(params):
             poses=p_test,
             mode="_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
 
     # visualise embeddings
@@ -251,6 +256,7 @@ def evaluate(params):
             classes_list,
             "_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
             embedding=latent_embedding[eval_mask],
         )
 
@@ -260,6 +266,7 @@ def evaluate(params):
             np.array(y_test),
             mode="_eval",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
         )
 
     # ############################# Predict #############################
@@ -288,6 +295,7 @@ def evaluate(params):
             classes_list,
             "_train_eval_comparison",
             vis_format=params.vis_format,
+            vis_print=params.vis_print,
             embedding=latent_embedding,
         )
 
@@ -306,17 +314,18 @@ def evaluate(params):
         classifier=params.classifier,
     )
     logging.info(
-        "------------------->>> Accuracy: Train: %f | Val : %f | Val with unseen labels: %f\n"
+        "------------------->>> Overall accuracy: Train: %f | Val : %f | Val with unseen labels: %f\n"
         % (train_acc, val_acc_selected, val_acc)
     )
-    vis.accuracy_plot(
+    vis.confusion_plot(
         np.array(latents_training_id),
         ypred_train,
         y_test,
         ypred_val,
         params.classes,
         mode="_eval",
         vis_format=params.vis_format,
+        vis_print=params.vis_print,
     )
     vis.f1_plot(
         np.array(latents_training_id),
@@ -325,6 +334,7 @@ def evaluate(params):
         ypred_val,
         mode="_eval",
         vis_format=params.vis_format,
+        vis_print=params.vis_print,
     )
     logging.info("Saving meta files with evaluation data.")