extension.py

# Copyright CNRS/Inria/UNS
# Contributor(s): Eric Debreuve (since 2019), Morgane Nadal (2020)
#
# eric.debreuve@cnrs.fr
#
# This software is governed by the CeCILL  license under French law and
# abiding by the rules of distribution of free software.  You can  use,
# modify and/ or redistribute the software under the terms of the CeCILL
# license as circulated by CEA, CNRS and INRIA at the following URL
# "http://www.cecill.info".
#
# As a counterpart to the access to the source code and  rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty  and the software's author,  the holder of the
# economic rights,  and the successive licensors  have only  limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using,  modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean  that it is complicated to manipulate,  and  that  also
# therefore means  that it is reserved for developers  and  experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and,  more generally, to use and operate it in the
# same conditions as regards security.
#
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.

from __future__ import annotations

import brick.processing.frangi3 as fg_
import brick.processing.map_labeling as ml_
import brick.processing.input as in_
from brick.component.glial_cmp import glial_cmp_t
from brick.general.type import array_t, site_h

from typing import Optional, Sequence, Tuple

import numpy as np_
import skimage.filters as fl_
import skimage.measure as ms_
import skimage.morphology as mp_
from scipy import ndimage as im_


class extension_t(glial_cmp_t):
    #
    # soma_uid: connected to a soma somewhere upstream
    #
    __slots__ = ("end_points", "scales", "soma_uid", "__cache__")

    def __init__(self):
        #
        super().__init__()
        self.end_points = None
        self.scales = None
        self.soma_uid = None
        self.__cache__ = None

    @classmethod
    def FromMap(cls, lmp: array_t, scales: array_t, uid: int) -> extension_t:
        #
        instance = cls()

        bmp = lmp == uid
        instance.InitializeFromMap(bmp, uid)
        end_point_map = cls.EndPointMap(bmp)
        instance.end_points = end_point_map.nonzero()
        instance.scales = scales[instance.sites]
        instance.__cache__ = {}

        return instance

    @property
    def is_unconnected(self) -> bool:
        #
        return self.soma_uid is None

    @property
    def end_points_as_array(self) -> array_t:
        #
        pty_name = 'end_points_as_array'
        if pty_name not in self.__cache__:
            self.__cache__[pty_name] = np_.array(self.end_points)

        return self.__cache__[pty_name]

    def EndPointsForSoma(
        self, soma_uid: int, influence_map: array_t
    ) -> Tuple[site_h, ...]:
        #
        ep_bmp = influence_map[self.end_points] == soma_uid  # bmp=boolean map
        if ep_bmp.any():
            end_point_idc = ep_bmp.nonzero()[0]
            end_points = self.end_points_as_array[:, end_point_idc]

            return tuple(zip(*end_points.tolist()))

        return ()

    def BackReferenceSoma(self, glial_cmp: glial_cmp_t) -> None:
        #
        if isinstance(glial_cmp, extension_t):
            self.soma_uid = glial_cmp.soma_uid
        else:
            self.soma_uid = glial_cmp.uid

    def __str__(self) -> str:
        #
        if self.extensions is None:
            n_extensions = 0
        else:
            n_extensions = self.extensions.__len__()

        return (
            f"Ext.{self.uid}, "
            f"sites={self.sites[0].__len__()}, "
            f"endpoints={self.end_points[0].__len__()}, "
            f"soma={self.soma_uid}, "
            f"extensions={n_extensions}"
        )

    @staticmethod
    def ExtensionContainingSite(
        extensions: Sequence[extension_t], site: site_h
    ) -> Optional[extension_t]:
        #
        for extension in extensions:
            if site in tuple(zip(*extension.sites)):
                return extension

        return None

    @staticmethod
    def EnhancedForDetection(
        image: array_t, in_parallel: bool = False
    ) -> Tuple[array_t, array_t]:
        #
        # import os.path as ph_

        # if ph_.exists("./__runtime__/frangi.npz"):
        #     print("/!\\ Reading from precomputed data file")
        #     loaded = np_.load("./frangi.npz")
        #     enhanced_img = loaded["enhanced_img"]
        #     scale_map = loaded["scale_map"]
        #
        #     return enhanced_img, scale_map

        preprocessed_img = im_.morphology.white_tophat(
            image, size=2, mode="constant", cval=0.0, origin=0
        )

        enhanced_img, scale_map = fg_.FrangiEnhancement(
            preprocessed_img,
            scale_range=(0.1, 3.1),
            scale_step=1.0,
            alpha=0.8,
            beta=0.5,
            frangi_c=500.0,
            bright_on_dark=True,
            in_parallel=in_parallel,
            method="c",
        )

        # np_.savez_compressed(
        #     "./runtime/frangi.npz", enhanced_img=enhanced_img, scale_map=scale_map
        # )

        return enhanced_img, scale_map

    @staticmethod
    def CoarseMap(image: array_t, low: float, high: float, selem: array_t) -> array_t:
        #
        result = __HysterisisImage__(image, low, high)
        result = __MorphologicalCleaning__(result, selem)

        return result

    @staticmethod
    def FilteredCoarseMap(map_: array_t, ext_min_area_c: int) -> array_t:
        #
        result = map_.copy()
        lmp = ms_.label(map_)

        for region in ms_.regionprops(lmp):
            if region.area <= ext_min_area_c:
                region_sites = (lmp == region.label).nonzero()
                result[region_sites] = 0
                lmp[region_sites] = 0

        return result, lmp

    @staticmethod
    def FineMapFromCoarseMap(coarse_map: array_t) -> array_t:
        #
        # Might contain True-voxels that could be removed w/o breaking connectivity
        result = mp_.skeletonize_3d(coarse_map.astype(np_.uint8, copy=False))

        return result.astype(np_.int8, copy=False)

    @staticmethod
    def EndPointMap(map_: array_t) -> array_t:
        #
        part_map = ml_.PartLMap(map_)
        result = part_map == 1

        return result.astype(np_.int8)


def __HysterisisImage__(image: array_t, low: float, high: float) -> array_t:
    #
    # low = 0.02
    # high = 0.04
    nonzero_sites = (image > 0).nonzero()
    nonzero_values = image[nonzero_sites]

    low = low * nonzero_values.min()
    high = high * image.max()
    # lowt = low*(max_image_f-min_image_f)+max_image_f
    # hight = high*(max_image_f- min_image_f)+min_image_f
    # lowt = (image_f >lowt).astype(int)
    # hight = (image_f <hight).astype(int)

    result = fl_.apply_hysteresis_threshold(image, low, high)
    result = result.astype(np_.int8, copy=False)

    return result


def __MorphologicalCleaning__(image: array_t, selem) -> array_t:
    #
    result = image.copy()

    for dep in range(result.shape[0]):
        result[dep, :, :] = mp_.closing(result[dep, :, :], selem)
        result[dep, :, :] = mp_.opening(result[dep, :, :], selem)

    return result