Merge branch 'main' into bump-patch

Author: constantinpape

bioimageio/core/__main__.py

@@ -70,14 +70,19 @@ def test_model(
     model_rdf: str = typer.Argument(
         ..., help="Path or URL to the model resource description file (rdf.yaml) or zipped model."
     ),
-    weight_format: Optional[WeightFormatEnum] = typer.Argument(None, help="The weight format to use."),
-    devices: Optional[List[str]] = typer.Argument(None, help="Devices for running the model."),
-    decimal: int = typer.Argument(4, help="The test precision."),
+    weight_format: Optional[WeightFormatEnum] = typer.Option(None, help="The weight format to use."),
+    devices: Optional[List[str]] = typer.Option(None, help="Devices for running the model."),
+    decimal: int = typer.Option(4, help="The test precision."),
 ) -> int:
     # this is a weird typer bug: default devices are empty tuple although they should be None
     if len(devices) == 0:
         devices = None
-    summary = resource_tests.test_model(model_rdf, weight_format=weight_format, devices=devices, decimal=decimal)
+    summary = resource_tests.test_model(
+        model_rdf,
+        weight_format=None if weight_format is None else weight_format.value,
+        devices=devices,
+        decimal=decimal,
+    )
     if summary["error"] is None:
         print(f"Model test for {model_rdf} has passed.")
         return 0
@@ -95,14 +100,16 @@ def test_resource(
     rdf: str = typer.Argument(
         ..., help="Path or URL to the resource description file (rdf.yaml) or zipped resource package."
     ),
-    weight_format: Optional[WeightFormatEnum] = typer.Argument(None, help="(for model only) The weight format to use."),
-    devices: Optional[List[str]] = typer.Argument(None, help="(for model only) Devices for running the model."),
-    decimal: int = typer.Argument(4, help="(for model only) The test precision."),
+    weight_format: Optional[WeightFormatEnum] = typer.Option(None, help="(for model only) The weight format to use."),
+    devices: Optional[List[str]] = typer.Option(None, help="(for model only) Devices for running the model."),
+    decimal: int = typer.Option(4, help="(for model only) The test precision."),
 ) -> int:
     # this is a weird typer bug: default devices are empty tuple although they should be None
     if len(devices) == 0:
         devices = None
-    summary = resource_tests.test_resource(rdf, weight_format=weight_format, devices=devices, decimal=decimal)
+    summary = resource_tests.test_resource(
+        rdf, weight_format=None if weight_format is None else weight_format.value, devices=devices, decimal=decimal
+    )
     if summary["error"] is None:
         print(f"Resource test for {rdf} has passed.")
         return 0
@@ -129,10 +136,10 @@ def predict_image(
     # tiling: Optional[Union[str, bool]] = typer.Argument(
     #     None, help="Padding to apply in each dimension passed as json encoded string."
     # ),
-    padding: Optional[bool] = typer.Argument(None, help="Whether to pad the image to a size suited for the model."),
-    tiling: Optional[bool] = typer.Argument(None, help="Whether to run prediction in tiling mode."),
-    weight_format: Optional[str] = typer.Argument(None, help="The weight format to use."),
-    devices: Optional[List[str]] = typer.Argument(None, help="Devices for running the model."),
+    padding: Optional[bool] = typer.Option(None, help="Whether to pad the image to a size suited for the model."),
+    tiling: Optional[bool] = typer.Option(None, help="Whether to run prediction in tiling mode."),
+    weight_format: Optional[WeightFormatEnum] = typer.Option(None, help="The weight format to use."),
+    devices: Optional[List[str]] = typer.Option(None, help="Devices for running the model."),
 ) -> int:
 
     if isinstance(padding, str):
