Support single-channel annotated z-stack exports from RoiSet

model_server/base/annotators.py

@@ -4,19 +4,26 @@ from PIL import Image, ImageDraw, ImageFont
 from model_server.base.process import rescale
 
 def draw_boxes_on_3d_image(roiset, draw_full_depth=False, **kwargs):
-    _, _, chroma, nz = roiset.acc_raw.shape
+    h, w, chroma, nz = roiset.acc_raw.shape
     font_size = kwargs.get('font_size', 18)
     linewidth = kwargs.get('linewidth', 4)
 
-    annotated = np.zeros(roiset.acc_raw.shape, dtype=roiset.acc_raw.dtype)
+    if ck := kwargs.get('channel'):
+        channels = [ck]
+        shape = (h, w, 1, nz)
+    else:
+        channels = range(0, chroma)
+        shape = roiset.acc_raw.shape
+
+    annotated = np.zeros(shape, dtype=roiset.acc_raw.dtype)
 
     for zi in range(0, nz):
         if draw_full_depth:
             subset = roiset.get_df()
         else:
             subset = roiset.get_df().query(f'zi == {zi}')
-        for c in range(0, chroma):
-            pilimg = Image.fromarray(roiset.acc_raw.data[:, :, c, zi])
+        for ci in range(0, len(channels)):
+            pilimg = Image.fromarray(roiset.acc_raw.data[:, :, channels[ci], zi])
             draw = ImageDraw.Draw(pilimg)
             draw.font = ImageFont.truetype(font="arial.ttf", size=font_size)
 
@@ -26,7 +33,7 @@ def draw_boxes_on_3d_image(roiset, draw_full_depth=False, **kwargs):
 
                 if kwargs.get('draw_label') is True:
                     draw.text((xm, roi.y0), f'{roi.label:04d}', fill='white', anchor='mb')
-            annotated[:, :, c, zi] = pilimg
+            annotated[:, :, ci, zi] = pilimg
 
 
     if clip := kwargs.get('rescale_clip'):

tests/test_roiset.py

@@ -229,5 +229,13 @@ class TestRoiSetMultichannelProducts(BaseTestRoiSetMonoProducts, unittest.TestCa
         self.assertEqual(result.chroma, self.stack.chroma)
         self.assertEqual(result.nz, self.stack.nz)
 
-
+    def test_export_single_channel_annotated_zstack(self):
+        file = self.roiset.export_annotated_zstack(
+            output_path / 'annotated_zstack',
+            channel=3,
+        )
+        result = generate_file_accessor(Path(file['location']) / file['filename'])
+        self.assertEqual(result.hw, self.roiset.acc_raw.hw)
+        self.assertEqual(result.nz, self.roiset.acc_raw.nz)
+        self.assertEqual(result.chroma, 1)