siibra 0.5a2__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

siibra/volumes/{neuroglancer.py → providers/neuroglancer.py}

@@ -1,4 +1,4 @@
-# Copyright 2018-2021
+# Copyright 2018-2024
 # Institute of Neuroscience and Medicine (INM-1), Forschungszentrum Jülich GmbH
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,28 +12,66 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Handles reading and preparing neuroglancer files."""
 
-from . import volume
-
-from ..commons import logger, MapType, merge_meshes, NEUROGLANCER_MAX_GIB
-from ..retrieval import requests, cache
-from ..locations import boundingbox as _boundingbox
-
-from neuroglancer_scripts.precomputed_io import get_IO_for_existing_dataset
-from neuroglancer_scripts.accessor import get_accessor_for_url
+from . import provider as _provider
+
+from ...commons import (
+    logger,
+    MapType,
+    merge_meshes,
+    SIIBRA_MAX_FETCH_SIZE_GIB,
+    QUIET,
+    resample_img_to_img
+)
+from ...retrieval import requests, cache
+from ...locations import boundingbox as _boundingbox
+
+from neuroglancer_scripts.precomputed_io import get_IO_for_existing_dataset, PrecomputedIO
+from neuroglancer_scripts.http_accessor import HttpAccessor
 from neuroglancer_scripts.mesh import read_precomputed_mesh, affine_transform_mesh
 from io import BytesIO
 import nibabel as nib
 import os
 import numpy as np
 from typing import Union, Dict, Tuple
+import json
 
 
-class NeuroglancerProvider(volume.VolumeProvider, srctype="neuroglancer/precomputed"):
+def shift_ng_transfrom(
+    transform_nm: np.ndarray, scale_resolution_nm: np.ndarray, max_resolution_nm: np.ndarray
+) -> np.ndarray:
+    """
+    Helper method to get nifti standard affine.
+
+    transfrorm.json stored with neuroglancer precomputed images and meshes
+    are meant to be used for neuroglancer viewers and hence they are not
+    representative of the affine in other tools. This method shifts back
+    half a voxel in each axis.
+    (see https://neuroglancer-scripts.readthedocs.io/en/latest/neuroglancer-info.html#different-conventions-for-coordinate-transformations)
+
+    Parameters
+    ----------
+    transform_nm: np.ndarray
+        Transform array created for dispalying an image correctly from
+        neuroglancer precomputed format in neuroglancer viewer.
+    max_resolution_nm: np.ndarray
+        The voxel resolution of the highest level of resolution.
+
+    Returns
+    -------
+    np.ndarray
+        Standard affine in nm
+    """
+    scaling = np.diag(np.r_[scale_resolution_nm, 1.0])
+    affine = np.dot(transform_nm, scaling)
+    affine[:3, 3] += (max_resolution_nm * 0.5)
+    return affine
+
+
+class NeuroglancerProvider(_provider.VolumeProvider, srctype="neuroglancer/precomputed"):
 
     def __init__(self, url: Union[str, Dict[str, str]]):
-        volume.VolumeProvider.__init__(self)
+        _provider.VolumeProvider.__init__(self)
         self._init_url = url
         # TODO duplicated code to giftimesh
         if isinstance(url, str):  # one single image to load
@@ -63,10 +101,17 @@ class NeuroglancerProvider(volume.VolumeProvider, srctype="neuroglancer/precompu
             The name of a fragment volume to fetch, if any. For example,
             some volumes are split into left and right hemisphere fragments.
             See :func:`~siibra.volumes.Volume.fragments`
-        resolution_mm: float
-            Specify the resolution
+        resolution_mm: float, default: None (i.e, lowest)
+            Desired resolution in millimeters.
+            Tip
+            ---
+            Set to -1 to get the highest resolution. (might need to set max_bytes)
         voi: BoundingBox
             optional specification of a volume of interest to fetch.
+        max_bytes: float: Default: NeuroglancerVolume.MAX_BYTES
+            Maximum allowable size (in bytes) for downloading the image. siibra
+            will attempt to find the highest resolution image with a size less
+            than this value.
         """
 
         result = None
@@ -79,10 +124,11 @@ class NeuroglancerProvider(volume.VolumeProvider, srctype="neuroglancer/precompu
 
         if len(self._fragments) > 1:
             if fragment is None:
-                raise RuntimeError(
-                    f"Merging of fragments not yet implemented in {self.__class__.__name__}. "
-                    f"Specify one of [{', '.join(self._fragments.keys())}] using fetch(fragment=<name>). "
+                logger.info(
+                    f"Merging fragments [{', '.join(self._fragments.keys())}]. "
+                    f"You can select one using `fragment` kwarg."
                 )
+                result = self._merge_fragments(resolution_mm=resolution_mm, voi=voi, **kwargs)
             else:
                 matched_names = [n for n in self._fragments if fragment.lower() in n.lower()]
                 if len(matched_names) != 1:
@@ -113,81 +159,117 @@ class NeuroglancerProvider(volume.VolumeProvider, srctype="neuroglancer/precompu
                 label = None
             if label is not None:
                 result = nib.Nifti1Image(
-                    (result.get_fdata() == label).astype('uint8'),
-                    result.affine
+                    (np.asanyarray(result.dataobj) == label).astype('uint8'),
+                    result.affine,
+                    dtype='uint8',
                 )
 
         return result
 
-    @property
-    def boundingbox(self):
+    def get_boundingbox(self, **fetch_kwargs) -> "_boundingbox.BoundingBox":
         """
-        Return the bounding box in physical coordinates
-        of the union of fragments in this neuroglancer volume.
+        Return the bounding box in physical coordinates of the union of
+        fragments in this neuroglancer volume.
+
+        Parameters
+        ----------
+        fetch_kwargs:
+            key word arguments that are used for fetchin volumes,
+            such as voi or resolution_mm.
         """
         bbox = None
         for frag in self._fragments.values():
-            next_bbox = _boundingbox.BoundingBox((0, 0, 0), frag.shape, space=None) \
-                .transform(frag.affine)
+            if len(frag.shape) > 3:
+                logger.warning(
+                    f"N-D Neuroglancer volume has shape {frag.shape}, but "
+                    f"bounding box considers only {frag.shape[:3]}"
+                )
+            resolution_mm = fetch_kwargs.get("resolution_mm")
+            if resolution_mm is None:
+                affine = frag.affine
+                shape = frag.shape[:3]
+            else:
+                scale = frag._select_scale(resolution_mm=resolution_mm)
+                affine = scale.affine
+                shape = scale.size[:3]
+            next_bbox = _boundingbox.BoundingBox(
+                (0, 0, 0), shape, space=None
+            ).transform(affine)
             bbox = next_bbox if bbox is None else bbox.union(next_bbox)
         return bbox
 
-    def _merge_fragments(self) -> nib.Nifti1Image:
-        # TODO this only performs nearest neighbor interpolation, optimized for float types.
-        bbox = self.boundingbox
+    def _merge_fragments(
+        self,
+        resolution_mm: float = None,
+        voi: _boundingbox.BoundingBox = None,
+        **kwargs
+    ) -> nib.Nifti1Image:
+        with QUIET:
+            bbox = self.get_boundingbox(
+                clip=False,
+                background=0,
+                resolution_mm=resolution_mm,
+                voi=voi
+            )
+
         num_conflicts = 0
         result = None
-
-        for loader in self._img_loaders.values():
-            img = loader()
+        for frag_vol in self._fragments.values():
+            frag_scale = frag_vol._select_scale(
+                resolution_mm=resolution_mm,
+                bbox=voi,
+                max_bytes=kwargs.pop("maxbytes", NeuroglancerVolume.MAX_BYTES)
+            )
+            img = frag_scale.fetch(voi=voi)
             if result is None:
                 # build the empty result image with its own affine and voxel space
                 s0 = np.identity(4)
                 s0[:3, -1] = list(bbox.minpoint.transform(np.linalg.inv(img.affine)))
                 result_affine = np.dot(img.affine, s0)  # adjust global bounding box offset to get global affine
-                voxdims = np.asanyarray(bbox.transform(result_affine).shape, dtype="int")
+                voxdims = np.asanyarray(np.ceil(
+                    bbox.transform(np.linalg.inv(result_affine)).shape  # transform to the voxel space
+                ), dtype="int")
                 result_arr = np.zeros(voxdims, dtype=img.dataobj.dtype)
                 result = nib.Nifti1Image(dataobj=result_arr, affine=result_affine)
 
-            arr = np.asanyarray(img.dataobj)
-            Xs, Ys, Zs = np.where(arr != 0)
-            Xt, Yt, Zt, _ = np.split(
-                (np.dot(
-                    np.linalg.inv(result_affine),
-                    np.dot(img.affine, np.c_[Xs, Ys, Zs, Zs * 0 + 1].T)
-                ) + .5).astype('int'),
-                4, axis=0
-            )
-            num_conflicts += np.count_nonzero(result_arr[Xt, Yt, Zt])
-            result_arr[Xt, Yt, Zt] = arr[Xs, Ys, Zs]
+            # resample to merge template and update it
+            resampled_img = resample_img_to_img(source_img=img, target_img=result)
+            arr = np.asanyarray(resampled_img.dataobj)
+            nonzero_voxels = arr != 0
+            num_conflicts += np.count_nonzero(result_arr[nonzero_voxels])
+            result_arr[nonzero_voxels] = arr[nonzero_voxels]
 
         if num_conflicts > 0:
             num_voxels = np.count_nonzero(result_arr)
-            logger.warning(f"Merging fragments required to overwrite {num_conflicts} conflicting voxels ({num_conflicts/num_voxels*100.:2.1f}%).")
+            logger.warning(
+                f"Merging fragments required to overwrite {num_conflicts} "
+                f"conflicting voxels ({num_conflicts / num_voxels * 100.:2.3f}%)."
+            )
 
         return result
 
 
 class NeuroglancerVolume:
-    # Wether to keep fetched data in local cache
-    USE_CACHE = False
 
-    @property
-    def MAX_BYTES(self):
-        "Number of bytes at which an image array is considered to large to fetch"
-        return NEUROGLANCER_MAX_GIB * 1024 ** 3
+    USE_CACHE = False  # Whether to keep fetched data in local cache
+    MAX_BYTES = SIIBRA_MAX_FETCH_SIZE_GIB * 1024 ** 3  # Number of bytes at which an image array is considered to large to fetch
 
     def __init__(self, url: str):
         # TODO do we still need VolumeProvider.__init__ ? given it's not a subclass of VolumeProvider?
-        volume.VolumeProvider.__init__(self)
+        _provider.VolumeProvider.__init__(self)
         assert isinstance(url, str)
         self.url = url
         self._scales_cached = None
-        self._io = None
+        self._info = None
         self._transform_nm = None
+        self._io: PrecomputedIO = None
 
     @property
-    def transform_nm(self):
+    def transform_nm(self) -> np.ndarray:
+        """
+        This is the transformation matrix created to cater neuroglancer viewer
+        for a neuroglancer precomputed images.
+        """
         if self._transform_nm is not None:
             return self._transform_nm
         try:
@@ -207,12 +289,19 @@ class NeuroglancerVolume:
         self._transform_nm = val
 
     @property
-    def map_type(self):
+    def io(self) -> PrecomputedIO:
         if self._io is None:
+            accessor = HttpAccessor(self.url)
+            self._io = get_IO_for_existing_dataset(accessor)
+        return self._io
+
+    @property
+    def map_type(self):
+        if self._info is None:
             self._bootstrap()
         return (
             MapType.LABELLED
-            if self._io.info.get("type") == "segmentation"
+            if self._info.get("type") == "segmentation"
             else MapType.STATISTICAL
         )
 
@@ -220,21 +309,20 @@ class NeuroglancerVolume:
     def map_type(self, val):
         if val is not None:
             logger.debug(
-                "NeuroglancerVolume can determine its own maptype from self._io.info.get('type')"
+                "NeuroglancerVolume can determine its own maptype from self._info.get('type')"
             )
 
     def _bootstrap(self):
-        accessor = get_accessor_for_url(self.url)
-        self._io = get_IO_for_existing_dataset(accessor)
+        self._info = requests.HttpRequest(f"{self.url}/info", func=lambda b: json.loads(b.decode())).get()
         self._scales_cached = sorted(
-            [NeuroglancerScale(self, i) for i in self._io.info["scales"]]
+            [NeuroglancerScale(self, i) for i in self._info["scales"]]
         )
 
     @property
     def dtype(self):
-        if self._io is None:
+        if self._info is None:
             self._bootstrap()
-        return np.dtype(self._io.info["data_type"])
+        return np.dtype(self._info["data_type"])
 
     @property
     def scales(self):
@@ -252,19 +340,30 @@ class NeuroglancerVolume:
         # return the affine matrix of the scale 0 data
         return self.scales[0].affine
 
-    def fetch(self, resolution_mm: float = None, voi: _boundingbox.BoundingBox = None, **kwargs):
+    def fetch(
+        self,
+        resolution_mm: float = None,
+        voi: _boundingbox.BoundingBox = None,
+        max_bytes: float = MAX_BYTES,
+        **kwargs
+    ):
         # the caller has to make sure voi is defined in the correct reference space
-        scale = self._select_scale(resolution_mm=resolution_mm, bbox=voi)
-        return scale.fetch(voi, **kwargs)
+        scale = self._select_scale(resolution_mm=resolution_mm, bbox=voi, max_bytes=max_bytes)
+        return scale.fetch(voi=voi, **kwargs)
 
-    def get_shape(self, resolution_mm=None):
-        scale = self._select_scale(resolution_mm)
+    def get_shape(self, resolution_mm=None, max_bytes: float = MAX_BYTES):
+        scale = self._select_scale(resolution_mm=resolution_mm, max_bytes=max_bytes)
         return scale.size
 
     def is_float(self):
         return self.dtype.kind == "f"
 
-    def _select_scale(self, resolution_mm: float, bbox: _boundingbox.BoundingBox = None):
+    def _select_scale(
+        self,
+        resolution_mm: float,
+        max_bytes: float = MAX_BYTES,
+        bbox: _boundingbox.BoundingBox = None
+    ) -> 'NeuroglancerScale':
         if resolution_mm is None:
             suitable = self.scales
         elif resolution_mm < 0:
@@ -276,23 +375,31 @@ class NeuroglancerVolume:
             scale = suitable[-1]
         else:
             scale = self.scales[0]
-            logger.warning(
-                f"Requested resolution {resolution_mm}mm is not available for "
-                f"all axes. Falling back to the highest possible resolution of "
-                f"({', '.join(map('{:.4f}mm'.format, scale.res_mm))})."
-            )
+            xyz_res = ['{:.6f}'.format(r).rstrip('0') for r in scale.res_mm]
+            if all(r.startswith(str(resolution_mm)) for r in xyz_res):
+                logger.info(f"Closest resolution to requested is {', '.join(xyz_res)} mm.")
+            else:
+                logger.warning(
+                    f"Requested resolution {resolution_mm} is not available. "
+                    f"Falling back to the highest possible resolution of "
+                    f"{', '.join(xyz_res)} mm."
+                )
 
         scale_changed = False
-        while scale._estimate_nbytes(bbox) > self.MAX_BYTES:
+        while scale._estimate_nbytes(bbox) > max_bytes:
             scale = scale.next()
             scale_changed = True
             if scale is None:
                 raise RuntimeError(
                     f"Fetching bounding box {bbox} is infeasible "
-                    f"relative to the limit of {self.MAX_BYTES/1024**3}GiB."
+                    f"relative to the limit of {max_bytes / 1024**3}GiB."
                 )
         if scale_changed:
-            logger.warning(f"Resolution was reduced to {scale.res_mm} to provide a feasible volume size")
+            logger.warning(
+                f"Resolution was reduced to {scale.res_mm} to provide a "
+                f"feasible volume size of {max_bytes / 1024**3} GiB. Set `max_bytes` to"
+                f" fetch in the resolution requested."
+            )
         return scale
 
 
@@ -301,7 +408,7 @@ class NeuroglancerScale:
 
     color_warning_issued = False
 
-    def __init__(self, volume: NeuroglancerProvider, scaleinfo: dict):
+    def __init__(self, volume: NeuroglancerVolume, scaleinfo: dict):
         self.volume = volume
         self.chunk_sizes = np.array(scaleinfo["chunk_sizes"]).squeeze()
         self.encoding = scaleinfo["encoding"]
@@ -337,8 +444,8 @@ class NeuroglancerScale:
         result = self.volume.dtype.itemsize * bbox_.volume
         logger.debug(
             f"Approximate size for fetching resolution "
-            f"({', '.join(map('{:.2f}'.format, self.res_mm))}) mm "
-            f"is {result/1024**3:.2f} GiB."
+            f"({', '.join(map('{:.6f}'.format, self.res_mm))}) mm "
+            f"is {result / 1024**3:.5f} GiB."
         )
         return result
 
@@ -361,10 +468,13 @@ class NeuroglancerScale:
 
     @property
     def affine(self):
-        scaling = np.diag(np.r_[self.res_nm, 1.0])
-        affine = np.dot(self.volume.transform_nm, scaling)
-        affine[:3, :] /= 1e6
-        return affine
+        affine_ = shift_ng_transfrom(
+            transform_nm=self.volume.transform_nm,
+            scale_resolution_nm=self.res_nm,
+            max_resolution_nm=self.volume.scales[0].res_nm[0],
+        )
+        affine_[:3, :] /= 1e6
+        return affine_
 
     def _point_to_lower_chunk_idx(self, xyz):
         return (
@@ -395,7 +505,7 @@ class NeuroglancerScale:
         y0 = gy * self.chunk_sizes[1]
         z0 = gz * self.chunk_sizes[2]
         x1, y1, z1 = np.minimum(self.chunk_sizes + [x0, y0, z0], self.size)
-        chunk_czyx = self.volume._io.read_chunk(self.key, (x0, x1, y0, y1, z0, z1))
+        chunk_czyx = self.volume.io.read_chunk(self.key, (x0, x1, y0, y1, z0, z1))
         if channel is None:
             channel = 0
             if chunk_czyx.shape[0] > 1 and not self.color_warning_issued:
@@ -424,9 +534,9 @@ class NeuroglancerScale:
         for dim in range(3):
             if bbox_.shape[dim] < 1:
                 logger.warning(
-                    f"Bounding box in voxel space will be enlarged to voxel size 1 along axis {dim}."
+                    f"Bounding box in voxel space will be enlarged to by {self.res_mm[dim]} along axis {dim}."
                 )
-                bbox_.maxpoint[dim] = bbox_.maxpoint[dim] + 1
+                bbox_.maxpoint[dim] = bbox_.maxpoint[dim] + self.res_mm[dim]
 
         # extract minimum and maximum the chunk indices to be loaded
         gx0, gy0, gz0 = self._point_to_lower_chunk_idx(tuple(bbox_.minpoint))
@@ -448,8 +558,8 @@ class NeuroglancerScale:
         # determine the remaining offset from the "chunk mosaic" to the
         # exact bounding box requested, to cut off undesired borders
         data_min = np.array([gx0, gy0, gz0]) * self.chunk_sizes
-        x0, y0, z0 = (np.array(tuple(bbox_.minpoint)) - data_min).astype("int")
-        xd, yd, zd = np.array(bbox_.shape).astype("int")
+        x0, y0, z0 = (np.array(bbox_.minpoint) - data_min).astype("int")
+        xd, yd, zd = np.ceil((np.array(bbox_.maxpoint))).astype(int) - np.floor((np.array(bbox_.minpoint))).astype(int)
         offset = tuple(bbox_.minpoint)
 
         # build the nifti image
@@ -461,14 +571,14 @@ class NeuroglancerScale:
         )
 
 
-class NeuroglancerMesh(volume.VolumeProvider, srctype="neuroglancer/precompmesh"):
+class NeuroglancerMesh(_provider.VolumeProvider, srctype="neuroglancer/precompmesh"):
     """
     A surface mesh provided as neuroglancer precomputed mesh.
     """
 
     @staticmethod
     def _fragmentinfo(url: str) -> Dict[str, Union[str, np.ndarray, Dict]]:
-        """ Prepare basic mesh fragment information from url. """
+        """Prepare basic mesh fragment information from url."""
         return {
             "url": url,
             "transform_nm": np.array(requests.HttpRequest(f"{url}/transform.json").data),
@@ -490,10 +600,12 @@ class NeuroglancerMesh(volume.VolumeProvider, srctype="neuroglancer/precompmesh"
     def _url(self) -> Union[str, Dict[str, str]]:
         return self._init_url
 
-    @property
-    def boundingbox(self) -> _boundingbox.BoundingBox:
+    def get_boundingbox(self, clip=False, background=0.0, **fetch_kwargs) -> '_boundingbox.BoundingBox':
+        """
+        Bounding box calculation is not yet implemented for meshes.
+        """
         raise NotImplementedError(
-            f"Fast bounding box access to {self.__class__.__name__} objects not yet implemented."
+            f"Bounding box access to {self.__class__.__name__} objects not yet implemented."
         )
 
     def _get_fragment_info(self, meshindex: int) -> Dict[str, Tuple[str, ]]: