🔀 Undo unwanted changes during merge.

Author: shaneahmed

tests/models/test_models_abc.py

@@ -6,15 +6,15 @@
 
 import pytest
 import torch
+import torchvision.models as torch_models
 from torch import nn
 
-import tiatoolbox.models
 from tiatoolbox import rcParam, utils
 from tiatoolbox.models.architecture import (
     fetch_pretrained_weights,
     get_pretrained_model,
 )
-from tiatoolbox.models.models_abc import ModelABC
+from tiatoolbox.models.models_abc import ModelABC, model_to
 from tiatoolbox.utils import env_detection as toolbox_env
 
 if TYPE_CHECKING:
@@ -154,19 +154,16 @@ def test_model_abc() -> None:
 
 def test_model_to() -> None:
     """Test for placing model on device."""
-    import torchvision.models as torch_models
-    from torch import nn
-
     # Test on GPU
-    # no GPU on Travis so this will crash
+    # no GPU on GitHub Actions so this will crash
     if not utils.env_detection.has_gpu():
         model = torch_models.resnet18()
         with pytest.raises((AssertionError, RuntimeError)):
-            _ = tiatoolbox.models.models_abc.model_to(device="cuda", model=model)
+            _ = model_to(device="cuda", model=model)
 
     # Test on CPU
     model = torch_models.resnet18()
-    model = tiatoolbox.models.models_abc.model_to(device="cpu", model=model)
+    model = model_to(device="cpu", model=model)
     assert isinstance(model, nn.Module)
 
 

tests/test_utils.py

@@ -1669,6 +1669,17 @@ def test_patch_pred_store() -> None:
     with pytest.raises(ValueError, match="coordinates"):
         misc.dict_to_store(patch_output, (1.0, 1.0))
 
+    patch_output = {
+        "predictions": [1, 0, 1],
+        "coordinates": [(0, 0, 1, 1), (1, 1, 2, 2), (2, 2, 3, 3)],
+        "other": "other",
+    }
+
+    store = misc.dict_to_store(patch_output, (1.0, 1.0))
+
+    # Check that it is an SQLiteStore containing the expected annotations
+    assert isinstance(store, SQLiteStore)
+
 
 def test_patch_pred_store_cdict() -> None:
     """Test patch_pred_store with a class dict."""

tiatoolbox/cli/common.py

@@ -94,8 +94,10 @@ def cli_output_type(
     input_type: click.Choice | None = None,
 ) -> Callable:
     """Enables --file-types option for cli."""
-    if input_type is None:
-        input_type = click.Choice(["zarr", "AnnotationStore"], case_sensitive=False)
+    click_choices = click.Choice(
+        choices=["zarr", "AnnotationStore"], case_sensitive=False
+    )
+    input_type = click_choices if input_type is None else input_type
     return click.option(
         "--output-type",
         help=add_default_to_usage_help(usage_help, default),
@@ -410,20 +412,6 @@ def cli_yaml_config_path(
     )
 
 
-def cli_on_gpu(
-    usage_help: str = "Run the model on GPU.",
-    *,
-    default: bool = False,
-) -> Callable:
-    """Enables --on-gpu option for cli."""
-    return click.option(
-        "--on-gpu",
-        type=bool,
-        default=default,
-        help=add_default_to_usage_help(usage_help, default),
-    )
-
-
 def cli_num_loader_workers(
     usage_help: str = "Number of workers to load the data. Please note that they will "
     "also perform preprocessing.",

tiatoolbox/cli/patch_predictor.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import click
-
 from tiatoolbox.cli.common import (
     cli_batch_size,
     cli_device,
@@ -45,7 +43,6 @@
 @cli_num_loader_workers(default=0)
 @cli_output_type(
     default="AnnotationStore",
-    input_type=click.Choice(["zarr", "AnnotationStore"], case_sensitive=False),
 )
 @cli_patch_mode(default=False)
 @cli_return_probabilities(default=True)