test_model.py

import unittest
from model_server.conf.testing import czifile
from model_server.base.accessors import CziImageFileAccessor
from model_server.base.models import DummySemanticSegmentationModel, DummyInstanceSegmentationModel, CouldNotLoadModelError

class TestCziImageFileAccess(unittest.TestCase):
    def setUp(self) -> None:
        self.cf = CziImageFileAccessor(czifile['path'])

    def test_instantiate_model(self):
        model = DummySemanticSegmentationModel(params=None)
        self.assertTrue(model.loaded)

    def test_instantiate_model_with_nondefault_kwarg(self):
        model = DummySemanticSegmentationModel(autoload=False)
        self.assertFalse(model.autoload, 'Could not override autoload flag in subclass of Model.')

    def test_raise_error_if_cannot_load_model(self):
        class UnloadableDummyImageToImageModel(DummySemanticSegmentationModel):
            def load(self):
                return False

        with self.assertRaises(CouldNotLoadModelError):
            mi = UnloadableDummyImageToImageModel()

    def test_dummy_pixel_segmentation(self):
        model = DummySemanticSegmentationModel()
        img = self.cf.get_one_channel_data(0)
        mask = model.label_pixel_class(img)

        w = czifile['w']
        h = czifile['h']

        self.assertEqual(
            mask.shape,
            (h, w, 1, 1),
            'Inferred image is not the expected shape'
        )

        self.assertEqual(
            mask.data[int(w/2), int(h/2)],
            255,
            'Middle pixel is not white as expected'
        )

        self.assertEqual(
            mask.data[0, 0],
            0,
            'First pixel is not black as expected'
        )
        return img, mask

    def test_dummy_instance_segmentation(self):
        img, mask = self.test_dummy_pixel_segmentation()
        model = DummyInstanceSegmentationModel()
        obmap = model.label_instance_class(img, mask)