API References

class GPUContext(enum.Enum):
    CUPY = 1
    NUMPY = 2

class ItemDictNonSingleton():      
    def __init__(self, **kwargs):

        # Note For Developers check setattr otherwise youll include your variables as dict keys and this will mess the whole dict up
        self._index = 0
        self._dict = {}
        self._item_class = kwargs.get('item_class', Item)

        #Default Values for using the plugins system 
        self._module = 'tomobase'
        self._folder = 'plugins'    
        self._hook = 'default'

        reserved_keys = ['_index', '_dict', '_module', '_folder', '_hook', '_item_class', 'item_class']
        for key, value in kwargs.items():
            if key not in reserved_keys:
                newkey = key.upper()
                newkey = newkey.replace(' ', '_')
                if value is None:
                    value = self._index
                self._dict[newkey] = self._item_class(value, key)
                self._index += 1

    def __setattr__(self, key, value):
        if key in ['_index', '_dict', '_module', '_folder', '_hook', '_item_class']:
            # Allow these keys to be set as attributes
            super().__setattr__(key, value)
        else:
            # Add other keys to the internal dictionary
            self._dict[key] = value

    def __getattr__(self, key):
        try:
            return self._dict[key]
        except KeyError:
            logger.warning(f" object has no attribute '{key}'")

    def __getitem__(self, key):
        return self._dict[key]

    def __setitem__(self, key, value):
        if key in self._dict:
            pass
        else:
            if value is None:
                value = self._index
            newkey = key.upper()
            newkey = newkey.replace(' ', '_')
            self._dict[newkey] = self._item_class(value, key)
            self._index += 1

    def __len__(self):
        return self._index

    def loc(self, index):
        for key, item in self._dict.items():
            if item.value == index:
                return self._dict[key]
        logger.warning(f"Index {index} not found in the dictionary")

    def items(self):
        return self._dict.items()

    def key(self, index):
        for key, item in self._dict.items():
            if item.value == index:
                return key
        logger.warning(f"Index {index} not found in the dictionary")

    def append(self, **kwargs):
        for key, value in kwargs.items():
            self[key] = value

    def help(self):
        msg = "\n"
        for key, value in self._dict.items():
            msg += f"{Fore.BLUE}{key}{Style.RESET_ALL}: {value.name}, {value.value}\n"
        logger.info(msg)


    def update(self):
        spec = importlib.util.find_spec(self._module)
        if spec is None or spec.origin is None:
            raise ImportError(f"Cannot find the {self._module} package")

        path = os.path.dirname(spec.origin)
        tiltscheme_path = os.path.join(path, self._folder)
        for root, _, files in os.walk(tiltscheme_path):
            for filename in files:
                if filename.endswith('.py'):
                    module_path = os.path.relpath(os.path.join(root, filename), start=path)
                    module_name = self._module+'.'+ module_path.replace(os.sep, '.')[:-3]
                    module = importlib.import_module(module_name)
                    for name, obj in inspect.getmembers(module):
                        if inspect.isclass(obj) or inspect.isfunction(obj):
                            if hasattr(obj, self._hook):
                                self._update_item(obj)

    def _update_item(self, obj):
        self[obj.tomobase_name] = obj

class ItemDict(ItemDictNonSingleton):
    _instances = {}

    def __new__(cls, *args, **kwargs):
        if cls not in cls._instances:
            cls._instances[cls] = super(ItemDict, cls).__new__(cls)
        return cls._instances[cls]

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

class Item(object):

    def __init__(self, value, name):
        self._value = value    
        self._name = name

    def __call__(self, *args, **kwargs):
        if callable(self._value):
            return self._value(*args, **kwargs)
        else:
            logger.warning(f"Item {self._name} is not callable!")

    def __setitem__(self, name, value):
        if isinstance(self._value, dict) or isinstance(self._value, ItemDictNonSingleton):
            self._value[name] = value
        else:
            super().__setattr__(name, value)

    def __getitem__(self, key):
        if isinstance(self._value, dict) or isinstance(self._value, ItemDictNonSingleton):
            return self._value[key]
        else:
            super().__getattr__(key)

    @property
    def value(self):
        return self._value

    @property
    def name(self):
        return self._name

    @name.setter
    def name(self, value):
        self._name = value

    @value.setter
    def value(self, value):
        self._value = value

    def items(self):
        if isinstance(self._value, dict) or isinstance(self._value, ItemDictNonSingleton) or isinstance(self._value, ItemDict):
            return self._value.items()
        else:
            _dict = {}
            return _dict.items()

