fix pytest mock model to be same as cp model. Rename tests accordingly, add slow test to check full model with 2d grayscale img

Author: mrariden

conftest.py

@@ -56,7 +56,7 @@ def data_dir(image_names):
 
 
 @pytest.fixture()
-def cellposemodel_fixture_2D():
+def cellposemodel_fixture_24layer():
     """ This is functionally identical to CellposeModel but uses mock class """
     use_gpu = torch.cuda.is_available()
     use_mps = 'mps' if torch.backends.mps.is_available() else False
@@ -66,10 +66,12 @@ def cellposemodel_fixture_2D():
 
 
 @pytest.fixture()
-def cellposemodel_fixture_3D():
+def cellposemodel_fixture_2layer():
     """ This is only uses 2 transformer blocks for speed """
-    use_gpu = torch.cuda.is_available() # Turn of gpu for mac 3d
-    model = MockCellposeModel(2, gpu=use_gpu)
+    use_gpu = torch.cuda.is_available()
+    use_mps = 'mps' if torch.backends.mps.is_available() else False
+    gpu = use_gpu or use_mps
+    model = MockCellposeModel(n_keep_layers=2, gpu=gpu)
     yield model
 
 
@@ -107,11 +109,12 @@ def __init__(self, n_keep_layers=2, gpu=False):
         super().__init__(gpu=gpu)
 
         self.net = MockTransformer(n_keep_layers)
+        self.net.to(self.device)
         self.net.load_model(Path().home() / '.cellpose/models/cpsam', device=self.device)
 
-    def eval(self, x, **kwargs):
+    def eval(self, *args, **kwargs):
         tic = time.time()
-        res = super().eval(x, **kwargs)
+        res = super().eval(*args, **kwargs)
         toc = time.time()
 
         print(f'eval() time elapsed: {toc-tic}')

tests/test_output.py

@@ -1,4 +1,4 @@
-from cellpose import io, metrics, utils
+from cellpose import io, metrics, utils, models
 import pytest
 from subprocess import check_output, STDOUT
 import os
@@ -38,38 +38,53 @@ def clear_output(data_dir, image_names):
             os.remove(npy_output)
 
 
-
-def test_class_2D(data_dir, image_names, cellposemodel_fixture_2D):
+def test_class_2D_one_img(data_dir, image_names, cellposemodel_fixture_24layer):
     clear_output(data_dir, image_names)
+        
+    img_file = data_dir / '2D' / image_names[0]
+
+    img = io.imread_2D(img_file)
+    # flowps = io.imread(img_file.parent / (img_file.stem + "_cp4_gt_flowps.tif"))
+
+    masks_pred, _, _ = cellposemodel_fixture_24layer.eval(img, normalize=True)
+    io.imsave(data_dir / '2D' / (img_file.stem + "_cp_masks.png"), masks_pred)
+    # flowsp_pred = np.concatenate([flows_pred[1], flows_pred[2][None, ...]], axis=0)
+    # mse = np.sqrt((flowsp_pred - flowps) ** 2).sum()
+    # assert mse.sum() < 1e-8, f"MSE of flows is too high: {mse.sum()} on image {image_name}"
+    # print("MSE of flows is %f" % mse.mean())
+
+    compare_masks_cp4(data_dir, image_names[0], "2D")
+    # clear_output(data_dir, image_names)
+
 
+@pytest.mark.slow
+def test_class_2D_all_imgs(data_dir, image_names, cellposemodel_fixture_24layer):
+    clear_output(data_dir, image_names)
     for image_name in image_names:
-        
         img_file = data_dir / '2D' / image_name
 
         img = io.imread_2D(img_file)
         # flowps = io.imread(img_file.parent / (img_file.stem + "_cp4_gt_flowps.tif"))
 
-        masks_pred, _, _ = cellposemodel_fixture_2D.eval(img, normalize=True)
+        masks_pred, _, _ = cellposemodel_fixture_24layer.eval(img, normalize=True)
         io.imsave(data_dir / '2D' / (img_file.stem + "_cp_masks.png"), masks_pred)
         # flowsp_pred = np.concatenate([flows_pred[1], flows_pred[2][None, ...]], axis=0)
         # mse = np.sqrt((flowsp_pred - flowps) ** 2).sum()
         # assert mse.sum() < 1e-8, f"MSE of flows is too high: {mse.sum()} on image {image_name}"
         # print("MSE of flows is %f" % mse.mean())
 
