test: bioimage.io core prediction task

tests/conftest.py

@@ -3,15 +3,42 @@
 import shutil
 from pathlib import Path
 
+import numpy as np
+import pooch
 import pytest
+import skimage.transform as skt
 import torch
 import yaml
 
+from plantseg.io.io import smart_load
+
 TEST_FILES = Path(__file__).resolve().parent / "resources"
 VOXEL_SIZE = (0.235, 0.15, 0.15)
 KEY_ZARR = "volumes/new"
 
 IS_CUDA_AVAILABLE = torch.cuda.is_available()
+CELLPOSE_TEST_IMAGE_RGB_3D = 'http://www.cellpose.org/static/data/rgb_3D.tif'
+
+
+@pytest.fixture
+def raw_zcyx_75x2x75x75(tmpdir) -> np.ndarray:
+    path_rgb_3d_75x2x75x75 = Path(pooch.retrieve(CELLPOSE_TEST_IMAGE_RGB_3D, path=tmpdir, known_hash=None))
+    return smart_load(path_rgb_3d_75x2x75x75)
+
+
+@pytest.fixture
+def raw_zcyx_96x2x96x96(raw_zcyx_75x2x75x75):
+    return skt.resize(raw_zcyx_75x2x75x75, (96, 2, 96, 96), order=1)
+
+
+@pytest.fixture
+def raw_cell_3d_100x128x128(raw_zcyx_75x2x75x75):
+    return skt.resize(raw_zcyx_75x2x75x75[:, 1], (100, 128, 128), order=1)
+
+
+@pytest.fixture
+def raw_cell_2d_96x96(raw_cell_3d_100x128x128):
+    return raw_cell_3d_100x128x128[48]
 
 
 @pytest.fixture

tests/functionals/prediction/test_prediction_biio.py

@@ -0,0 +1,18 @@
+import pytest
+
+from plantseg.functionals.prediction.prediction import biio_prediction
+
+
+@pytest.mark.parametrize(
+    "raw_fixture_name, input_layout, model_id",
+    (
+        ('raw_zcyx_96x2x96x96', 'ZCYX', 'philosophical-panda'),
+        ('raw_cell_3d_100x128x128', 'ZYX', 'emotional-cricket'),
+        ('raw_cell_2d_96x96', 'YX', 'pioneering-rhino'),
+    ),
+)
+def test_biio_prediction(raw_fixture_name, input_layout, model_id, request):
+    named_pmaps = biio_prediction(request.getfixturevalue(raw_fixture_name), input_layout, model_id)
+    for key, pmap in named_pmaps.items():
+        assert pmap is not None, f"Prediction map for {key} is None"
+        assert pmap.ndim == 4, f"Prediction map for {key} has {pmap.ndim} dimensions"

tests/tasks/test_prediction_tasks.py

@@ -3,7 +3,7 @@
 
 from plantseg.core.image import ImageLayout, ImageProperties, PlantSegImage, SemanticType
 from plantseg.io.voxelsize import VoxelSize
-from plantseg.tasks.prediction_tasks import unet_prediction_task
+from plantseg.tasks.prediction_tasks import biio_prediction_task, unet_prediction_task
 
 
 @pytest.mark.parametrize(
@@ -13,7 +13,7 @@
         ((64, 64), ImageLayout.YX, 'confocal_2D_unet_ovules_ds2x'),
     ],
 )
-def test_unet_prediction(shape, layout, model_name):
+def test_unet_prediction_task(shape, layout, model_name):
     mock_data = np.random.rand(*shape).astype('float32')
 
     property = ImageProperties(
@@ -25,7 +25,12 @@ def test_unet_prediction(shape, layout, model_name):
     )
     image = PlantSegImage(data=mock_data, properties=property)
 
-    result = unet_prediction_task(image=image, model_name=model_name, model_id=None, device='cpu')
+    result = unet_prediction_task(
+        image=image,
+        model_name=model_name,
+        model_id=None,
+        device='cpu',
+    )
 
     assert len(result) == 1
     result = result[0]
@@ -34,3 +39,32 @@ def test_unet_prediction(shape, layout, model_name):
     assert result.image_layout == property.image_layout
     assert result.voxel_size == property.voxel_size
     assert result.shape == mock_data.shape
+
+
+@pytest.mark.parametrize(
+    "raw_fixture_name, input_layout, model_id",
+    (
+        ('raw_zcyx_96x2x96x96', 'ZCYX', 'philosophical-panda'),
+        ('raw_cell_3d_100x128x128', 'ZYX', 'emotional-cricket'),
+        ('raw_cell_2d_96x96', 'YX', 'pioneering-rhino'),
+    ),
+)
+def test_biio_prediction_task(raw_fixture_name, input_layout, model_id, request):
+    image = PlantSegImage(
+        data=request.getfixturevalue(raw_fixture_name),
+        properties=ImageProperties(
+            name='test',
+            voxel_size=VoxelSize(voxels_size=(1.0, 1.0, 1.0), unit='um'),
+            semantic_type=SemanticType.RAW,
+            image_layout=input_layout,
+            original_voxel_size=VoxelSize(voxels_size=(1.0, 1.0, 1.0), unit='um'),
+        ),
+    )
+    result = biio_prediction_task(
+        image=image,
+        model_id=model_id,
+        suffix="_biio_prediction",
+    )
+    for new_image in result:
+        assert new_image.semantic_type == SemanticType.PREDICTION
+        assert '_biio_prediction' in new_image.name