class Data(ABC):
    """
    Abstract base class for microscopy and tomography datasets. To implement a child of this class you must:

    - implement methods to read data from a file, these should be class methods
      that return an instance of the class

    - implement methods to write data to a file

    - create the class variables ``_readers`` and ``_writers`` which are
      dictionaries that link each supported filetype to the correct reader or
      writer respectively, the keys should be the extensions in lowercase

    Attributes:
        pixelsize (float): The size of the pixels in the dataset
        metadata (dict): A dictionary containing metadata about the dataset

    """


    def __init__(self, pixelsize: float=1.0, metadata: dict = {}):
        """Initialize the Data object

        Args:
            pixelsize (float, optional): The size of the pixels in the dataset. Defaults to 1.0 nm
            metadata (dict, optional): A dictionary containing metadata about the dataset. Defaults to {}.
        """
        self.pixelsize = pixelsize
        self.metadata = metadata

        # Note these are used to swap between numpy and cupy context
        self._context = GPUContext.NUMPY
        self._device = 0

    @classmethod
    def from_file(cls, filename: pathlib.Path | None = None, **kwargs):
        """Read a dataset from a file

        Args:
            filename (pathlib.Path | None): The name of the file from which to read the data, a GUI will be opened to select a file if none is specified (default: None)

        Exceptions:
            ValueError: If the file type is not supported or the file cannot be found

        Returns:
            Data: An instance of the Data class containing the read data
        """
        if filename is None:
            app = QApplication([])
            filename, _ = QFileDialog.getOpenFileName(None, "Select File", "", 
                                    [(f"{ext.upper()} files", f"*.{ext}") for ext in cls._readers.keys()])
            if not filename:
                raise Exception("No file selected or file could not be found")
            app.quit()

        _, ext = os.path.splitext(filename)
        ext = ext[1:]  # remove the dot

        try:
            reader = cls._readers[ext]
            reader(filename, **kwargs)
            data = reader(filename, **kwargs)
            return data
        except KeyError:
            raise ValueError(f"The given file type {ext.upper()} is not supported.")


    def to_file(self, filename: pathlib.Path | None = None, **kwargs):
        """Save the data to a file

        Arguments:
            filename (pathlib.Path | None): The name of the file to which to write the data, a GUI will be opened to select a file if none is specified (default: None)

        Exceptions:
            ValueError: If the file type is not supported or the file is not found
        """
        if filename is None:
            app = QApplication([])
            filename, _ = QFileDialog.getSaveFileName(None, "Save File", "", ";;".join([f"{ext.upper()} files (*.{ext})"
                                                                            for ext in self._writers.keys()]))
            if not filename:
                raise Exception("No file selected or file could not be found")
            app.quit()

        _, ext = os.path.splitext(filename)
        ext = ext[1:]  # remove the dot

        try:
            writer = self._writers[ext]
            writer(self, filename, **kwargs)
        except KeyError:
            raise ValueError(f"The given file type {ext.upper()} is not supported.")


    @property
    @abstractmethod
    def _writers(self):
        pass

    @property
    @abstractmethod
    def _readers(self):
        pass

    @classmethod
    def _get_type_id(cls) -> int:
        # Only use with napari used to identify the object from the layer etc
        class_name_upper = cls.__name__.upper()
        return TOMOBASE_DATATYPES[class_name_upper].value

    def _set_context(self, context:GPUContext | None = None, device:int | None = None):
        # used to set the context 
        if context is None:
            context = xp.context
        if device is None:
            device = xp.device

        self.data = xp.asarray(self.data, context, device)
        self._context = context
        self._device = device

    def layer_metadata(self, metadata: dict = {}):
        """Get the layer metadata in the format required for napari implementation

        Args:
            metadata (dict): the associated metadata for the layer

        Returns:
            dict: the layer metadata in the required format. By default the dict is wrapped in another dict of key 'ct metadata'. This avoids the metadata from conflicting with other packages.
        """

        meta = {}
        if 'ct metadata' not in metadata:
            meta['ct metadata'] = {}
        else:
            meta['ct metadata'] = metadata['ct metadata']

        for key, value in self.metadata.items():
            meta['ct metadata'][key] = value

        #must replace with the correct type in subclasses
        meta['ct metadata']['type'] = TOMOBASE_DATATYPES.DATA.value
        return meta

    def layer_attributes(self, attributes: dict = {}):
        """Get the layer attributes in the format required for napari implementation

        Args:
            attributes (dict): the associated attributes for the layer. 

        Returns:
            dict: the layer attributes in the required format. By default the structure is as follows:

        """
        attr = {}
        attr['name'] = attributes.get('name', 'Data')
        attr['scale'] = attributes.get('pixelsize' ,(self.pixelsize, self.pixelsize, self.pixelsize))
        attr['colormap'] = attributes.get('colormap', 'gray')
        attr['contrast_limits'] = attributes.get('contrast_limits', [0, np.max(self.data)*1.5])
        return attr

    def to_data_tuple(self, attributes:dict={}, metadata:dict={}):
        """ 
        Builds a Napari Layer Data Tuple

        Args:
            attributes (dict): the associated attributes for the layer. Defaults to {}.
            metadata (dict): the associated metadata for the layer. Defaults to {}.

        Returns:
            tuple: A tuple of (data, attributes, 'image') where data is the data array, attributes is a dictionary of attributes and 'image' is the layer type for napari.
        """
        attributes = self.layer_attributes(attributes)
        metadata = self.layer_metadata(metadata)
        attributes['metadata'] = metadata
        layerdata = (self.data, attributes, 'image')
        return layerdata

    @classmethod
    def from_data_tuple(cls, layer, attributes:dict|None=None):
        """Create an instance of the class from a Napari layer data tuple.

        Args:
            layer (napari layer): The Napari layer to convert.
            attributes (dict | None, optional): The associated attributes for the layer. Defaults to None.

        Returns:
           Data: An instance of the Data class.
        """

        if attributes is None:
            data = layer.data
            scale = layer.scale[0]
            layer_metadata = layer.metadata['ct metadata']
        else:
            data = layer
            scale = attributes['scale'][0]
            layer_metadata = attributes['metadata']['ct metadata']

        return cls(data, scale, layer_metadata)

