test_ilastik.py

from pathlib import Path
from shutil import copyfile
import unittest

import numpy as np

from model_server.base.accessors import CziImageFileAccessor, generate_file_accessor, InMemoryDataAccessor, PatchStack, write_accessor_data_to_file
from model_server.base.api import app
from model_server.extensions.ilastik import models as ilm
from model_server.extensions.ilastik.pipelines import px_then_ob
from model_server.extensions.ilastik.router import router
from model_server.base.roiset import RoiSet, RoiSetMetaParams
from model_server.base.pipelines import segment
import model_server.conf.testing as conf

data = conf.meta['image_files']
output_path = conf.meta['output_path']
params = conf.meta['roiset']
czifile = conf.meta['image_files']['czifile']
ilastik_classifiers = conf.meta['ilastik_classifiers']

app.include_router(router)

def _random_int(*args):
    return np.random.randint(0, 2 ** 8, size=args, dtype='uint8')

class TestIlastikPixelClassification(unittest.TestCase):
    def setUp(self) -> None:
        self.cf = CziImageFileAccessor(czifile['path'])
        self.channel = 0
        self.model = ilm.IlastikPixelClassifierModel(
            params=ilm.IlastikPixelClassifierParams(project_file=ilastik_classifiers['px']['path'].__str__())
        )
        self.mono_image = self.cf.get_mono(self.channel)


    def test_raise_error_if_autoload_disabled(self):
        model = ilm.IlastikPixelClassifierModel(
            params=ilm.IlastikPixelClassifierParams(project_file=ilastik_classifiers['px']['path'].__str__()),
            autoload=False
        )
        w = 512
        h = 256

        input_img = InMemoryDataAccessor(data=np.random.rand(w, h, 1, 1))

        with self.assertRaises(AttributeError):
            mask = model.label_pixel_class(input_img)


    def test_run_pixel_classifier_on_random_data(self):
        w = 512
        h = 256

        input_img = InMemoryDataAccessor(data=np.random.rand(h, w, 1, 1))

        mask = self.model.label_pixel_class(input_img)
        self.assertEqual(mask.shape, (h, w, 1, 1))


    def test_run_pixel_classifier(self):
        self.assertEqual(self.mono_image.shape_dict['X'], czifile['w'])
        self.assertEqual(self.mono_image.shape_dict['Y'], czifile['h'])
        self.assertEqual(self.mono_image.shape_dict['C'], 1)
        self.assertEqual(self.mono_image.shape_dict['Z'], 1)

        mask = self.model.label_pixel_class(self.mono_image)

        self.assertTrue(mask.is_mask())
        self.assertEqual(mask.shape[0:2], self.cf.shape[0:2])
        self.assertEqual(mask.shape_dict['C'], 1)
        self.assertEqual(mask.shape_dict['Z'], 1)
        self.assertTrue(
            write_accessor_data_to_file(
                output_path / 'seg' / f'seg_{self.cf.fpath.stem}_ch{self.channel}.tif',
                mask
            )
        )

    def test_label_pixels_with_params(self):
        def _run_seg(tr, sig):
            mod = ilm.IlastikPixelClassifierModel(
                params=ilm.IlastikPixelClassifierParams(
                    project_file=ilastik_classifiers['px']['path'].__str__(),
                    px_prob_threshold=tr,
                    px_smoothing=sig,
                ),
            )
            mask = mod.label_pixel_class(self.mono_image)
            write_accessor_data_to_file(
                output_path / 'seg' / f'seg_tr{int(10*tr)}_sig{int(10*sig)}.tif',
                mask
            )
            return mask
        mask1 = _run_seg(0.5, 0.0)
        mask2 = _run_seg(0.5, 0.2)
        self.assertEqual(mask1.shape, mask2.shape)


    def test_pixel_classifier_enforces_input_shape(self):
        self.assertEqual(self.model.model_chroma, 1)
        self.assertEqual(self.model.model_3d, False)

        # correct data
        self.assertIsInstance(
            self.model.label_pixel_class(
                InMemoryDataAccessor(
                    _random_int(512, 256, 1, 1)
                )
            ),
            InMemoryDataAccessor
        )

        # raise except with input of multiple channels
        with self.assertRaises(ilm.IlastikInputShapeError):
            mask = self.model.label_pixel_class(
                InMemoryDataAccessor(
                    _random_int(512, 256, 3, 1)
                )
            )

        # raise except with input of multiple channels
        with self.assertRaises(ilm.IlastikInputShapeError):
            mask = self.model.label_pixel_class(
                InMemoryDataAccessor(
                    _random_int(512, 256, 1, 15)
                )
            )

    def test_ilastik_infer_pxmap_from_patchstack(self):

        def _r(h):
            return np.random.randint(0, 2 ** 8, size=(h, 512, 1, 1), dtype='uint8')

        acc = PatchStack([_r(256), _r(512), _r(256)])
        self.assertEqual(acc.hw, (512, 512))
        self.assertEqual(acc.iat(0, crop=True).hw, (256, 512))

        mask = self.model.label_patch_stack(acc)
        self.assertEqual(mask.dtype, bool)
        self.assertEqual(mask.chroma, 1)
        self.assertEqual(mask.hw, acc.hw)
        self.assertEqual(mask.nz, acc.nz)
        self.assertEqual(mask.count, acc.count)

        pxmap, _ = self.model.infer_patch_stack(acc)
        self.assertEqual(pxmap.dtype, float)
        self.assertEqual(pxmap.chroma, len(self.model.labels))
        self.assertEqual(pxmap.hw, acc.hw)
        self.assertEqual(pxmap.nz, acc.nz)
        self.assertEqual(pxmap.count, acc.count)

    def test_run_object_classifier_from_pixel_predictions(self):
        self.test_run_pixel_classifier()
        fp = czifile['path']
        model = ilm.IlastikObjectClassifierFromPixelPredictionsModel(
            params=ilm.IlastikParams(project_file=ilastik_classifiers['pxmap_to_obj']['path'].__str__())
        )
        mask = self.model.label_pixel_class(self.mono_image)
        objmap, _ = model.infer(self.mono_image, mask)

        self.assertTrue(
            write_accessor_data_to_file(
                output_path / f'obmap_{fp.stem}.tif',
                objmap,
            )
        )
        self.assertEqual(objmap.data.max(), 2)


    def test_run_object_classifier_from_segmentation(self):
        self.test_run_pixel_classifier()
        fp = czifile['path']
        model = ilm.IlastikObjectClassifierFromSegmentationModel(
            params=ilm.IlastikParams(project_file=ilastik_classifiers['seg_to_obj']['path'].__str__())
        )
        mask = self.model.label_pixel_class(self.mono_image)
        objmap = model.label_instance_class(self.mono_image, mask)

        self.assertTrue(
            write_accessor_data_to_file(
                output_path / f'obmap_from_seg_{fp.stem}.tif',
                objmap,
            )
        )
        self.assertEqual(objmap.data.max(), 2)

    def test_ilastik_pixel_classification_as_workflow(self):
        res = segment.segment_pipeline(
            accessors={
                'accessor': generate_file_accessor(czifile['path'])
            },
            models={
                'model': ilm.IlastikPixelClassifierModel(
                    params=ilm.IlastikPixelClassifierParams(
                        project_file=ilastik_classifiers['px']['path'].__str__()
                    ),
                ),
            },
            channel=0,
        )
        self.assertGreater(res.times['inference'], 0.1)


