test_zstack.py

import unittest

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.params import RoiSetExportParams, RoiSetMetaParams
from extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack, export_patch_masks_from_zstack
from extensions.chaeo.workflows import infer_object_map_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 IlastikPixelClassifierModel
from model_server.models import DummyInstanceSegmentationModel

class TestZStackDerivedDataProducts(unittest.TestCase):

    def setUp(self) -> None:

        # need test data incl obj map
        self.stack = generate_file_accessor(multichannel_zstack['path'])
        self.stack_ch_seg = self.stack.get_one_channel_data(pipeline_params['segmentation_channel'])
        self.stack_ch_pa = self.stack.get_one_channel_data(pipeline_params['patches_channel'])

        self.pxmodel = IlastikPixelClassifierModel(
            {'project_file': pixel_classifier['path']},
        )
        mip = InMemoryDataAccessor(
            self.stack_ch_seg.data.max(axis=-1, keepdims=True)
        )
        pxmap, _ = self.pxmodel.infer(mip)

        write_accessor_data_to_file(output_path / 'pxmap.tif', pxmap)
        self.seg_mask = InMemoryDataAccessor(
            pxmap.get_one_channel_data(
                pipeline_params['pxmap_channel']
            ).data > pipeline_params['pxmap_threshold']
        )
        write_accessor_data_to_file(output_path / 'seg_mask.tif', self.seg_mask)

    def test_zmask_makes_correct_boxes(self, mask_type='boxes', **kwargs):
        zmask, meta, df, interm = build_zmask_from_object_mask(
            self.seg_mask,
            self.stack_ch_pa,
            mask_type=mask_type,
            **kwargs,
        )
        zmask_acc = InMemoryDataAccessor(zmask)
        self.assertTrue(zmask_acc.is_mask())

        # assert dimensionality of zmask
        self.assertGreater(zmask_acc.shape_dict['Z'], 1)
        self.assertEqual(zmask_acc.shape_dict['C'], 1)
        write_accessor_data_to_file(output_path / 'zmask.tif', zmask_acc)

        # mask values are not just all True or all False
        self.assertTrue(np.any(zmask))
        self.assertFalse(np.all(zmask))

        # assert non-trivial meta info in boxes
        self.assertGreater(len(meta), 1)
        sh = meta[1]['mask'].shape
        ar = meta[1]['info'].area
        self.assertGreaterEqual(sh[0] * sh[1], ar)

        # assert dimensionality of intermediate data products
        self.assertEqual(interm['label_map'].shape, zmask.shape[0:2])
        self.assertEqual(interm['argmax'].shape, zmask.shape[0:2])

        return zmask, meta

    def test_zmask_works_on_non_zstacks(self, **kwargs):
        acc_zstack_slice = InMemoryDataAccessor(self.stack_ch_pa.data[:, :, :, 0])
        self.assertEqual(acc_zstack_slice.nz, 1)
        zmask, meta, df, interm = build_zmask_from_object_mask(
            self.seg_mask,
            acc_zstack_slice,
            mask_type='boxes',
            **kwargs,
        )
        zmask_acc = InMemoryDataAccessor(zmask)
        self.assertTrue(zmask_acc.is_mask())

    def test_zmask_makes_correct_contours(self):
        return self.test_zmask_makes_correct_boxes(mask_type='contours')

    def test_zmask_makes_correct_boxes_with_filters(self):
        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))

    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)
        )
        files = export_patches_from_zstack(
            output_path / '2d_patches',
            self.stack_ch_pa,
            meta,
            draw_bounding_box=True,
        )
        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),
        )
        files = export_patches_from_zstack(
            output_path / '3d_patches',
            self.stack_ch_pa,
            meta,
            make_3d=True)
        self.assertGreaterEqual(len(files), 1)

    def test_flatten_image(self):
        zmask, meta, df, interm = build_zmask_from_object_mask(
            self.seg_mask,
            self.stack_ch_pa,
            mask_type='boxes',
        )

        from extensions.chaeo.zmask import project_stack_from_focal_points

        dff = df[df['keeper']]

        img = project_stack_from_focal_points(
            dff['centroid-0'].to_numpy(),
            dff['centroid-1'].to_numpy(),
            dff['zi'].to_numpy(),
            self.stack,
            degree=4,
        )

        self.assertEqual(img.shape[0:2], self.stack.shape[0:2])

        write_accessor_data_to_file(
            output_path / 'flattened.tif',
            InMemoryDataAccessor(img)
        )

    def test_make_multichannel_2d_patches_from_zmask(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(128, 2)
        )
        files = export_multichannel_patches_from_zstack(
            output_path / '2d_patches_chlorophyl_bbox_overlay',
            InMemoryDataAccessor(self.stack.data),
            meta,
            ch_white=4,
            draw_bounding_box=True,
            bounding_box_channel=1,
        )
        self.assertGreaterEqual(len(files), 1)

    def test_make_multichannel_2d_patches_with_mask_overlay(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(128, 2)
        )
        files = export_multichannel_patches_from_zstack(
            output_path / '2d_patches_chlorophyl_mask_overlay',
            InMemoryDataAccessor(self.stack.data),
            meta,
            ch_white=4,
            ch_rgb_overlay=(3, None, None),
            draw_mask=True,
            mask_channel=0,
            overlay_gain=(0.1, 1.0, 1.0)
        )
        self.assertGreaterEqual(len(files), 1)

    def test_make_multichannel_2d_patches_with_contour_overlay(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(128, 2)
        )
        files = export_multichannel_patches_from_zstack(
            output_path / '2d_patches_chlorophyl_contour_overlay',
            InMemoryDataAccessor(self.stack.data),
            meta,
            ch_white=4,
            ch_rgb_overlay=(3, None, None),
            draw_contour=True,
            contour_channel=1,
            overlay_gain=(0.1, 1.0, 1.0)
        )
        self.assertGreaterEqual(len(files), 1)

    def test_make_binary_masks_from_zmask(self):
        zmask, meta = self.test_zmask_makes_correct_boxes(
            filters={'area': (1e3, 1e4)},
            expand_box_by=(128, 2)
        )
        files = export_patch_masks_from_zstack(
            output_path / '2d_mask_patches',
            InMemoryDataAccessor(self.stack.data),
            meta,
        )
        self.assertGreaterEqual(len(files), 1)

    def test_object_map_workflow(self):
        pp = pipeline_params
        models = [
            self.pxmodel,
            DummyInstanceSegmentationModel(),
        ]

        models = {
            'pixel_classifier': {
                'model': self.pxmodel,
                'params': {
                    'px_class': 0,
                    'px_prob_threshold': 0.6,
                }
            },
            'object_classifier': {
                'model': DummyInstanceSegmentationModel(),
            }
        }

        roi_params = RoiSetMetaParams(**{
            'mask_type': 'boxes',
            'filters': {
                'area': {'min': 1e3, 'max': 1e8}
            },
            'expand_box_by': [128, 2]
        })

        export_params = RoiSetExportParams(**{
            'pixel_probabilities': True,
            'patches_3d': {},
            'patches_2d_for_annotation': {
                'draw_bounding_box': True,
                'rgb_overlay_channels': [3, None, None],
                'rgb_overlay_weights': [0.2, 1.0, 1.0],
                'pad_to': 512,
            },
            'patches_2d_for_training': {
                'draw_bounding_box': False,
                'draw_mask': False,
            },
            'patch_masks': True,
            'annotated_zstacks': {},
            'object_classes': True,
            'dataframe': True,
        })

        infer_object_map_from_zstack(
            multichannel_zstack['path'],
            output_path / 'roiset' / 'workflow',
            models,
            segmentation_channel=pp['segmentation_channel'],
            patches_channel=pp['patches_channel'],
            export_params=export_params,
            roi_params=roi_params,
        )