Batch export of patches and annotated 2d images is now working

extensions/chaeo/annotators.py

@@ -3,8 +3,8 @@ from PIL import Image, ImageDraw, ImageFont
 
 from model_server.process import rescale
 
-def draw_boxes_on_2d_image(img, boxes, **kwargs):
-    pilimg = Image.fromarray(np.copy(img))  # drawing modifies array in-place
+def draw_boxes_on_2d_image(yx_img, boxes, **kwargs):
+    pilimg = Image.fromarray(np.copy(yx_img))  # drawing modifies array in-place
     draw = ImageDraw.Draw(pilimg)
     font_size = kwargs.get('font_size', 18)
     linewidth = kwargs.get('linewidth', 4)
@@ -29,15 +29,20 @@ def draw_boxes_on_2d_image(img, boxes, **kwargs):
     return pilimg
 
 
-def draw_boxes_on_3d_image(img, boxes, draw_full_depth=False, **kwargs):
-    annotated = np.zeros(img.shape, dtype=img.dtype)
+def draw_boxes_on_3d_image(yxcz_img, boxes, draw_full_depth=False, **kwargs):
+    assert len(yxcz_img.shape) == 4
+    nz = yxcz_img.shape[3]
+    assert yxcz_img.shape[2] == 1
+    assert nz > 1
 
-    for zi in range(0, img.shape[0]):
+    annotated = np.zeros(yxcz_img.shape, dtype=yxcz_img.dtype)
+
+    for zi in range(0, nz):
         if draw_full_depth:
             zi_boxes = boxes
         else:
             zi_boxes = [bb for bb in boxes if bb['info'].zi == zi]
-        annotated[zi, :, :] = draw_boxes_on_2d_image(img[zi, :, :], zi_boxes, **kwargs)
+        annotated[:, :, 0, zi] = draw_boxes_on_2d_image(yxcz_img[:, :, 0, zi], zi_boxes, **kwargs)
 
     if clip := kwargs.get('rescale_clip'):
         assert clip >= 0.0 and clip <= 1.0

extensions/chaeo/products.py

 from pathlib import Path
 
 import numpy as np
-from PIL import Image, ImageDraw
 
 from skimage.io import imsave
 from tifffile import imwrite

extensions/chaeo/workflows.py

+from pathlib import Path
+from typing import Dict
+
+from pydantic import BaseModel
+
+
+from extensions.ilastik.models import IlastikPixelClassifierModel
+from extensions.chaeo.annotators import draw_boxes_on_3d_image
+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 model_server.workflows import Timer
+
+def export_patches_from_multichannel_zstack(  # TODO: PyDantic model for arguments
+        input_zstack_path: Path,
+        px_model: IlastikPixelClassifierModel,
+        pxmap_threshold: float,
+        pixel_class: int,
+        zmask_channel: int,
+        patches_channel: int,
+        where_output: Path,
+        mask_type: str = 'boxes',
+        zmask_filters: Dict = None,
+        zmask_expand_box_by: int = None,
+) -> Dict:
+
+    ti = Timer()
+    stack = generate_file_accessor(input_zstack_path)
+    fstem = input_zstack_path.stem
+    ti.click('file_input')
+    assert stack.nz > 1, 'Expecting z-stack'
+
+    # MIP and classify pixels
+    mip = InMemoryDataAccessor(
+        stack.get_one_channel_data(channel=0).data.max(axis=-1, keepdims=True)
+    )
+    pxmap, _ = px_model.infer(mip)
+    ti.click('infer_pixel_probability')
+
+    write_accessor_data_to_file(
+        where_output / 'pixel_probabilities' / (fstem + '.tif'),
+        pxmap
+    )
+    ti.click('export_pixel_probability')
+
+    obmask = InMemoryDataAccessor(
+        pxmap.data > pxmap_threshold
+    )
+    ti.click('threshold_pixel_mask')
+
+    # make zmask
+    zmask, zmask_meta = build_zmask_from_object_mask(
+        obmask.get_one_channel_data(pixel_class),
+        stack.get_one_channel_data(zmask_channel),
+        mask_type=mask_type,
+        filters=zmask_filters,
+        expand_box_by=zmask_expand_box_by,
+    )
+    zmask_acc = InMemoryDataAccessor(zmask)
+    ti.click('generate_zmasks')
+
+    # export patches
+    files = export_patches_from_zstack(
+        where_output / '2d_patches',
+        stack.get_one_channel_data(patches_channel),
+        zmask_meta,
+        prefix=fstem,
+        draw_bounding_box=True,
+    )
+    ti.click('export_patches')
+
+    # export annotated zstack
+    annotated = InMemoryDataAccessor(
+        draw_boxes_on_3d_image(
+            stack.get_one_channel_data(patches_channel).data,
+            zmask_meta
+        )
+    )
+    write_accessor_data_to_file(
+        where_output / 'annotated_zstacks' / (fstem + '.tif'),
+        annotated
+    )
+    ti.click('export_annotated_zstack')
+
+    return {
+        'pixel_model_id': px_model.model_id,
+        'input_filepath': str(input_zstack_path),
+        'number_of_objects': len(zmask_meta),
+        'success': True,
+        'timer_results': ti.events,
+    }
\ No newline at end of file

extensions/ilastik/examples/batch_run_patches.py

+from pathlib import Path
+
+from extensions.chaeo.workflows import export_patches_from_multichannel_zstack
+from extensions.ilastik.models import IlastikPixelClassifierModel
+
+if __name__ == '__main__':
+    where_czi = Path(
+        'z:/rhodes/projects/proj0004-marine-photoactivation/data/exp0038/AutoMic/20230906-163415/Selection'
+    )
+    where_output = Path(
+        'c:/Users/rhodes/projects/proj0011-plankton-seg/exp0009/batch_output'
+    )
+    px_ilp = Path.home() / 'model-server' / 'ilastik' / 'AF405-bodies_boundaries.ilp'
+    px_model = IlastikPixelClassifierModel(
+        params={'project_file': px_ilp}
+    )
+    for ff in where_czi.iterdir():
+        print(ff)
+        if not ff.suffix.upper() == '.CZI':
+            continue
+
+        export_patches_from_multichannel_zstack(
+            input_zstack_path=where_czi/ff,
+            px_model=px_model,
+            pxmap_threshold=0.6,
+            pixel_class=0,
+            zmask_channel=0,
+            patches_channel=4,
+            where_output=where_output,
+            mask_type='boxes',
+            zmask_filters={'area': (1e2, 1e5)},
+            zmask_expand_box_by=(64, 3),
+        )
+