-        break # Just test one image for now
-
-    compare_masks_cp4(data_dir, image_names[0], "2D")
+    compare_masks_cp4(data_dir, image_names, "2D")
     clear_output(data_dir, image_names)
 
 
 @pytest.mark.slow
-def test_cyto2_to_seg(data_dir, image_names, cellposemodel_fixture_2D):
+def test_cyto2_to_seg(data_dir, image_names, cellposemodel_fixture_24layer):
     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)
+    masks, flows, _ = cellposemodel_fixture_24layer.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)
@@ -79,13 +94,13 @@ def test_cyto2_to_seg(data_dir, image_names, cellposemodel_fixture_2D):
     clear_output(data_dir, image_names)
 
 
-def test_class_3D(data_dir, image_names_3d, cellposemodel_fixture_3D):
+def test_class_3D(data_dir, image_names_3d, cellposemodel_fixture_2layer):
     clear_output(data_dir, image_names_3d)
 
     for image_name in image_names_3d:
         img_file = data_dir / '3D' / image_name
         img = io.imread_3D(img_file)
-        masks_pred, flows_pred, _ = cellposemodel_fixture_3D.eval(img, do_3D=True, channel_axis=-1, z_axis=0)
+        masks_pred, flows_pred, _ = cellposemodel_fixture_2layer.eval(img, do_3D=True, channel_axis=-1, z_axis=0)
         # io.imsave(data_dir / "3D" / (img_file.stem + "_cp_masks.tif"), masks)
 
         assert img.shape[:-1] == masks_pred.shape, f'mask incorrect shape for {image_name}, {masks_pred.shape=}'
@@ -129,15 +144,15 @@ def test_cli_3D_diam(data_dir, image_names_3d):
 
 
 @pytest.mark.slow
-def test_outlines_list(data_dir, image_names, cellposemodel_fixture_2D):
+def test_outlines_list(data_dir, image_names, cellposemodel_fixture_24layer):
     """ 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)
+    masks, _, _ = cellposemodel_fixture_24layer.eval(img, diameter=30)
     outlines_single = utils.outlines_list(masks, multiprocessing=False)
     outlines_multi = utils.outlines_list(masks, multiprocessing=True)
 

tests/test_shape.py

@@ -3,56 +3,55 @@
 
 
 #################### 2D Tests ####################
-@pytest.mark.slow
-def test_shape_2D_grayscale(cellposemodel_fixture_2D):
+def test_shape_2D_grayscale(cellposemodel_fixture_24layer):
     img = np.zeros((224, 224))
-    masks, _, _ = cellposemodel_fixture_2D.eval(img)
+    masks, _, _ = cellposemodel_fixture_24layer.eval(img)
     assert masks.shape == (224, 224)
 
 
-def test_shape_2D_chan_first_diam_resize(cellposemodel_fixture_2D):
+def test_shape_2D_chan_first_diam_resize(cellposemodel_fixture_24layer):
     img = np.zeros((1, 224, 224))
-    masks, flows, _ = cellposemodel_fixture_2D.eval(img, diameter=50)
+    masks, flows, _ = cellposemodel_fixture_24layer.eval(img, diameter=50)
     assert masks.shape == (224, 224), 'mask shape mismatch'
     assert flows[1].shape == (2, 224, 224), 'dP shape mismatch'
     assert flows[2].shape == (224, 224), 'cellprob shape mismatch'
 
 
 @pytest.mark.slow
-def test_shape_2D_chan_diam_resize(cellposemodel_fixture_2D):
+def test_shape_2D_chan_diam_resize(cellposemodel_fixture_24layer):
     img = np.zeros((1, 224, 224))
-    masks, _, _ = cellposemodel_fixture_2D.eval(img, diameter=50)
+    masks, _, _ = cellposemodel_fixture_24layer.eval(img, diameter=50)
     assert masks.shape == (224, 224)
 
 
-def test_shape_2D_chan_last(cellposemodel_fixture_2D):
+def test_shape_2D_chan_last(cellposemodel_fixture_24layer):
     img = np.zeros((224, 224, 2))
-    masks, flows, _ = cellposemodel_fixture_2D.eval(img)
+    masks, flows, _ = cellposemodel_fixture_24layer.eval(img)
     assert masks.shape == (224, 224), 'mask shape mismatch'
     assert flows[1].shape == (2, 224, 224), 'dP shape mismatch'
     assert flows[2].shape == (224, 224), 'cellprob shape mismatch'
 
 
 
 @pytest.mark.slow
-def test_shape_2D_chan_specify(cellposemodel_fixture_2D):
+def test_shape_2D_chan_specify(cellposemodel_fixture_24layer):
     img = np.zeros((224, 224, 2))
-    masks, _, _ = cellposemodel_fixture_2D.eval(img, channel_axis=-1)
+    masks, _, _ = cellposemodel_fixture_24layer.eval(img, channel_axis=-1)
     assert masks.shape == (224, 224)
 
 
-def test_shape_2D_2chan_specify(cellposemodel_fixture_2D):
+def test_shape_2D_2chan_specify(cellposemodel_fixture_24layer):
     img = np.zeros((224, 5, 224))
-    masks, flows, _ = cellposemodel_fixture_2D.eval(img, channels=[2, 1], channel_axis=1)
+    masks, flows, _ = cellposemodel_fixture_24layer.eval(img, channels=[2, 1], channel_axis=1)
     assert masks.shape == (224, 224), 'mask shape mismatch'
     assert flows[1].shape == (2, 224, 224), 'dP shape mismatch'
     assert flows[2].shape == (224, 224), 'cellprob shape mismatch'
 
 
 #################### 3D Tests ####################
-def test_shape_stitch(cellposemodel_fixture_3D):
+def test_shape_stitch(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80, 2)) # 5 layer 3d input, 2 channels
-    masks, flows, _ = cellposemodel_fixture_3D.eval(img, channels=[0, 0],
+    masks, flows, _ = cellposemodel_fixture_2layer.eval(img, channels=[0, 0],
                                     stitch_threshold=0.9, 
                                     channel_axis=3, z_axis=0, 
                                     do_3D=False)
@@ -63,52 +62,52 @@ def test_shape_stitch(cellposemodel_fixture_3D):
 
 
 @pytest.mark.slow
-def test_shape_3D(cellposemodel_fixture_3D):
+def test_shape_3D(cellposemodel_fixture_2layer):
     img = np.zeros((80, 80, 5, 1))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=3, z_axis=2, do_3D=True)
+    masks, _, _ = cellposemodel_fixture_2layer.eval(img, channel_axis=3, z_axis=2, do_3D=True)
     assert masks.shape == (5, 80, 80)
 
 
 @pytest.mark.slow
-def test_shape_3D_1ch(cellposemodel_fixture_3D):
+def test_shape_3D_1ch(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80, 1))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=3, z_axis=0, do_3D=True)
+    masks, _, _ = cellposemodel_fixture_2layer.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_3ndim(cellposemodel_fixture_3D):
+def test_shape_3D_1ch_3ndim(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=None, z_axis=0, do_3D=True)
+    masks, _, _ = cellposemodel_fixture_2layer.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_diam(cellposemodel_fixture_3D):
+def test_shape_3D_1ch_3ndim_diam(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, channel_axis=None, diameter=50, z_axis=0, do_3D=True)
+    masks, _, _ = cellposemodel_fixture_2layer.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):
+def test_shape_3D_2ch(cellposemodel_fixture_2layer):
     img = np.zeros((80, 2, 80, 4))
 
-    masks, flows, _ = cellposemodel_fixture_3D.eval(img, z_axis=-1, channel_axis=1, do_3D=True)
+    masks, flows, _ = cellposemodel_fixture_2layer.eval(img, z_axis=-1, channel_axis=1, do_3D=True)
     assert masks.shape == (4, 80, 80), 'mask shape mismatch'
     assert flows[1].shape == (3, 4, 80, 80), 'dP shape mismatch'
     assert flows[2].shape == (4, 80, 80), 'cellprob shape mismatch'
 
 
 @pytest.mark.slow
-def test_shape_3D_rgb_diam(cellposemodel_fixture_3D):
+def test_shape_3D_rgb_diam(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80, 3))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, diameter=50, channels=[0, 0],
+    masks, _, _ = cellposemodel_fixture_2layer.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(cellposemodel_fixture_3D):
+def test_shape_3D_rgb(cellposemodel_fixture_2layer):
     img = np.zeros((5, 80, 80, 3))
-    masks, _, _ = cellposemodel_fixture_3D.eval(img, channels=[0, 0],
+    masks, _, _ = cellposemodel_fixture_2layer.eval(img, channels=[0, 0],
                                     channel_axis=3, z_axis=0, do_3D=True)
     assert masks.shape == (5, 80, 80)