Consolidated 2d and 3d patch exporting; tests pass

extensions/chaeo/products.py

-from PIL import Image, ImageDraw, ImageFont
+from pathlib import Path
 
-def draw_boxes_on_2d_image(img, boxes, **kwargs):
-    pilimg = Image.fromarray(np.copy(img))  # drawing modifies array in-place
-    draw = ImageDraw.Draw(pilimg)
-    font_size = kwargs.get('font_size', 18)
-    linewidth = kwargs.get('linewidth', 4)
+import numpy as np
+from PIL import Image, ImageDraw
 
-    draw.font = ImageFont.truetype(font="arial.ttf", size=font_size)
+from skimage.io import imsave
+from tifffile import imwrite
 
-    for box in boxes:
-        y0 = box['info'].y0
-        y1 = box['info'].y1
-        x0 = box['info'].x0
-        x1 = box['info'].x1
-        xm = round((x0 + x1) / 2)
+from model_server.accessors import GenericImageDataAccessor
+from model_server.process import pad, rescale, resample_to_8bit
 
-        la = box['info'].label
-        zi = box['info'].zi
+def _write_patch_to_file(where, fname, data):
+    ext = fname.split('.')[-1].upper()
+    where.mkdir(parents=True, exist_ok=True)
 
-        draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=linewidth)
+    if ext == 'PNG':
+        assert data.dtype == 'uint8', f'Invalid data type {data.dtype}'
+        assert data.shape[2] == 1, f'Cannot export multichannel images as PNGs; RGB not supported'
+        assert data.shape[3] == 1, f'Cannot export z-stacks as PNGs'
+        imsave(where / fname, data[:, :, 0, 0], check_contrast=False)
+        return True
 
-        if kwargs.get('add_label') is True:
-            draw.text((xm, y0), f'Z{zi:04d}-L{la:04d}', fill='white', anchor='mb')
+    elif ext in ['TIF', 'TIFF']:
+        zcyx = np.moveaxis(data, [3, 2, 0, 1], [0, 1, 2, 3])
+        imwrite(where / fname, zcyx, imagej=True)
+        return True
 
-    return pilimg
+    else:
+        raise Exception(f'Unsupported file extension: {ext}')
 
-def generate_patches(
-        desc, stack, boxes, rescale_clip=0.0,
-        pad_to=256,
-        proj=lambda x: x.max(axis=0),
+
+def export_patches_from_zstack(
+        where: Path,
+        stack: GenericImageDataAccessor,
+        zmask_meta: list,
+        rescale_clip: float = 0.0,
+        pad_to: int = 256,
+        make_3d: bool = False,
         prefix='patch',
         **kwargs
 ):
-    patch_dir = root / 'output' / 'patches' / desc
-    patch_dir.mkdir(parents=True, exist_ok=True)
+    assert stack.chroma == 1, 'Expecting monochromatic image data'
+    assert stack.nz > 1, 'Expecting z-stack'
+
+    exported = []
+    for mi in zmask_meta:
+        obj = mi['info']
+        sl = mi['slice']
+        rbb = mi['relative_bounding_box']
 
-    for box in boxes:
-        obj = box['info']
-        sl = box['slice']
-        rbb = box['relative_bounding_box']
+        if make_3d:
+            patch = stack.data[sl]
+        else:
+            patch = np.max(stack.data[sl], axis=3, keepdims=True)
 
-        patch = proj(stack[sl])
-        patch_fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}'
+        assert len(patch.shape) == 4
+        assert patch.shape[2] == stack.chroma
+        # assert patch.shape[3] == 1 # should not get to zstacks by this point
 
         if rescale_clip is not None:
             patch = rescale(patch, rescale_clip)
@@ -52,46 +66,17 @@ def generate_patches(
             x1 = rbb['x1']
             y1 = rbb['y1']
 
-            pilimg = Image.fromarray(patch)  # drawing modifies array in-place
-            draw = ImageDraw.Draw(pilimg)
-            draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=kwargs.get('linewidth', 1))
-            patch = np.array(pilimg)
+            for zi in range(0, patch.shape[3]):
+                pilimg = Image.fromarray(patch[:, :, 0, zi])  # drawing modifies array in-place
+                draw = ImageDraw.Draw(pilimg)
+                draw.rectangle([(x0, y0), (x1, y1)], outline='white', width=kwargs.get('linewidth', 1))
+                patch[:, :, 0, zi] = np.array(pilimg)
 
         if pad_to:
             patch = pad(patch, pad_to)
 
