add --runslow option to pytest

conftest.py

@@ -1,4 +1,5 @@
 import time
+import numpy as np
 import pytest
 from cellpose import utils, models, vit_sam
 import zipfile
@@ -7,6 +8,26 @@
 from pathlib import Path
 
 
+def pytest_addoption(parser):
+    parser.addoption(
+        "--runslow", action="store_true", default=False, help="run slow tests"
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "slow: mark test as slow to run")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--runslow"):
+        # --runslow given in cli: do not skip slow tests
+        return
+    skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+    for item in items:
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)
+
+
 @pytest.fixture()
 def image_names():
     image_names = ['rgb_2D_tif.tif', 'rgb_2D.png', 'gray_2D.png']
@@ -58,7 +79,7 @@ def __init__(self, use_layers: int):
         super().__init__()
 
         self.use_layers = use_layers
-
+        self.layer_idxs = np.linspace(0, 23, self.use_layers, dtype=int)
 
     def forward(self, x):
         # same progression as SAM until readout
@@ -68,10 +89,8 @@ def forward(self, x):
             x = x + self.encoder.pos_embed
 
         # only use self.use_layers layers
-        for i, block in enumerate(self.encoder.blocks):
-            if i == self.use_layers:
-                break
-            x = block(x)
+        for layer_idx in self.layer_idxs:
+            x = self.encoder.blocks[layer_idx](x)
 
         x = self.encoder.neck(x.permute(0, 3, 1, 2))
 

tests/test_import.py

@@ -14,8 +14,6 @@ def test_gui_imports_without_error():
 
 
 def test_gpu_check():
-    #     from cellpose import models
-    #     models.use_gpu()
     from cellpose import core
     core.use_gpu()
 

tests/test_output.py

@@ -1,6 +1,5 @@
-from cellpose import io, metrics
-from cellpose.models import CellposeModel
-import platform
+from cellpose import io, metrics, utils
+import pytest
 from subprocess import check_output, STDOUT
 import os
 import numpy as np
@@ -62,22 +61,22 @@ def test_class_2D(data_dir, image_names, cellposemodel_fixture_2D):
     compare_masks_cp4(data_dir, image_names[0], "2D")
     clear_output(data_dir, image_names)
 
-# def test_cyto2_to_seg(data_dir, image_names):
-#     clear_output(data_dir, image_names)
-#     use_gpu = torch.cuda.is_available()
-#     file_names = [data_dir / "2D" / n for n in image_names]
-#     imgs = [io.imread_2D(file_name) for file_name in file_names]
-#     model = models.CellposeModel(gpu=use_gpu)
 
-#     # masks, flows, styles = model.eval(imgs, diameter=30)  # Errors during SAM stuff
-#     masks, flows, _ = model.eval(imgs, bsize=256, batch_size=64, normalize=True)
+@pytest.mark.slow
+def test_cyto2_to_seg(data_dir, image_names, cellposemodel_fixture_2D):
+    clear_output(data_dir, image_names)
+    file_names = [data_dir / "2D" / n for n in image_names]
+    imgs = [io.imread_2D(file_name) for file_name in file_names]
+
+    # masks, flows, styles = model.eval(imgs, diameter=30)  # Errors during SAM stuff
+    masks, flows, _ = cellposemodel_fixture_2D.eval(imgs, bsize=256, batch_size=64, normalize=True)
 
-#     for file_name, mask in zip(file_names, masks):
-#         io.imsave(data_dir/'2D'/(file_name.stem + '_cp_masks.png'), mask)
+    for file_name, mask in zip(file_names, masks):
+        io.imsave(data_dir/'2D'/(file_name.stem + '_cp_masks.png'), mask)
 
-#     io.masks_flows_to_seg(imgs, masks, flows, file_names)
-#     compare_masks_cp4(data_dir, image_names, "2D")
-#     clear_output(data_dir, image_names)
+    io.masks_flows_to_seg(imgs, masks, flows, file_names)
+    compare_masks_cp4(data_dir, image_names, "2D")
+    clear_output(data_dir, image_names)
 
 
 def test_class_3D(data_dir, image_names_3d, cellposemodel_fixture_3D):
@@ -113,51 +112,50 @@ def test_cli_2D(data_dir, image_names):
     clear_output(data_dir, image_names)
 
 