class Image(Data):
    """ A class for single image datasets.

    Supported File Formats:
        - .png
        - .jpg
        - .jpeg
        - .bmp
        - .tif
        - .tiff

    Attributes:
        data (numpy.ndarray | cupy.ndarray): The image data. The data is indexed using the (rows, columns, channels) standard, which corresponds to (x, y, color)
    """

    def __init__(self, data, pixelsize: float = 1.0, metadata: dict = {}):
        """
        Initialize the Image object

        Args:
            data (numpy.ndarray) : the image data
            pixelsize (float): The size of the pixels in the dataset. Defaults to 1.0 nm
            metadata (dict): A dictionary containing metadata about the dataset. Defaults to {}.
        """

        self.data = data
        super().__init__(pixelsize, metadata=metadata)

    @staticmethod
    def _read_emi(filename, **kwargs):
        # TODO make this method independent of Hyperspy
        """
        content = hs.load(filename, lazy=False, reader='emi')
        data = np.transpose(np.asarray(content.data, dtype=float), (1, 0))   # With transpose we make sure that
                                                                # the orientation is correct in the
                                                                # case where the scanning rotation
                                                                # was set to -90 (which is the
                                                                # default of the FEI tomo software)
        pixelsize = content.axes_manager[0].scale  # Hyperspy uses nm
        im = Image(data, pixelsize)
        im.metadata['alpha_tilt'] = content.metadata.Acquisition_instrument.TEM.Stage.tilt_alpha
        return im
        """

    @staticmethod
    def _read_image(filename, **kwargs):
        return Image(xp.asarray(iio.imread(filename), dtype=float))

    def _write_image(self, filename, **kwargs):
        iio.imwrite(filename, self.data)

    def layer_metadata(self, metadata={}):
        meta = super().layer_metadata(metadata)
        meta['ct metadata']['type'] = TOMOBASE_DATATYPES.IMAGE.value
        meta['ct metadata']['axis'] = ['Signal', 'y', 'x'] if len(self.data.shape) == 3 else ['y', 'x']
        return meta

    def layer_attributes(self, attributes={}):
        attr = super().layer_attributes(attributes)
        attr['name'] = attr.get('name', 'Image')
        attr['scale'] = attr.get('pixelsize' ,(self.pixelsize, self.pixelsize))
        attr['contrast_limits'] = attr.get('contrast_limits', [0, np.max(self.data)*1.5])
        return attr


    _readers = {}
    _writers = {
        'png': _write_image,
        'bmp': _write_image,
        'tif': _write_image,
        'tiff': _write_image,
    }