class TestServerTestCase(conf.TestServerBaseClass):
    app_name = 'tests.test_ilastik.test_ilastik:app'
    input_data = czifile


class TestIlastikOverApi(TestServerTestCase):
    def test_httpexception_if_incorrect_project_file_loaded(self):
        self.assertPutFailure(
            'ilastik/seg/load/',
            500,
            body={'project_file': 'improper.ilp'},
        )


    def test_load_ilastik_pixel_model(self):
        mid = self.assertPutSuccess(
            'ilastik/seg/load/',
            body={'project_file': str(ilastik_classifiers['px']['path'])},
        )['model_id']
        rl = self.assertGetSuccess('models')
        self.assertEqual(rl[mid]['class'], 'IlastikPixelClassifierModel')
        return mid

    def test_load_another_ilastik_pixel_model(self):
        self.test_load_ilastik_pixel_model()
        self.assertEqual(len(self.assertGetSuccess('models')), 1)
        self.assertPutSuccess(
            'ilastik/seg/load/',
            body={'project_file': str(ilastik_classifiers['px']['path']), 'duplicate': True},
        )
        self.assertEqual(len(self.assertGetSuccess('models')), 2)
        self.assertPutSuccess(
            'ilastik/seg/load/',
            body={'project_file': str(ilastik_classifiers['px']['path']), 'duplicate': False},
        )
        self.assertEqual(len(self.assertGetSuccess('models')), 2)

    def test_load_ilastik_pixel_model_with_params(self):
        params = {
            'project_file': str(ilastik_classifiers['px']['path']),
            'px_class': 0,
            'px_prob_threshold': 0.5
        }
        mid = self.assertPutSuccess(
            'ilastik/seg/load/',
            body=params,
        )['model_id']
        mods = self.assertGetSuccess('models')
        self.assertEqual(len(mods), 1)
        self.assertEqual(mods[mid]['params']['px_prob_threshold'], 0.5)


    def test_load_ilastik_pxmap_to_obj_model(self):
        mid = self.assertPutSuccess(
            'ilastik/pxmap_to_obj/load/',
            body={'project_file': str(ilastik_classifiers['pxmap_to_obj']['path'])},
        )['model_id']
        rl = self.assertGetSuccess('models')
        self.assertEqual(rl[mid]['class'], 'IlastikObjectClassifierFromPixelPredictionsModel')
        return mid

    def test_load_ilastik_model_with_model_id(self):
        nmid = 'new_model_id'
        rmid = self.assertPutSuccess(
            'ilastik/pxmap_to_obj/load/',
            query={
                'model_id': nmid,
            },
            body={
                'project_file': str(ilastik_classifiers['pxmap_to_obj']['path']),
            },
        )['model_id']
        self.assertEqual(rmid, nmid)

    def test_load_ilastik_seg_to_obj_model(self):
        mid = self.assertPutSuccess(
            'ilastik/seg_to_obj/load/',
            body={'project_file': str(ilastik_classifiers['seg_to_obj']['path'])},
        )['model_id']
        rl = self.assertGetSuccess('models')
        self.assertEqual(rl[mid]['class'], 'IlastikObjectClassifierFromSegmentationModel')
        return mid

    def test_ilastik_infer_pixel_probability(self):
        fname = self.copy_input_file_to_server()
        mid = self.test_load_ilastik_pixel_model()
        acc_id = self.assertPutSuccess(f'accessors/read_from_file/{fname}')

        self.assertPutSuccess(
            f'pipelines/segment',
            body={'model_id': mid, 'accessor_id': acc_id, 'channel': 0},
        )


    def test_ilastik_infer_px_then_ob(self):
        fname = self.copy_input_file_to_server()
        px_model_id = self.test_load_ilastik_pixel_model()
        ob_model_id = self.test_load_ilastik_pxmap_to_obj_model()

        in_acc_id = self.assertPutSuccess(f'accessors/read_from_file/{fname}')

        self.assertPutSuccess(
            'ilastik/pipelines/pixel_then_object_classification/infer/',
            body={
                'px_model_id': px_model_id,
                'ob_model_id': ob_model_id,
                'accessor_id': in_acc_id,
                'channel': 0,
            }
        )