-# def test_cli_3D_diam(data_dir, image_names_3d):
-#     clear_output(data_dir, image_names_3d)
-#     use_gpu = torch.cuda.is_available()
-#     gpu_string = "--use_gpu" if use_gpu else ""
-#     cmd = f"python -m cellpose --image_path {str(data_dir / "3D" / image_names_3d[0])} --do_3D --diameter 25 --save_tif {gpu_string} --verbose"
-#     try:
-#         cmd_stdout = check_output(cmd, stderr=STDOUT, shell=True).decode()
-#         print(cmd_stdout)
-#     except Exception as e:
-#         print(e)
-#         raise ValueError(e)
-#     compare_masks_cp4(data_dir, image_names_3d[0], "3D")
-#     clear_output(data_dir, image_names_3d)
-
-
-# def test_outlines_list(data_dir, image_names):
-#     """ test both single and multithreaded by comparing them"""
-#     clear_output(data_dir, image_names)
-#     model_type = "cyto"
-#     channels = [2, 1]
-#     image_name = "rgb_2D.png"
-
-#     file_name = str(data_dir.joinpath("2D").joinpath(image_name))
-#     img = io.imread(file_name)
-
-#     model = models.CellposeModel(model_type=model_type)
-#     masks, _, _ = model.eval(img, diameter=30)
-#     outlines_single = utils.outlines_list(masks, multiprocessing=False)
-#     outlines_multi = utils.outlines_list(masks, multiprocessing=True)
-
-#     assert len(outlines_single) == len(outlines_multi)
-
-#     # Check that the outlines are the same, but not necessarily in the same order
-#     outlines_matched = [False] * len(outlines_single)
-#     for i, outline_single in enumerate(outlines_single):
-#         for j, outline_multi in enumerate(outlines_multi):
-#             if not outlines_matched[j] and np.array_equal(outline_single,
-#                                                           outline_multi):
-#                 outlines_matched[j] = True
-#                 break
-#         else:
-#             assert False, "Outline not found in outlines_multi: {}".format(
-#                 outline_single)
-
-#     assert all(outlines_matched), "Not all outlines in outlines_multi were matched"
+@pytest.mark.slow
+def test_cli_3D_diam(data_dir, image_names_3d):
+    clear_output(data_dir, image_names_3d)
+    use_gpu = torch.cuda.is_available()
+    gpu_string = "--use_gpu" if use_gpu else ""
+    cmd = f"python -m cellpose --image_path {str(data_dir / '3D' / image_names_3d[0])} --do_3D --diameter 25 --save_tif {gpu_string} --verbose"
+    try:
+        cmd_stdout = check_output(cmd, stderr=STDOUT, shell=True).decode()
+        print(cmd_stdout)
+    except Exception as e:
+        print(e)
+        raise ValueError(e)
+    compare_masks_cp4(data_dir, image_names_3d[0], "3D")
+    clear_output(data_dir, image_names_3d)
+
+
+@pytest.mark.slow
+def test_outlines_list(data_dir, image_names, cellposemodel_fixture_2D):
+    """ test both single and multithreaded by comparing them"""
+    clear_output(data_dir, image_names)
+    image_name = "rgb_2D.png"
+
+    file_name = str(data_dir.joinpath("2D").joinpath(image_name))
+    img = io.imread(file_name)
+
+    masks, _, _ = cellposemodel_fixture_2D.eval(img, diameter=30)
+    outlines_single = utils.outlines_list(masks, multiprocessing=False)
+    outlines_multi = utils.outlines_list(masks, multiprocessing=True)
+
+    assert len(outlines_single) == len(outlines_multi)
+
+    # Check that the outlines are the same, but not necessarily in the same order
+    outlines_matched = [False] * len(outlines_single)
+    for i, outline_single in enumerate(outlines_single):
+        for j, outline_multi in enumerate(outlines_multi):
+            if not outlines_matched[j] and np.array_equal(outline_single,
+                                                          outline_multi):
+                outlines_matched[j] = True
+                break
+        else:
+            assert False, "Outline not found in outlines_multi: {}".format(
+                outline_single)
+
+    assert all(outlines_matched), "Not all outlines in outlines_multi were matched"
 
 
 def compare_masks_cp4(data_dir, image_names, runtype):

tests/test_shape.py

@@ -1,14 +1,13 @@
-import platform
-from cellpose.models import CellposeModel
 import numpy as np
+import pytest
 
 
 #################### 2D Tests ####################
-# def test_shape_2D_grayscale():
-#     img = np.zeros((224, 224))
-#     model = models.CellposeModel()
-#     masks, _, _ = model.eval(img)
-#     assert masks.shape == (224, 224)
+@pytest.mark.slow
+def test_shape_2D_grayscale(cellposemodel_fixture_2D):
+    img = np.zeros((224, 224))
+    masks, _, _ = cellposemodel_fixture_2D.eval(img)
+    assert masks.shape == (224, 224)
 
 
 def test_shape_2D_chan_first_diam_resize(cellposemodel_fixture_2D):
