Add more tests for weight_format

Author: constantinpape

tests/test_cli.py

@@ -16,7 +16,7 @@ def test_cli_test_model(unet2d_nuclei_broad_model):
     assert ret.returncode == 0
 
 
-def test_cli_test_model_with_specific_weight_format(unet2d_nuclei_broad_model):
+def test_cli_test_model_with_weight_format(unet2d_nuclei_broad_model):
     ret = subprocess.run(
         ["bioimageio", "test-model", unet2d_nuclei_broad_model, "--weight-format", "pytorch_state_dict"]
     )
@@ -28,25 +28,34 @@ def test_cli_test_resource(unet2d_nuclei_broad_model):
     assert ret.returncode == 0
 
 
-def test_cli_test_resource_with_specific_weight_format(unet2d_nuclei_broad_model):
+def test_cli_test_resource_with_weight_format(unet2d_nuclei_broad_model):
     ret = subprocess.run(
         ["bioimageio", "test-model", unet2d_nuclei_broad_model, "--weight-format", "pytorch_state_dict"]
     )
     assert ret.returncode == 0
 
 
-def test_cli_predict_image(unet2d_nuclei_broad_model, tmp_path):
-    spec = load_resource_description(unet2d_nuclei_broad_model)
+def _test_cli_predict_image(model, tmp_path, extra_kwargs=None):
+    spec = load_resource_description(model)
     in_path = spec.test_inputs[0]
     out_path = tmp_path.with_suffix(".npy")
-    ret = subprocess.run(
-        ["bioimageio", "predict-image", unet2d_nuclei_broad_model, "--inputs", str(in_path), "--outputs", str(out_path)]
-    )
+    cmd = ["bioimageio", "predict-image", model, "--inputs", str(in_path), "--outputs", str(out_path)]
+    if extra_kwargs is not None:
+        cmd.extend(extra_kwargs)
+    ret = subprocess.run(cmd)
     assert ret.returncode == 0
     assert out_path.exists()
 
 
-def test_cli_predict_images(unet2d_nuclei_broad_model, tmp_path):
+def test_cli_predict_image(unet2d_nuclei_broad_model, tmp_path):
+    _test_cli_predict_image(unet2d_nuclei_broad_model, tmp_path)
+
+
+def test_cli_predict_image_with_weight_format(unet2d_nuclei_broad_model, tmp_path):
+    _test_cli_predict_image(unet2d_nuclei_broad_model, tmp_path, ["--weight-format", "pytorch_state_dict"])
+
+
+def _test_cli_predict_images(model, tmp_path, extra_kwargs=None):
     n_images = 3
     shape = (1, 1, 128, 128)
     expected_shape = (1, 1, 128, 128)
@@ -64,14 +73,25 @@ def test_cli_predict_images(unet2d_nuclei_broad_model, tmp_path):
         expected_outputs.append(out_folder / f"im-{i}.npy")
 
     input_pattern = str(in_folder / "*.npy")
-    ret = subprocess.run(["bioimageio", "predict-images", unet2d_nuclei_broad_model, input_pattern, str(out_folder)])
+    cmd = ["bioimageio", "predict-images", model, input_pattern, str(out_folder)]
+    if extra_kwargs is not None:
+        cmd.extend(extra_kwargs)
+    ret = subprocess.run(cmd)
     assert ret.returncode == 0
 
     for out_path in expected_outputs:
         assert out_path.exists()
         assert np.load(out_path).shape == expected_shape
 
 
+def test_cli_predict_images(unet2d_nuclei_broad_model, tmp_path):
+    _test_cli_predict_images(unet2d_nuclei_broad_model, tmp_path)
+
+
+def test_cli_predict_images_with_weight_format(unet2d_nuclei_broad_model, tmp_path):
+    _test_cli_predict_images(unet2d_nuclei_broad_model, tmp_path, ["--weight-format", "pytorch_state_dict"])
+
+
 def test_torch_to_torchscript(unet2d_nuclei_broad_model, tmp_path):
     out_path = tmp_path.with_suffix(".pt")
     ret = subprocess.run(

tests/test_prediction.py

@@ -39,6 +39,24 @@ def test_predict_image(unet2d_fixed_shape_or_not, tmpdir):
         assert_array_almost_equal(res, exp, decimal=4)
 
 
+def test_predict_image_with_weight_format(unet2d_fixed_shape_or_not, tmpdir):
+    from bioimageio.core.prediction import predict_image
+
+    spec = load_resource_description(unet2d_fixed_shape_or_not)
+    assert isinstance(spec, Model)
+    inputs = spec.test_inputs
+
+    outputs = [Path(tmpdir) / f"out{i}.npy" for i in range(len(spec.test_outputs))]
+    predict_image(unet2d_fixed_shape_or_not, inputs, outputs, weight_format="pytorch_state_dict")
+    for out_path in outputs:
+        assert out_path.exists()
+
+    result = [np.load(str(p)) for p in outputs]
+    expected = [np.load(str(p)) for p in spec.test_outputs]
+    for res, exp in zip(result, expected):
+        assert_array_almost_equal(res, exp, decimal=4)
+
+
 def test_predict_image_with_padding(unet2d_fixed_shape_or_not, tmp_path):
     any_model = unet2d_fixed_shape_or_not  # todo: replace 'unet2d_fixed_shape_or_not' with 'any_model'
     from bioimageio.core.prediction import predict_image