@@ -145,7 +152,9 @@ def predict_image(
     # this is a weird typer bug: default devices are empty tuple although they should be None
     if len(devices) == 0:
         devices = None
-    prediction.predict_image(model_rdf, inputs, outputs, padding, tiling, weight_format, devices)
+    prediction.predict_image(
+        model_rdf, inputs, outputs, padding, tiling, None if weight_format is None else weight_format.value, devices
+    )
     return 0
 
 
@@ -167,10 +176,10 @@ def predict_images(
     # tiling: Optional[Union[str, bool]] = typer.Argument(
     #     None, help="Padding to apply in each dimension passed as json encoded string."
     # ),
-    padding: Optional[bool] = typer.Argument(None, help="Whether to pad the image to a size suited for the model."),
-    tiling: Optional[bool] = typer.Argument(None, help="Whether to run prediction in tiling mode."),
-    weight_format: Optional[str] = typer.Argument(None, help="The weight format to use."),
-    devices: Optional[List[str]] = typer.Argument(None, help="Devices for running the model."),
+    padding: Optional[bool] = typer.Option(None, help="Whether to pad the image to a size suited for the model."),
+    tiling: Optional[bool] = typer.Option(None, help="Whether to run prediction in tiling mode."),
+    weight_format: Optional[WeightFormatEnum] = typer.Option(None, help="The weight format to use."),
+    devices: Optional[List[str]] = typer.Option(None, help="Devices for running the model."),
 ) -> int:
     input_files = glob(input_pattern)
     input_names = [os.path.split(infile)[1] for infile in input_files]
@@ -194,7 +203,7 @@ def predict_images(
         output_files,
         padding=padding,
         tiling=tiling,
-        weight_format=weight_format,
+        weight_format=None if weight_format is None else weight_format.value,
         devices=devices,
         verbose=True,
     )
@@ -213,8 +222,8 @@ def convert_torch_weights_to_onnx(
         ),
         output_path: Path = typer.Argument(..., help="Where to save the onnx weights."),
         opset_version: Optional[int] = typer.Argument(12, help="Onnx opset version."),
-        use_tracing: bool = typer.Argument(True, help="Whether to use torch.jit tracing or scripting."),
-        verbose: bool = typer.Argument(True, help="Verbosity"),
+        use_tracing: bool = typer.Option(True, help="Whether to use torch.jit tracing or scripting."),
+        verbose: bool = typer.Option(True, help="Verbosity"),
     ) -> int:
         return torch_converter.convert_weights_to_onnx(model_rdf, output_path, opset_version, use_tracing, verbose)
 
@@ -226,7 +235,7 @@ def convert_torch_weights_to_torchscript(
             ..., help="Path to the model resource description file (rdf.yaml) or zipped model."
         ),
         output_path: Path = typer.Argument(..., help="Where to save the torchscript weights."),
-        use_tracing: bool = typer.Argument(True, help="Whether to use torch.jit tracing or scripting."),
+        use_tracing: bool = typer.Option(True, help="Whether to use torch.jit tracing or scripting."),
     ) -> int:
         return torch_converter.convert_weights_to_pytorch_script(model_rdf, output_path, use_tracing)
 

tests/test_cli.py

@@ -16,18 +16,46 @@ def test_cli_test_model(unet2d_nuclei_broad_model):
     assert ret.returncode == 0
 
 
-def test_cli_predict_image(unet2d_nuclei_broad_model, tmp_path):
-    spec = load_resource_description(unet2d_nuclei_broad_model)
-    in_path = spec.test_inputs[0]
-    out_path = tmp_path.with_suffix(".npy")
+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"]
+    )
+    assert ret.returncode == 0
+
+
+def test_cli_test_resource(unet2d_nuclei_broad_model):
+    ret = subprocess.run(["bioimageio", "test-model", unet2d_nuclei_broad_model])
+    assert ret.returncode == 0
+
+
+def test_cli_test_resource_with_weight_format(unet2d_nuclei_broad_model):
     ret = subprocess.run(
-        ["bioimageio", "predict-image", unet2d_nuclei_broad_model, "--inputs", str(in_path), "--outputs", str(out_path)]
+        ["bioimageio", "test-model", unet2d_nuclei_broad_model, "--weight-format", "pytorch_state_dict"]
     )
     assert ret.returncode == 0
+
+
+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")
+    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)
@@ -45,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