class TestIlastikOnMultichannelInputs(TestServerTestCase):
    def setUp(self) -> None:
        super(TestIlastikOnMultichannelInputs, self).setUp()
        self.pa_px_classifier = ilastik_classifiers['px_color_zstack']['path']
        self.pa_ob_pxmap_classifier = ilastik_classifiers['ob_pxmap_color_zstack']['path']
        self.pa_ob_seg_classifier = ilastik_classifiers['ob_seg_color_zstack']['path']
        self.pa_input_image = data['multichannel_zstack_raw']['path']
        self.pa_mask = data['multichannel_zstack_mask3d']['path']


    def test_classify_pixels(self):
        img = generate_file_accessor(self.pa_input_image)
        self.assertGreater(img.chroma, 1)
        mod = ilm.IlastikPixelClassifierModel(ilm.IlastikPixelClassifierParams(project_file=self.pa_px_classifier.__str__()))
        pxmap = mod.infer(img)[0]
        self.assertEqual(pxmap.hw, img.hw)
        self.assertEqual(pxmap.nz, img.nz)
        return pxmap

    def test_classify_objects(self):
        pxmap = self.test_classify_pixels()
        img = generate_file_accessor(self.pa_input_image)
        mod = ilm.IlastikObjectClassifierFromPixelPredictionsModel(
            ilm.IlastikParams(project_file=self.pa_ob_pxmap_classifier.__str__())
        )
        obmap = mod.infer(img, pxmap)[0]
        self.assertEqual(obmap.hw, img.hw)
        self.assertEqual(obmap.nz, img.nz)

    def test_workflow(self):
        """
        Test calling pixel then object map classification pipeline function directly
        """
        def _call_workflow(channel):
            return px_then_ob.pixel_then_object_classification_pipeline(
                accessors={
                    'accessor': generate_file_accessor(self.pa_input_image)
                },
                models={
                    'px_model': ilm.IlastikPixelClassifierModel(
                        ilm.IlastikParams(project_file=self.pa_px_classifier.__str__()),
                    ),
                    'ob_model': ilm.IlastikObjectClassifierFromPixelPredictionsModel(
                        ilm.IlastikParams(project_file=self.pa_ob_pxmap_classifier.__str__()),
                    )
                },
                channel=channel,
            )

        with self.assertRaises(ilm.IlastikInputShapeError):
            _call_workflow(channel=0)
        res = _call_workflow(channel=None)
        acc_input = generate_file_accessor(self.pa_input_image)
        acc_obmap = res['ob_map']
        self.assertEqual(acc_obmap.hw, acc_input.hw)
        self.assertEqual(len(acc_obmap.unique()[1]), 3)


    def test_api(self):
        """
        Test calling pixel then object map classification pipeline over API
        """
        copyfile(
            self.pa_input_image,
            Path(self.assertGetSuccess('paths')['inbound_images']) / self.pa_input_image.name
        )

        in_acc_id = self.assertPutSuccess(f'accessors/read_from_file/{self.pa_input_image.name}')

        px_model_id = self.assertPutSuccess(
            'ilastik/seg/load/',
            body={'project_file': str(self.pa_px_classifier)},
        )['model_id']

        ob_model_id  = self.assertPutSuccess(
            'ilastik/pxmap_to_obj/load/',
            body={'project_file': str(self.pa_ob_pxmap_classifier)},
        )['model_id']

        # run the pipeline
        obmap_id = self.assertPutSuccess(
            'ilastik/pipelines/pixel_then_object_classification/infer/',
            body={
                'accessor_id': in_acc_id,
                'px_model_id': px_model_id,
                'ob_model_id': ob_model_id,
            }
        )['output_accessor_id']

        # save output object map to file and compare
        obmap_acc = self.get_accessor(obmap_id)
        self.assertEqual(obmap_acc.shape_dict['C'], 1)

        # compare dimensions to input image
        self.assertEqual(obmap_acc.hw, generate_file_accessor(self.pa_input_image).hw)