class Sinogram(Data):
    """
    The sinogram is a stack of projection images, indexed using the
    (n, x, y) orientation. 

    Supported File Types:
        - .h5
        - .mrc
        - .emi
        - .mat (experimental)

    Attributes:
        data (numpy.ndarray): The sinogram data, indexed using the (n, x, y) orientation.
        angles (numpy.ndarray): The tilt angles in degrees corresponding to the projection images.
        times (numpy.ndarray): The times of acquisition corresponding to the projection images. This defaults to the projection index starting at 1. Otherwise it should be provided in seconds


    """

    def __init__(self, data, angles: np.ndarray, pixelsize: float = 1.0, times: np.ndarray | None = None, metadata: dict = {}):
        """Initialize a sinogram class

        Arguments:
            data (numpy.ndarray): The sinogram data, indexed using the (n, x, y) orientation.
            angles (numpy.ndarray): The tilt angles in degrees corresponding to the projection images.
            pixelsize (float): The width of the pixels in nanometer (default 1.0)
            times (numpy.ndarray | None): The times of acquisition corresponding to the projection images. This defaults to the projection index starting at 1. Otherwise it should be provided in seconds.
            metadata (dict): Additional metadata to store with the sinogram.

        """

        if len(angles) != data.shape[0]:
            raise ValueError(("There should be the same number of projection images as tilt angles."))
        if times is None:
            times = np.linspace(1, len(angles), len(angles))
        elif len(times) != len(angles):
            raise ValueError(("There should be the same number of projection images as times."))

        self.times = times
        self.data = data
        super().__init__(pixelsize, metadata)
        self.angles = np.asarray(angles)
        self.dim_default = 3

    def sort(self, bytime:bool = False):
        """
        Sort the sinogram by angles or by time

        Args:
            bytime (bool): Sort by time instead of angles
        """
        if bytime:
            indices = np.argsort(self.times)
            self.times = self.times[indices]
            self.angles = self.angles[indices] 
            self.data = self.data[indices,:,:]
        else:
            indices = np.argsort(self.angles)
            self.angles = self.angles[indices]
            self.times = self.times[indices]
            self.data = self.data[indices,:,:]

    def insert(self, img: np.ndarray, angle: float, time: float | None = None):
        """
        Insert a new image into the sinogram

        Args:
            img (numpy.ndarray): The image to insert
            angle (float): The angle of the image
            time (float): The time of acquisition
        """
        if time is None:
            time = self.times[-1] + 1

        if self.data.ndim == 2:
            # If `self.data` is 2D, add a new first axis and stack `img`
            self.data = np.stack((self.data, img), axis=0)
        else:
            # If `self.data` is already 3D, concatenate along the first axis
            self.data = np.concatenate(    (self.data, img), axis=0)
        #self.data = np.dstack((self.data, img))
        self.angles = np.append(self.angles, angle)
        self.times = np.append(self.times, time)

    def remove(self, index: int):
        """
        Remove an image from the sinogram

        Args:
            index (int): The index of the image to remove
        """
        self.data = np.delete(self.data, index, axis=0)
        self.angles = np.delete(self.angles, index)
        self.times = np.delete(self.times, index)

    @staticmethod
    def _read_h5(filename):
        f = h5py.File(filename, 'r')
        nt = len(f.keys())-2
        times = np.zeros(nt)
        angles = np.zeros(nt)
        for i in range(nt):
            key = 'image '+str(i)
            if i == 0:
                nx, ny = f[key]['HAADF'].shape
                data = xp.zeros([nx,ny,nt])
            data[:,:,i] = f[key]['HAADF']
            times[i] = np.array(f[key]['acquisition timee (s)']).item()
            angles[i] = np.array(f[key]['alpha tilt (deg)']).item()
        return Sinogram(data, angles, times=times)

    @staticmethod
    def _read_mrc(filename, **kwargs):
        data, metadata = mrcz.readMRC(filename)
        data = xp.asarray(data)
        pixelsize = metadata['pixelsize'][0]
        angles = metadata['angles']
        if 'times' in metadata:
            times = metadata['times']
        else:
            times = np.linspace(1, len(angles), len(angles)+1)
        return Sinogram(data, angles, pixelsize, times)


    @staticmethod
    def _read_mat(path):
        obj = loadmat(path)
        #if obj has a series or stack key

        if 'series' in obj:
            key = 'series'
        elif 'stack' in obj:
            key = 'stack'
        else:
            key = 'obj'


        d = obj[key]['data'][0][0]
        a = obj[key]['angles'][0][0]
        p = obj[key]['pixelsize'][0][0]
        if 'times' in obj[key]:
            t = obj[key]['times'][0][0]
        else:
            t = np.linspace(1, len(a), len(a)+1)

        ts = Sinogram(xp.asarray(d.squeeze()), a.squeeze(), p.squeeze(), t.squeeze())
        return ts

    def _write_mrc(self, filename, **kwargs):
        mrcz.writeMRC(self.data, filename, meta={'angles': self.angles, 'times': self.times}, pixelsize=[self.pixelsize, self.pixelsize, self.pixelsize])

    def _write_mat(self, filename, **kwargs):
        myrec = {'data':self.data, 'angles':self.angles, 'pixelsize':self.pixelsize, 'times':self.times} 
        savemat(filename, {'obj': myrec})

    @staticmethod
    def _read_emi_stack(filename, **kwargs):
        # List all EMI files in the directory of the selected file
        dirname = os.path.dirname(os.path.realpath(filename))
        filenames = glob.glob(dirname + "*.emi")
        # Read the first file for image size and metadata
        im = Image.from_file(filenames[0])
        data = np.zeros((len(filenames), im.data.shape[0], im.data.shape[1] ))
        angles = np.zeros(len(filenames))
        # Set the contents of the first file
        data[0, :, :] = im.data
        angles[0] = im.metadata['alpha_tilt']
        pixelsize = im.pixelsize
        # Loop over the other files to get all the projection images
        for i, filename in enumerate(filenames[1:], start=1):
            im = Image.from_file(filenames[0])
            data[i, :, :] = im.data
            angles[i] = im.metadata['alpha_tilt']
        # Sort the images by tilt angle and return the sinogram
        sorted_indices = np.argsort(angles)
        data = data[:, :, sorted_indices]
        angles = angles[sorted_indices]
        return Sinogram(data, angles, pixelsize)

    _readers = {}
    _writers = {
        'mrc': _write_mrc,
        'mat': _write_mat,
        'ali': _write_mrc,
    }



    def layer_attributes(self, attributes={}):
        attr = super().layer_attributes(attributes)
        attr['name'] = attributes.get('name', 'Sinogram')
        return attr

    def layer_metadata(self, metadata={}):
        meta = super().layer_metadata(metadata)
        meta['ct metadata']['type'] = TOMOBASE_DATATYPES.SINOGRAM.value
        meta['ct metadata']['angles'] = self.angles
        meta['ct metadata']['times'] = self.times
        meta['ct metadata']['axis'] = ['Projection', 'y', 'x'] if len(self.data.shape) == 3 else ['Projection', 'Signal', 'y', 'x']

        return meta

    @classmethod
    def from_data_tuple(cls, layer, attributes=None):
        if attributes is None:
            data = layer.data
            scale = layer.scale[0]
            times = layer.metadata['ct metadata']['times']
            angles = layer.metadata['ct metadata']['angles']
            layer_metadata = layer.metadata['ct metadata']
        else:
            data = layer
            scale = attributes['scale'][0]
            layer_metadata = attributes['metadata']['ct metadata']
            times =attributes['metadata']['ct metadata']['times']
            angles = attributes['metadata']['ct metadata']['angles']

        layer_metadata.pop('times', None)
        layer_metadata.pop('angles', None)

        return cls(data, angles, scale, times, layer_metadata)