@@ -19,11 +18,11 @@ def test_shape_2D_chan_first_diam_resize(cellposemodel_fixture_2D):
     assert flows[2].shape == (224, 224), 'cellprob shape mismatch'
 
 
-# def test_shape_2D_chan_diam_resize():
-#     img = np.zeros((1, 224, 224))
-#     model = models.CellposeModel()
-#     masks, _, _ = model.eval(img, diameter=50)
-#     assert masks.shape == (224, 224)
+@pytest.mark.slow
+def test_shape_2D_chan_diam_resize(cellposemodel_fixture_2D):
+    img = np.zeros((1, 224, 224))
+    masks, _, _ = cellposemodel_fixture_2D.eval(img, diameter=50)
+    assert masks.shape == (224, 224)
 
 
 def test_shape_2D_chan_last(cellposemodel_fixture_2D):
@@ -35,11 +34,11 @@ def test_shape_2D_chan_last(cellposemodel_fixture_2D):
 
 
 
-# def test_shape_2D_chan_specify():
-#     img = np.zeros((224, 224, 2))
-#     model = models.CellposeModel()
-#     masks, _, _ = model.eval(img, channel_axis=-1)
-#     assert masks.shape == (224, 224)
+@pytest.mark.slow
+def test_shape_2D_chan_specify(cellposemodel_fixture_2D):
+    img = np.zeros((224, 224, 2))
+    masks, _, _ = cellposemodel_fixture_2D.eval(img, channel_axis=-1)
+    assert masks.shape == (224, 224)
 
 
 def test_shape_2D_2chan_specify(cellposemodel_fixture_2D):
@@ -63,36 +62,32 @@ def test_shape_stitch(cellposemodel_fixture_3D):
     assert flows[2].shape == (5, 80, 80), 'cellprob shape mismatch'
 
 
-# def test_shape_3D():
-#     img = np.zeros((80, 80, 5, 1))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, channel_axis=3, z_axis=2, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D(cellposemodel_fixture_3D):
+    img = np.zeros((80, 80, 5, 1))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=3, z_axis=2, do_3D=True)
+    assert masks.shape == (5, 80, 80)
 
 
-# def test_shape_3D_1ch():
-#     img = np.zeros((5, 80, 80, 1))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, channel_axis=3, z_axis=0, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D_1ch(cellposemodel_fixture_3D):
+    img = np.zeros((5, 80, 80, 1))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=3, z_axis=0, do_3D=True)
+    assert masks.shape == (5, 80, 80)
 
 
-# def test_shape_3D_1ch_3ndim():
-#     img = np.zeros((5, 80, 80))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, channel_axis=None, z_axis=0, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D_1ch_3ndim(cellposemodel_fixture_3D):
+    img = np.zeros((5, 80, 80))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=None, z_axis=0, do_3D=True)
+    assert masks.shape == (5, 80, 80)
 
 
-# def test_shape_3D_1ch_3ndim_diam():
-#     img = np.zeros((5, 80, 80))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, channel_axis=None, diameter=50, z_axis=0, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D_1ch_3ndim_diam(cellposemodel_fixture_3D):
+    img = np.zeros((5, 80, 80))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=None, diameter=50, z_axis=0, do_3D=True)
+    assert masks.shape == (5, 80, 80)
 
 
 def test_shape_3D_2ch(cellposemodel_fixture_3D):
@@ -104,18 +99,16 @@ def test_shape_3D_2ch(cellposemodel_fixture_3D):
     assert flows[2].shape == (4, 80, 80), 'cellprob shape mismatch'
 
 
-# def test_shape_3D_rgb_diam():
-#     img = np.zeros((5, 80, 80, 3))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, diameter=50, channels=[0, 0],
-#                                     channel_axis=3, z_axis=0, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D_rgb_diam(cellposemodel_fixture_3D):
+    img = np.zeros((5, 80, 80, 3))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, diameter=50, channels=[0, 0],
+                                    channel_axis=3, z_axis=0, do_3D=True)
+    assert masks.shape == (5, 80, 80)
 
-# def test_shape_3D_rgb():
-#     img = np.zeros((5, 80, 80, 3))
-#     use_gpu = torch.cuda.is_available()
-#     model = models.CellposeModel(gpu=use_gpu)
-#     masks, _, _ = model.eval(img, channels=[0, 0],
-#                                     channel_axis=3, z_axis=0, do_3D=True)
-#     assert masks.shape == (5, 80, 80)
+@pytest.mark.slow
+def test_shape_3D_rgb(cellposemodel_fixture_3D):
+    img = np.zeros((5, 80, 80, 3))
+    masks, _, _ = cellposemodel_fixture_3D.eval(img, channels=[0, 0],
+                                    channel_axis=3, z_axis=0, do_3D=True)
+    assert masks.shape == (5, 80, 80)