accessors.py

from pathlib import Path

import numpy as np

from model_server.accessors import generate_file_accessor, InMemoryDataAccessor

class MonoPatchStack(InMemoryDataAccessor):

    def __init__(self, data):
        """
        A sequence of n monochrome images of the same size
        :param data: either np.ndarray of dimensions YXn, or a list of np.ndarrays of size YX
        """

        if isinstance(data, np.ndarray): # interpret as YXZ
            assert data.ndim == 3
            self._data = np.expand_dims(data, 2)
        elif isinstance(data, list): # list of YX patches
            nda = np.array(data)
            assert nda.squeeze().ndim == 3
            self._data = np.expand_dims(
                np.moveaxis(
                    nda,
                    [1, 2, 0],
                    [0, 1, 2]),
                2
            )
        else:
            raise InvalidDataForPatchStackError(f'Cannot create accessor from {type(data)}')

    def make_tczyx(self):
        assert self.chroma == 1
        tyx = np.moveaxis(
            self.data[:, :, 0, :], # YX(C)Z
            [2, 0, 1],
            [0, 1, 2]
        )
        return np.expand_dims(tyx, (1, 2))

    @property
    def count(self):
        return self.nz

    def iat(self, i):
        return self.data[:, :, 0, i]

    def get_list(self):
        n = self.nz
        return [self.data[:, :, 0, zi] for zi in range(0, n)]


class MonoPatchStackFromFile(MonoPatchStack):
    def __init__(self, fpath):
        if not Path(fpath).exists():
            raise FileNotFoundError(f'Could not find {fpath}')
        self.file_acc = generate_file_accessor(fpath)
        super().__init__(self.file_acc.data[:, :, 0, :])


class Error(Exception):
    pass

class InvalidDataForPatchStackError(Error):
    pass

class FileNotFoundError(Error):
    pass