Also detect prediction layer & add refresh_metadata function (#75)

* also detect prediction layer

* format

* add refresh_metadata method

* only override wkw resolutions for existing layers when using refresh_metadata

* black

* add CLI for refresh metadata

Author: philippotto

README.md

@@ -16,7 +16,7 @@ Created with [Python3](https://www.python.org/).
 * `wkcuber.convert_knossos`: Convert KNOSSOS cubes to WKW cubes
 * `wkcuber.downsampling`: Create downsampled magnifications (with `median`, `mode` and linear interpolation modes)
 * `wkcuber.compress`: Compress WKW cubes for efficient file storage (especially useful for segmentation data)
-* `wkcuber.metadata`: Create metadata (with guessing of most parameters)
+* `wkcuber.metadata`: Create (or refresh) metadata (with guessing of most parameters)
 * Most modules support multiprocessing
 
 ## Supported input formats
@@ -73,6 +73,9 @@ python -m wkcuber.compress --layer_name segmentation data/target data/target_com
 
 # Create metadata
 python -m wkcuber.metadata --name great_dataset --scale 11.24,11.24,25 data/target
+
+# Refresh metadata so that new layers and/or magnifications are picked up
+python -m wkcuber.metadata --refresh data/target
 ```
 
 ### Parallelization

wkcuber/metadata.py

@@ -12,20 +12,25 @@
 from typing import List
 
 
+def get_datasource_path(dataset_path):
+    return path.join(dataset_path, "datasource-properties.json")
+
+
 def create_parser():
     parser = ArgumentParser()
 
     parser.add_argument("path", help="Directory containing the dataset.")
 
-    parser.add_argument("--name", "-n", help="Name of the dataset")
-
+    parser.add_argument("--name", "-n", help="Name of the dataset", default=None)
     parser.add_argument(
         "--scale",
         "-s",
         help="Scale of the dataset (e.g. 11.2,11.2,25)",
         default="1,1,1",
     )
 
+    parser.add_argument("--refresh", "-r", default=False, action="store_true")
+
     group = parser.add_mutually_exclusive_group()
     group.add_argument(
         "--compute_max_id",
@@ -39,6 +44,23 @@ def create_parser():
     return parser
 
 
+def write_datasource_properties(dataset_path, datasource_properties):
+    datasource_properties_path = get_datasource_path(dataset_path)
+    with open(datasource_properties_path, "wt") as datasource_properties_file:
+        json.dump(datasource_properties, datasource_properties_file, indent=2)
+
+
+def read_datasource_properties(dataset_path):
+    with open(get_datasource_path(dataset_path), "r") as datasource_properties_file:
+        return json.load(datasource_properties_file)
+
+
+"""
+Creates a datasource-properties.json file with the specified properties
+for the given dataset path. Common layers are detected automatically.
+"""
+
+
 def write_webknossos_metadata(
     dataset_path,
     name,
@@ -50,28 +72,71 @@ def write_webknossos_metadata(
 
     # Generate a metadata file for webKnossos
     # Currently includes no source of information for team
-    datasource_properties_path = path.join(dataset_path, "datasource-properties.json")
     layers = list(
         detect_layers(dataset_path, max_id, compute_max_id, exact_bounding_box)
     )
-    with open(datasource_properties_path, "wt") as datasource_properties_json:
-        json.dump(
-            {
-                "id": {"name": name, "team": "<unknown>"},
-                "dataLayers": layers,
-                "scale": scale,
-            },
-            datasource_properties_json,
-            indent=2,
+    write_datasource_properties(
+        dataset_path,
+        {
+            "id": {"name": name, "team": "<unknown>"},
+            "dataLayers": layers,
+            "scale": scale,
+        },
+    )
+
+
+"""
+Updates the datasource-properties.json file for a given dataset.
+Use this method if you added (or removed) layers and/or changed magnifications for
+existing layers.
+
+Raises an exception if the datasource-properties.json file does not exist, yet.
+In this case, use write_webknossos_metadata instead.
+"""
+
+
+def refresh_metadata(
+    wkw_path, max_id=0, compute_max_id=False, exact_bounding_box: Optional[dict] = None
+):
+    dataset_path = get_datasource_path(wkw_path)
+    if not path.exists(dataset_path):
+        raise Exception(
+            "datasource-properties.json file could not be found. Please use write_webknossos_metadata to create it."
         )
 
+    datasource_properties = read_datasource_properties(wkw_path)
+    existing_layers_dict = {
+        layer["name"]: layer for layer in datasource_properties["dataLayers"]
+    }
+
+    new_layers = list(
+        detect_layers(wkw_path, max_id, compute_max_id, exact_bounding_box)
+    )
+
+    # Merge the freshly read layers with the existing layer information, so that information,
+    # such as bounding boxes, are not lost for existing layers.
+    # For existing layers, only the resolutions will be updated.
+    merged_layers = []
+    for new_layer in new_layers:
+        layer_name = new_layer["name"]
+        if layer_name in existing_layers_dict:
+            existing_layer = existing_layers_dict[layer_name]
+            # Update the resolutions
+            existing_layer["wkwResolutions"] = new_layer["wkwResolutions"]
+            merged_layers.append(existing_layer)
+        else:
+            merged_layers.append(new_layer)
+
+    datasource_properties["dataLayers"] = merged_layers
+    write_datasource_properties(wkw_path, datasource_properties)
+
 
 def read_metadata_for_layer(wkw_path, layer_name):
-    datasource_properties = json.load(
-        open(path.join(wkw_path, "datasource-properties.json"), "r")
-    )
+    datasource_properties = read_datasource_properties(wkw_path)
+
     layers = datasource_properties["dataLayers"]
     layer_info = next(layer for layer in layers if layer["name"] == layer_name)
+
     dtype = np.dtype(layer_info["elementClass"])
     bounding_box = layer_info["boundingBox"]
     origin = bounding_box["topLeft"]
@@ -146,13 +211,15 @@ def detect_resolutions(dataset_path, layer) -> List[Mag]:
 def detect_standard_layer(dataset_path, layer_name, exact_bounding_box=None):
     # Perform metadata detection for well-known layers
 
+    mags = list(detect_resolutions(dataset_path, layer_name))
+    mags = sorted(mags)
+    assert len(mags) > 0, "No resolutions found"
+
     if exact_bounding_box is None:
-        bbox = detect_bbox(dataset_path, layer_name)
+        bbox = detect_bbox(dataset_path, layer_name, mags[0])
     else:
         bbox = exact_bounding_box
 
-    mags = list(detect_resolutions(dataset_path, layer_name))
-    mags = sorted(mags)
     resolutions = [
         {
             "resolution": mag.to_array(),
@@ -161,7 +228,6 @@ def detect_standard_layer(dataset_path, layer_name, exact_bounding_box=None):
         for mag in mags
     ]
 
-    assert len(mags) > 0, "No resolutions found"
     dtype = detect_dtype(dataset_path, layer_name, mags[0])
 
     return {
@@ -205,9 +271,10 @@ def detect_segmentation_layer(
 
 
 def detect_layers(dataset_path, max_id, compute_max_id, exact_bounding_box=None):
-    # Detect metadata for well-known layers, e.g. color and segmentation
-    if path.exists(path.join(dataset_path, "color")):
-        yield detect_standard_layer(dataset_path, "color", exact_bounding_box)
+    # Detect metadata for well-known layers (i.e., color, prediction and segmentation)
+    for layer_name in ["color", "prediction"]:
+        if path.exists(path.join(dataset_path, layer_name)):
+            yield detect_standard_layer(dataset_path, layer_name, exact_bounding_box)
     if path.exists(path.join(dataset_path, "segmentation")):
         yield detect_segmentation_layer(
             dataset_path, "segmentation", max_id, compute_max_id, exact_bounding_box
@@ -217,7 +284,17 @@ def detect_layers(dataset_path, max_id, compute_max_id, exact_bounding_box=None)
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
     args = create_parser().parse_args()
-    scale = tuple(float(x) for x in args.scale.split(","))
-    write_webknossos_metadata(
-        args.path, args.name, scale, args.max_id, args.compute_max_id
-    )
+    if not args.refresh:
+        assert (
+            args.name is not None
+        ), "Please provide a name via --name to create meta data."
+        scale = tuple(float(x) for x in args.scale.split(","))
+        write_webknossos_metadata(
+            args.path, args.name, scale, args.max_id, args.compute_max_id
+        )
+    else:
+        if args.name is not None:
+            logging.warn(
+                "The --name argument is ignored, since --refresh was provided."
+            )
+        refresh_metadata(args.path, args.max_id, args.compute_max_id)