-        imsave(
-            patch_dir / (patch_fname + '.png'),
-            resample(patch),
-            check_contrast=False,
-        )
-    print(f'Successfully wrote {len(boxes)} patches to:\n{patch_dir}')
-
-
-def generate_3d_patches(  # in extensions.chaeo.products
-        desc, stack, boxes, rescale_clip=0.0,
-        pad_to=256,
-        prefix='patch',
-        proj=lambda x: x,
-):
-    patch_dir = root / 'output' / '3d_patches' / desc
-    patch_dir.mkdir(parents=True, exist_ok=True)
-
-    for box in boxes:
-        obj = box['info']
-        sl = box['slice']
-        patch = proj(stack[sl])
-        patch_fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}'
-
-        if rescale_clip is not None:
-            patch = rescale(patch, rescale_clip)
-
-        if pad_to:
-            patch = pad_3d(patch, pad_to)
-
-        imwrite(
-            patch_dir / (patch_fname + '.tif'),
-            patch,
-            imagej=True
-        )
-    print(f'Successfully wrote {len(boxes)} patches to:\n{patch_dir}')
\ No newline at end of file
+        ext = 'tif' if make_3d else 'png'
+        fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}.{ext}'
+        success = _write_patch_to_file(where, fname, resample_to_8bit(patch))
+        exported.append(fname)
+    return success, exported
\ No newline at end of file

extensions/chaeo/test_zstack.py → extensions/chaeo/tests/test_zstack.py

@@ -5,6 +5,7 @@ import numpy as np
 from conf.testing import output_path
 
 from extensions.chaeo.conf.testing import multichannel_zstack, pixel_classifier, pipeline_params
+from extensions.chaeo.products import export_patches_from_zstack
 from extensions.chaeo.zmask import build_zmask_from_object_mask
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from extensions.ilastik.models import IlastikObjectClassifierModel, IlastikPixelClassifierModel
@@ -51,6 +52,8 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         ar = meta[1]['info'].area
         self.assertGreaterEqual(sh[0] * sh[1], ar)
 
+        return zmask, meta
+
     def test_zmask_makes_correct_contours(self):
         return self.test_zmask_makes_correct_boxes(mask_type='contours')
 
@@ -58,4 +61,32 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         return self.test_zmask_makes_correct_boxes(filters={'area': (1e3, 1e4)})
 
     def test_zmask_makes_correct_expanded_boxes(self):
-        return self.test_zmask_makes_correct_boxes(expand_box_by=(64, 2))
\ No newline at end of file
+        return self.test_zmask_makes_correct_boxes(expand_box_by=(64, 2))
+
+    def test_make_2d_patches_from_zmask(self):
+        zmask, meta = self.test_zmask_makes_correct_boxes(
+            filters={'area': (1e3, 1e4)},
+            expand_box_by=(64, 2)
+        )
+        success, files = export_patches_from_zstack(
+            output_path / '2d_patches',
+            self.stack.get_one_channel_data(channel=1),
+            meta,
+            draw_bounding_box=True,
+        )
+        self.assertTrue(success)
+        self.assertGreaterEqual(len(files), 1)
+
+    def test_make_3d_patches_from_zmask(self):
+        zmask, meta = self.test_zmask_makes_correct_boxes(
+            filters={'area': (1e3, 1e4)},
+            expand_box_by=(64, 2),
+        )
+        success, files = export_patches_from_zstack(
+            output_path / '3d_patches',
+            self.stack.get_one_channel_data(channel=1),
+            meta,
+            make_3d=True)
+        self.assertTrue(success)
+        self.assertGreaterEqual(len(files), 1)
+

extensions/chaeo/zmask.py

@@ -5,7 +5,6 @@ from skimage.measure import find_contours, label, regionprops_table
 
 from model_server.accessors import GenericImageDataAccessor
 
-# build a single boolean 3d mask (objects v. bboxes) and return bounding boxes
 def build_zmask_from_object_mask(
         obmask: GenericImageDataAccessor,
         zstack: GenericImageDataAccessor,
@@ -97,7 +96,7 @@ def build_zmask_from_object_mask(
             'x1': ob.x1 - x0,
         }
 
-        sl = np.s_[y0: y1, x0: x1, 0, z0: z1 + 1]
+        sl = np.s_[y0: y1, x0: x1, :, z0: z1 + 1]
 
         # compute contours
         obmask = (lamap == ob.label)