Reject ambiguous TIFF axes in inference and postprocessing

Author: yanchaoz

Postprocess/FRMC_post.py

@@ -93,7 +93,10 @@ def _load_volume(path):
         return np.load(path, allow_pickle=False)
     if path.suffix.lower() in (".tif", ".tiff"):
         import tifffile
-        return tifffile.imread(path)
+        with tifffile.TiffFile(path) as image:
+            if len(image.series) != 1 or image.series[0].axes not in {"ZYX", "QYX", "IYX"}:
+                raise ValueError("TIFF must contain one grayscale ZYX stack")
+            return image.series[0].asarray()
     raise ValueError(f"Expected .npy, .tif or .tiff: {path}")
 
 

Train_and_Inference/inference.py

@@ -24,7 +24,7 @@ def load_volume(path):
         import tifffile
 
         with tifffile.TiffFile(path) as image:
-            if len(image.series) != 1 or any(axis in image.series[0].axes for axis in "CST"):
+            if len(image.series) != 1 or image.series[0].axes not in {"ZYX", "QYX", "IYX"}:
                 raise ValueError("TIFF must contain one grayscale ZYX stack, without color, channel or time axes")
             volume = image.series[0].asarray()
     else:
@@ -152,7 +152,7 @@ def sha256(path):
 def main(argv=None):
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--input", required=True, type=Path, help="uint8 3D ZYX TIFF or NPY raw volume")
-    parser.add_argument("--checkpoint", required=True, type=Path, help="MNet checkpoint (.pth or .pt)")
+    parser.add_argument("--checkpoint", required=True, type=Path, help="MNet checkpoint (.ckpt, .pth or .pt)")
     parser.add_argument("--output-dir", required=True, type=Path, help="New directory; existing paths are refused")
     parser.add_argument("--device", default="cpu", help="cpu (default) or cuda:N, e.g. cuda:0")
     args = parser.parse_args(argv)

tests/test_inference.py

@@ -85,6 +85,9 @@ def test_rejects_rgb_and_time_series_tiff(self):
             tifffile.imwrite(path, np.zeros((4, 8, 9), np.uint8), photometric="minisblack", metadata={"axes": "TYX"})
             with self.assertRaisesRegex(ValueError, "grayscale"):
                 inference.load_volume(path)
+            tifffile.imwrite(path, np.zeros((4, 8, 9), np.uint8), photometric="minisblack", metadata={"axes": "YZX"})
+            with self.assertRaisesRegex(ValueError, "grayscale"):
+                inference.load_volume(path)
 
     def test_checkpoint_formats_and_dataparallel_prefix(self):
         source = torch.nn.Conv3d(1, 1, 1)

tests/test_postprocess.py

@@ -21,6 +21,14 @@
 
 
 class PostprocessTests(unittest.TestCase):
+    def test_tiff_axes_are_not_silently_reinterpreted(self):
+        with tempfile.TemporaryDirectory() as temporary:
+            path = Path(temporary) / "head.tif"
+            tifffile.imwrite(path, np.zeros((4, 8, 9), dtype=np.float32),
+                             photometric="minisblack", metadata={"axes": "YZX"})
+            with self.assertRaisesRegex(ValueError, "ZYX"):
+                post._load_volume(path)
+
     def test_help_does_not_import_elf(self):
         code = (
             "import sys; sys.modules['elf'] = None; "