class Volume(Data):
    """
    A 3D volume that is the result of a tomographic reconstruction. 

    Supported file formats:
        - .rec
        - .tiff

    Attributes:
        data (numpy.ndarray): The data represented by voxels. The data is indexed using  (y, x, z) notation

    """

    def __init__(self, data:np.ndarray, pixelsize:float=1.0, metadata:dict={}):
        """
        Initialize a volume object.

        Arguments:
            data (numpy.ndarray): The volume data.
            pixelsize (float): The size of the pixels in nanometers (default 1.0).
            metadata (dict): Additional metadata for the volume (default empty).
        """
        self.data = data
        super().__init__(pixelsize, metadata)


    @staticmethod
    def _read_rec(filename, normalize=True, **kwargs):
        with open(filename, 'rb') as f:
            # Data dimensions and type
            nx, ny, nz = np.fromfile(f, count=3, dtype='int32')
            datatype = np.fromfile(f, count=1, dtype='int32')
            if datatype == 0:
                datatype = 'uint8'
            elif datatype == 1:
                datatype = 'int16'
            elif datatype == 2:
                datatype = 'float32'
            elif datatype == 6:
                datatype = 'uint16'
            else:
                raise ValueError("Unsupported datatype in REC data.")

            # Pixel size in nm
            f.seek(10)
            cell_size = np.fromfile(f, count=1, dtype='int32')
            pixelsize = cell_size.astype('float32') / nx

            # Skip header
            f.seek(92)
            header_size = np.fromfile(f, count=1, dtype='int32')
            f.seek(1024 + header_size.item())

            # Read data
            data = np.fromfile(f, count=nx*ny*nz, dtype=datatype)
            data = np.reshape(data, [nx, ny, nz], order='F')
            data = np.transpose(data, (1, 0, 2))

            if normalize:
                return _rescale(Volume(data.astype(float), pixelsize=1.0))
            else:
                return Volume(data, pixelsize)

    def _write_rec(self, filename, normalize=True, **kwargs):
        # Convert data to (X, Y, Z)
        data = np.transpose(self.data, (1, 0, 2))

        # Create MRC header
        header = np.zeros(256, dtype='int32')
        header[:3] = data.shape  # Array dimensions
        if data.dtype == np.uint8 or normalize:
            header[3] = 0
        elif data.dtype == np.int16:
            header[3] = 1
        elif data.dtype == np.float32:
            header[3] = 2
        elif data.dtype == np.uint16:
            header[3] = 6
        else:
            raise TypeError("Unsupported data type for writing in REC file.")
        # Sampling along X, Y and Z. Same as array dimensions
        header[7:10] = data.shape
        # Physical dimensions in nm. Preserve float32 data type
        dimensions = self.pixelsize * np.array(data.shape, dtype='float32')
        header[10:13] = dimensions.view('int32')

        data = data.flatten(order='F')
        if normalize:
            data = data.astype(np.float32)
            data -= data.min()
            data *= 255 / data.max()
            data = data.astype('uint8')

        with open(filename, 'wb') as f:
            header.tofile(f)
            data.tofile(f)

    @staticmethod
    def _read_tiff(filename, **kwargs):
        raise NotImplementedError

    def _write_tiff(self, filename, **kwargs):
        raise NotImplementedError

    _readers = {}
    _writers = {
        'rec': _write_rec,
        'tif': _write_tiff,
        'tiff': _write_tiff,
    }


    def layer_attributes(self, attributes={}):
        attr = super().layer_attributes(attributes)
        attr['name'] = attributes.get('name', 'Volume')
        attr['scale'] = attributes.get('pixelsize', (self.pixelsize, self.pixelsize, self.pixelsize))
        attr['colormap'] = attributes.get('colormap', 'magma')
        attr['rendering'] = attributes.get('rendering', 'attenuated_mip')
        attr['contrast_limits'] = attributes.get('contrast_limits', [0, np.max(self.data)*1.5])
        return attr



    def layer_metadata(self, metadata={}):
        meta = super().layer_metadata(metadata)
        meta['ct metadata']['type'] = TOMOBASE_DATATYPES.VOLUME.value
        meta['ct metadata']['axis'] = ['z', 'y', 'x'] if len(self.data.shape) == 3 else ['z', 'Signal', 'y', 'x']

        return meta