Added second segmentation pipeline with specific z-stack processing

model_server/base/pipelines/segment_zproj.py

+from typing import Dict
+
+from fastapi import APIRouter
+
+from .segment import SegmentParams, SegmentRecord, segment_pipeline
+from .shared import call_pipeline, PipelineTrace
+from ..accessors import GenericImageDataAccessor
+from ..models import Model
+
+from pydantic import Field
+
+router = APIRouter(
+    prefix='/pipelines',
+    tags=['pipelines'],
+)
+
+
+class SegmentZStackParams(SegmentParams):
+    zi: int = Field(None, description='z coordinate to use on input stack; apply MIP if empty')
+
+
+class SegmentZStackRecord(SegmentRecord):
+    pass
+
+
+@router.put('/segment_zproj')
+def segment_zproj(p: SegmentZStackParams) -> SegmentZStackRecord:
+    """
+    Run a semantic segmentation model to compute a binary mask from a projected input zstack
+    """
+    return call_pipeline(segment_zproj(), p)
+
+
+def segment_zproj_pipeline(
+        accessors: Dict[str, GenericImageDataAccessor],
+        models: Dict[str, Model],
+        **k
+) -> PipelineTrace:
+    d = PipelineTrace(accessors.get('accessor'))
+
+    if isinstance(k.get('zi'), int):
+        assert 0 < k['zi'] < d.last.nz
+        d['mip'] = d.last.get_zi(k['zi'])
+    else:
+        d['mip'] = d.last.get_mip()
+    return segment_pipeline({'accessor': d.last}, models, **k)
+

tests/base/test_pipelines.py

@@ -2,23 +2,24 @@ from pathlib import Path
 import unittest
 
 from model_server.base.accessors import generate_file_accessor, write_accessor_data_to_file
-from model_server.base.pipelines import segment
+from model_server.base.pipelines import segment, segment_zproj
 
 import model_server.conf.testing as conf
 from tests.base.test_model import DummySemanticSegmentationModel
 
 czifile = conf.meta['image_files']['czifile']
+zstack = conf.meta['image_files']['tifffile']
 output_path = conf.meta['output_path']
 
 
-class TestSegmentationPipeline(unittest.TestCase):
+class TestSegmentationPipelines(unittest.TestCase):
     def setUp(self) -> None:
         self.model = DummySemanticSegmentationModel()
 
-    def test_call_pipeline_function(self):
+    def test_call_segment_pipeline(self):
         acc = generate_file_accessor(czifile['path'])
         trace = segment.segment_pipeline({'accessor': acc}, {'model': self.model}, channel=2, smooth=3)
-        outfp = output_path / 'classify_pixels.tif'
+        outfp = output_path / 'pipelines' / 'segment_binary_mask.tif'
         write_accessor_data_to_file(outfp, trace.last)
 
         import tifffile
@@ -48,4 +49,19 @@ class TestSegmentationPipeline(unittest.TestCase):
             output_path / 'pipelines' / 'segment_interm',
             prefix=czifile['name']
         )
-        self.assertTrue([ofp.stem.split('_')[-1] for ofp in interm_fps] == ['mono', 'inference', 'smooth'])
\ No newline at end of file
+        self.assertTrue([ofp.stem.split('_')[-1] for ofp in interm_fps] == ['mono', 'inference', 'smooth'])
+
+    def test_call_segment_zproj_pipeline(self):
+        acc = generate_file_accessor(zstack['path'])
+
+        trace1 = segment_zproj.segment_zproj_pipeline({'accessor': acc}, {'model': self.model}, channel=0, smooth=3, zi=4)
+        self.assertEqual(trace1.last.chroma, 1)
+        self.assertEqual(trace1.last.nz, 1)
+
+        trace2 = segment_zproj.segment_zproj_pipeline({'accessor': acc}, {'model': self.model}, channel=0, smooth=3)
+        self.assertEqual(trace2.last.chroma, 1)
+        self.assertEqual(trace2.last.nz, 1)
+
+        trace3 = segment_zproj.segment_zproj_pipeline({'accessor': acc}, {'model': self.model})
+        self.assertEqual(trace3.last.chroma, 1)  # still == 1: model returns a single channel regardless of input
+        self.assertEqual(trace3.last.nz, 1)