Enable lazy data loading in prediction check script

Author: constantinpape

scripts/prediction/check_prediction.py

@@ -7,31 +7,36 @@
 from elf.io import open_file
 
 
+def _load_n5(path, key, lazy):
+    data = open_file(path, "r")[key]
+    if not lazy:
+        data = data[:]
+    return data
+
+
 # will not work for very large images
-def check_prediction(input_path, output_folder, check_downsampled=False, input_key=None):
+def check_prediction(input_path, output_folder, check_downsampled=False, input_key=None, lazy=False):
     if input_key is None:
         input_ = imageio.imread(input_path)
     else:
-        input_ = open_file(input_path, "r")[input_key][:]
+        input_ = _load_n5(input_path, input_key, lazy)
 
     if check_downsampled:
 
         pred_path = os.path.join(output_folder, "predictions.zarr")
-        with open_file(pred_path, "r") as f:
-            prediction = f["prediction"][:]
+        prediction = _load_n5(pred_path, "prediction", lazy)
 
         seg_path = os.path.join(output_folder, "seg_downscaled.zarr")
-        with open_file(seg_path, "r") as f:
-            segmentation = f["segmentation"][:]
+        segmentation = _load_n5(seg_path, "segmentation", lazy)
 
         scale = (0.5, 0.5, 0.5)
 
     else:
         seg_path = os.path.join(output_folder, "segmentation.zarr")
+        segmentation = _load_n5(seg_path, "segmentation", lazy)
         with open_file(seg_path, "r") as f:
-            segmentation = f["segmentation"][:]
             if "segmentation_postprocessed" in f:
-                seg_pp = f["segmentation_postprocessed"][:]
+                seg_pp = _load_n5(seg_path, "segmentation_postprocessed", lazy)
             else:
                 seg_pp = None
 
@@ -54,9 +59,13 @@ def main():
     parser.add_argument("-o", "--output_folder", required=True)
     parser.add_argument("-k", "--input_key", default=None)
     parser.add_argument("-c", "--check_downsampled", action="store_true")
+    parser.add_argument("--lazy", action="store_true")
 
     args = parser.parse_args()
-    check_prediction(args.input, args.output_folder, check_downsampled=args.check_downsampled, input_key=args.input_key)
+    check_prediction(
+        args.input, args.output_folder, check_downsampled=args.check_downsampled,
+        input_key=args.input_key, lazy=args.lazy
+    )
 
 
 if __name__ == "__main__":