class TestIlastikObjectClassification(unittest.TestCase):
    def setUp(self):
        stack = generate_file_accessor(data['multichannel_zstack_raw']['path'])
        stack_ch_pa = stack.get_mono(conf.meta['roiset']['patches_channel'])
        seg_mask = generate_file_accessor(data['multichannel_zstack_mask2d']['path'])

        self.roiset = RoiSet.from_binary_mask(
            stack_ch_pa,
            seg_mask,
            params=RoiSetMetaParams(
                mask_type='boxes',
                filters={'area': {'min': 1e3, 'max': 1e4}},
                expand_box_by=(64, 2)
            )
        )

        self.classifier = ilm.IlastikObjectClassifierFromSegmentationModel(
            params=ilm.IlastikParams(project_file=ilastik_classifiers['seg_to_obj']['path'].__str__()),
        )
        self.raw = self.roiset.get_patches_acc()
        self.masks = self.roiset.get_patch_masks_acc()


    def test_classify_patches(self):
        res = self.classifier.label_patch_stack(self.raw, self.masks)
        self.assertEqual(res.count, self.roiset.count)
        res.export_pyxcz(output_path / 'res_patches.tif')
        for pi in range(0, res.count):  # assert that there is only one nonzero label per patch
            la, ct = np.unique(res.iat(pi).data, return_counts=True)
            self.assertEqual(np.sum(ct > 1), 2)  # exclude single-pixel anomaly
            self.assertEqual(la[0], 0)