data_handling.py

"""
Created on Jan 3 12:30:00 2025

@author: Anna Grim
@email: anna.grim@alleninstitute.org

Routines for loading data during training and inference.

"""

from aind_exaspim_dataset_utils.s3_util import get_img_prefix
from bm4d import bm4d
from concurrent.futures import (
    as_completed, ProcessPoolExecutor, ThreadPoolExecutor,
)
from copy import deepcopy
from torch.utils.data import Dataset
from tqdm import tqdm

import fastremap
import numpy as np
import random
import tensorstore as ts
import torch

from aind_exaspim_image_compression.utils import img_util, util
from aind_exaspim_image_compression.utils.swc_util import Reader


# --- Custom Datasets ---
class TrainDataset(Dataset):
    """
    A PyTorch Dataset for sampling 3D patches from whole-brain images and
    applying the BM4D denoising algorithm. The dataset's __getitem__ method
    returns both the original and denoised patches. Optionally, the patch
    sampling maybe biased toward foreground regions.

    Attributes
    ----------
    """

    def __init__(
        self,
        patch_shape,
        anisotropy=(0.748, 0.748, 1.0),
        boundary_buffer=5000,
        foreground_sampling_rate=0.5,
        n_examples_per_epoch=300,
        normalization_percentiles=(0.5, 99.9),
        normalized_brightness_clip=8,
        prefetch_foreground_sampling=16,
        sigma_bm4d=16,
    ):
        # Call parent class
        super(TrainDataset, self).__init__()

        # Class attributes
        self.anisotropy = anisotropy
        self.boundary_buffer = boundary_buffer
        self.foreground_sampling_rate = foreground_sampling_rate
        self.n_examples_per_epoch = n_examples_per_epoch
        self.normalization_percentiles = normalization_percentiles
        self.normalized_brightness_clip = normalized_brightness_clip
        self.patch_shape = patch_shape
        self.prefetch_foreground_sampling = prefetch_foreground_sampling
        self.sigma_bm4d = sigma_bm4d
        self.swc_reader = Reader()

        # Data structures
        self.segmentations = dict()
        self.skeletons = dict()
        self.imgs = dict()

    # --- Ingest data ---
    def ingest_brain(self, brain_id, img_path, segmentation_path, swc_pointer):
        """
        Loads a brain image, label mask, and skeletons, then stores each in
        internal dictionaries.

        Parameters
        ----------
        brain_id : str
            Unique identifier for the brain corresponding to the image.
        img_path : str or Path
            Path to whole-brain image to be read.
        segmentation_path : str
            Path to segmentation.
        swc_path : str
            Path to SWC files.
        """
        self.imgs[brain_id] = img_util.read(img_path)
        self._load_segmentation(brain_id, segmentation_path)
        self._load_swcs(brain_id, swc_pointer)

    def _load_segmentation(self, brain_id, segmentation_path):
        """
        Reads a segmentation mask generated by Google Applied Sciences (GAS).

        Parameters
        ----------
        brain_id : str
            Unique identifier for the brain corresponding to the given path.
        segmentation_path : str
            Path to segmentation.
        """
        if segmentation_path:
            # Load image
            label_mask = ts.open(
                {
                    "driver": "neuroglancer_precomputed",
                    "kvstore": {
                        "driver": "gcs",
                        "bucket": "allen-nd-goog",
                        "path": segmentation_path,
                    },
                    "context": {
                        "cache_pool": {"total_bytes_limit": 1000000000},
                        "cache_pool#remote": {"total_bytes_limit": 1000000000},
                        "data_copy_concurrency": {"limit": 8},
                    },
                    "recheck_cached_data": "open",
                }
            ).result()

            # Permute axes to be consistent with raw image.
            label_mask = label_mask[ts.d["channel"][0]]
            label_mask = label_mask[ts.d[0].transpose[2]]
            label_mask = label_mask[ts.d[0].transpose[1]]
            self.segmentations[brain_id] = label_mask

    def _load_swcs(self, brain_id, swc_pointer):
        if swc_pointer:
            # Initializations
            swc_dicts = self.swc_reader.read(swc_pointer)
            n_points = np.sum([len(d["xyz"]) for d in swc_dicts])

            # Extract skeleton voxels
            if n_points > 0:
                start = 0
                skeletons = np.zeros((n_points, 3), dtype=np.int32)
                for swc_dict in swc_dicts:
                    end = start + len(swc_dict["xyz"])
                    skeletons[start:end] = self.to_voxels(swc_dict["xyz"])
                    start = end
                self.skeletons[brain_id] = skeletons

    # --- Sample Image Patches ---
    def __getitem__(self, dummy_input):
        """
        Returns a pair of noisy and BM4D-denoised image patches, normalized
        according to percentile-based scaling.

        Parameters
        ----------
        dummy_input : Any
            Dummy argument required by PyTorch's "Dataset" interface for
            indexing. Not used in the patch sampling procedure.

        Returns
        -------
        noise : numpy.ndarray
            Noisy image patch, normalized and clipped.
        denoised : numpy.ndarray
            Denoised image patch, normalized and clipped using the same scale
            as the noisy patch.
        (mn, mx) : Tuple[float]
            Lower and upper percentiles used for normalization.
        """
        # Get image patches
        brain_id = self.sample_brain()
        voxel = self.sample_voxel(brain_id)
        noise = self.read_patch(brain_id, voxel)
        mn, mx = np.percentile(noise, self.normalization_percentiles)
        denoised = bm4d(noise, self.sigma_bm4d)

        # Normalize image patches
        noise = self.normalize(noise, mn, mx)
        denoised = self.normalize(denoised, mn, mx)
        return noise, denoised, (mn, mx)

    def sample_brain(self):
        """
        Samples a brain ID from the loaded images.

        Returns
        -------
        brain_id : str
            Unique identifier of the sampled whole-brain.
        """
        return util.sample_once(self.imgs.keys())

    def sample_voxel(self, brain_id):
        """
        Samples a voxel from a brain volume, either foreground or interior.

        Parameters
        ----------
        brain_id : str
            Unique identifier of the sampled whole-brain.

        Returns
        -------
        Tuple[int]
            Voxel coordinate chosen according to the foreground or interior
            sampling strategy.
        """
        if random.random() < self.foreground_sampling_rate:
            return self.sample_foreground_voxel(brain_id)
        else:
            return self.sample_interior_voxel(brain_id)

    def sample_foreground_voxel(self, brain_id):
        """
        Samples a voxel likely to be part of the foreground of a neuron.

        Parameters
        ----------
        brain_id : str
            Unique identifier of a whole-brain.

        Returns
        -------
        Tuple[int]
            Voxel coordinate representing a likely foreground location.
        """
        if brain_id in self.skeletons and np.random.random() > 0.5:
            return self.sample_skeleton_voxel(brain_id)
        elif brain_id in self.segmentations:
            return self.sample_segmentation_voxel(brain_id)
        else:
            return self.sample_bright_voxel(brain_id)

    def sample_interior_voxel(self, brain_id):
        """
        Samples a random voxel coordinate from the interior of a 3D image
        volume, avoiding boundary regions.

        Parameters
        ----------
        brain_id : str
            Unique identifier of a whole-brain.

        Returns
        -------
        Tuple[int]
            Voxel coordinate sampled uniformly at random within the valid
            interior region of the image volume.
        """
        voxel = list()
        for s in self.imgs[brain_id].shape[2::]:
            upper = s - self.boundary_buffer
            voxel.append(random.randint(self.boundary_buffer, upper))
        return tuple(voxel)

    def sample_skeleton_voxel(self, brain_id):
        """
        Samples a voxel coordinate near a skeleton point.

        Parameters
        ----------
        brain_id : str
            Unique identifier of a whole-brain.

        Returns
        -------
        Tuple[int]
            Voxel coordinate near a skeleton point.
        """
        idx = random.randint(0, len(self.skeletons[brain_id]) - 1)
        shift = np.random.randint(0, 16, size=3)
        return tuple(self.skeletons[brain_id][idx] + shift)

    def sample_segmentation_voxel(self, brain_id):
        """
        Sample a voxel coordinate whose corresponding segmentation patch
        contains a sufficiently large object.

        Parameters
        ----------
        brain_id : str
            Identifier for the image volume which must be a key in
            "self.segmentations".

        Returns
        -------
        best_voxel : Tuple[int]
            Voxel coordinate whose patch contains a sufficiently large object
            or had the largest object after 5 * self.prefetch attempts.
        """
        best_volume = 0
        best_voxel = self.sample_interior_voxel(brain_id)
        cnt = 0
        with ThreadPoolExecutor() as executor:
            while best_volume < 1600:
                # Read random image patches
                pending = dict()
                for _ in range(self.prefetch_foreground_sampling):
                    voxel = self.sample_interior_voxel(brain_id)
                    thread = executor.submit(
                        self.read_precomputed_patch, brain_id, voxel
                    )
                    pending[thread] = voxel

                # Check if labels patch has large enough object
                for thread in as_completed(pending.keys()):
                    voxel = pending.pop(thread)
                    labels_patch = thread.result()
                    vals, cnts = fastremap.unique(
                        labels_patch, return_counts=True
                    )

                    if len(cnts) > 1:
                        volume = np.max(cnts[1:])
                        if volume > best_volume:
                            best_voxel = voxel
                            best_volume = volume

                # Check number of tries
                cnt += 1
                if cnt > 5:
                    break
        return best_voxel

    def sample_bright_voxel(self, brain_id):
        """
        Samples a voxel coordinate whose image patch is sufficiently bright.

        Parameters
        ----------
        brain_id : str
            Unique identifier of a whole-brain.

        Returns
        -------
        best_voxel : Tuple[int]
            Voxel coordinate whose patch is sufficiently bright or is the
            highest observed brightness after 4 * self.prefetch attempts.
        """
        best_brightness = 0
        best_voxel = self.sample_interior_voxel(brain_id)
        cnt = 0
        with ThreadPoolExecutor() as executor:
            while best_brightness < 1000:
                # Read random image patches
                pending = dict()
                for _ in range(self.prefetch_foreground_sampling):
                    voxel = self.sample_interior_voxel(brain_id)
                    thread = executor.submit(
                        self.read_patch, brain_id, voxel
                    )
                    pending[thread] = voxel

                # Check if image patch is bright enough
                for thread in as_completed(pending.keys()):
                    voxel = pending.pop(thread)
                    img_patch = thread.result()
                    brightness = np.sum(img_patch > 100)
                    if brightness > best_brightness:
                        best_voxel = voxel
                        best_brightness = brightness

                # Check number of tries
                cnt += 1
                if cnt > 5:
                    break
        return best_voxel

    # --- Helpers ---
    def __len__(self):
        """
        Gets the number of training examples used in each epoch.

        Returns
        -------
        int
            Number of training examples used in each epoch.
        """
        return self.n_examples_per_epoch

    def normalize(self, img, mn, mx):
        """
        Normalizes the given image using a percentile-based scheme and clips
        the max brightness.

        Parameters
        ----------
        img : numpy.ndarray
            Image to be normalized
        mn : float
            Lower percentile.
        mx : float
            Upper percentile

        Returns
        -------
        img : numpy.ndarray
            Normalized image.
        """
        img = (img - mn) / (mx - mn + 1e-8)
        img = np.clip(img, 0, self.normalized_brightness_clip)
        return img

    def read_patch(self, brain_id, center):
        """
        Reads an image patch from a Zarr array.

        Parameters
        ----------
        brain_id : str
            Unique identifier of the sampled brain.
        center : Tuple[int]
            Center of image patch to be read.

        Returns
        -------
        numpy.ndarray
            Image patch.
        """
        s = img_util.get_slices(center, self.patch_shape)
        return self.imgs[brain_id][(0, 0, *s)]

    def read_precomputed_patch(self, brain_id, center):
        """
        Reads an image patch from a Precomputed array.

        Parameters
        ----------
        brain_id : str
            Unique identifier of the sampled brain.
        center : Tuple[int]
            Center of image patch to be read.

        Returns
        -------
        numpy.ndarray
            Image patch.
        """
        s = img_util.get_slices(center, self.patch_shape)
        return self.segmentations[brain_id][s].read().result()

    def to_voxels(self, xyz_arr):
        """
        Converts 3D points from physical to voxel coordinates.

        Parameters
        ----------
        xyz_arr : numpy.ndarray
            Array with shape (n, 3) that contains 3D points.

        Returns
        -------
        numpy.ndarray
            3D Points converted to voxel coordinates.
        """
        for i in range(3):
            xyz_arr[:, i] = xyz_arr[:, i] / self.anisotropy[i]
        return np.flip(xyz_arr, axis=1).astype(int)


class ValidateDataset(Dataset):

    def __init__(
        self,
        patch_shape,
        normalization_percentiles=(0.5, 99.9),
        normalized_brightness_clip=8,
        sigma_bm4d=16,
    ):
        """
        Instantiates a ValidateDataset object.

        Parameters
        ----------
        patch_shape : Tuple[int]
            Shape of image patches to be extracted.
        normalization_percentiles : Tuple[float], optional
            Upper and lower percentiles used to normalize the input image.
            Default is (0.5, 99.5).
        normalized_brightness_clip : float, optional
            Brightness value used as an upper limit that normalized intensities
            are clipped to. Default is 10.
        sigma_bm4d : float, optional
            Smoothing parameter used in the BM4D denoising algorithm. Default
            is 16.
        """
        # Call parent class
        super(ValidateDataset, self).__init__()

        # Instance attributes
        self.normalization_percentiles = normalization_percentiles
        self.normalized_brightness_clip = normalized_brightness_clip
        self.patch_shape = patch_shape
        self.sigma_bm4d = sigma_bm4d

        # Data structures
        self.example_ids = list()
        self.imgs = dict()
        self.denoised = list()
        self.noise = list()
        self.mn_mxs = list()

    def __len__(self):
        """
        Counts the number of examples in the dataset.

        Returns
        -------
        int
            Number of examples in the dataset.
        """
        return len(self.example_ids)

    def ingest_brain(self, brain_id, img_path):
        """
        Loads a brain image and stores it in the internal image dictionary.

        Parameters
        ----------
        brain_id : str
            Unique identifier for the brain corresponding to the image.
        img_path : str or Path
            Path to whole-brain image to be read.
        """
        self.imgs[brain_id] = img_util.read(img_path)

    def ingest_example(self, brain_id, voxel):
        """
        Extracts, denoises, normalizes, and stores an image patch from a brain
        volume.

        Parameters
        ----------
        brain_id : hashable
            Unique identifier of the brain from which to extract the patch.
        voxel : Tuple[int]
            Voxel coordinates of the patch center in the brain volume.
        """
        # Get image patches
        noise = self.read_patch(brain_id, voxel)
        mn, mx = np.percentile(noise, self.normalization_percentiles)
        denoised = bm4d(noise, self.sigma_bm4d)

        # Normalize image patches
        noise = self.normalize(noise, mn, mx)
        denoised = self.normalize(denoised, mn, mx)

        # Store results
        self.example_ids.append((brain_id, voxel))
        self.noise.append(noise)
        self.denoised.append(denoised)
        self.mn_mxs.append((mn, mx))

    def __getitem__(self, idx):
        """
        Retrieves a single example from the dataset.

        Parameters
        ----------
        idx : int
            Index of the sample to retrieve.

        Returns
        -------
        noise : numpy.ndarray
            Noisy image patch at the given index.
        denoised : numpy.ndarray
            Corresponding denoised image patch.
        mn_mx : Tuple[int]
            Minimum and maximum values used for normalization of the image
            patches.
        """
        return self.noise[idx], self.denoised[idx], self.mn_mxs[idx]

    # --- Helpers ---
    def normalize(self, img, mn, mx):
        """
        Normalizes the given image using a percentile-based scheme and clips
        the max brightness.

        Parameters
        ----------
        img : numpy.ndarray
            Image to be normalized
        mn : float
            Lower percentile.
        mx : float
            Upper percentile

        Returns
        -------
        img : numpy.ndarray
            Normalized image.
        """
        img = (img - mn) / (mx - mn + 1e-8)
        img = np.clip(img, 0, self.normalized_brightness_clip)
        return img

    def read_patch(self, brain_id, center):
        """
        Reads an image patch from a Zarr array.

        Parameters
        ----------
        brain_id : str
            Unique identifier of the sampled brain.
        center : Tuple[int]
            Center of image patch to be read.

        Returns
        -------
        numpy.ndarray
            Image patch.
        """
        slices = img_util.get_slices(center, self.patch_shape)
        return self.imgs[brain_id][(0, 0, *slices)]


# --- Custom Dataloader ---
class DataLoader:
    """
    DataLoader that uses multithreading to fetch image patches from the cloud
    to form batches.

    Attributes
    ----------
    dataset : torch.utils.data.Dataset
        Dataset to iterated over.
    batch_size : int
        Number of examples in each batch.
    patch_shape : Tuple[int]
        Shape of image patch expected by the model.
    """

    def __init__(self, dataset, batch_size=16):
        """
        Instantiates a DataLoader object.

        Parameters
        ----------
        dataset : torch.utils.data.Dataset
            Dataset to iterated over.
        batch_size : int, optional
            Number of examples in each batch. Default is 16.
        """
        # Instance attributes
        self.dataset = dataset
        self.batch_size = batch_size
        self.patch_shape = dataset.patch_shape

    def __iter__(self):
        """
        Iterates over the dataset and yields batches of examples.

        Returns
        -------
        iterator
            Yields batches of examples.
        """
        for idx in range(0, len(self.dataset), self.batch_size):
            yield self._load_batch(idx)

    def _load_batch(self, start_idx):
        # Compute batch size
        n_remaining_examples = len(self.dataset) - start_idx
        batch_size = min(self.batch_size, n_remaining_examples)

        # Generate batch
        with ProcessPoolExecutor() as executor:
            # Assign processs
            processes = list()
            for idx in range(start_idx, start_idx + batch_size):
                processes.append(
                    executor.submit(self.dataset.__getitem__, idx)
                )

            # Process results
            shape = (batch_size, 1,) + self.patch_shape
            noise_patches = np.zeros(shape)
            denoised_patches = np.zeros(shape)
            mn_mxs = np.zeros((self.batch_size, 2))
            for i, process in enumerate(as_completed(processes)):
                noise, denoised, mn_mx = process.result()
                noise_patches[i, 0, ...] = noise
                denoised_patches[i, 0, ...] = denoised
                mn_mxs[i, :] = mn_mx
        return to_tensor(noise_patches), to_tensor(denoised_patches), mn_mxs


# --- Helpers ---
def init_datasets(
    brain_ids,
    img_paths_json,
    patch_shape,
    foreground_sampling_rate=0.5,
    n_train_examples_per_epoch=100,
    n_validate_examples=0,
    segmentation_prefixes_path=None,
    sigma_bm4d=16,
    swc_pointers=None
):
    # Initializations
    train_dataset = TrainDataset(
        patch_shape,
        foreground_sampling_rate=foreground_sampling_rate,
        n_examples_per_epoch=n_train_examples_per_epoch,
        sigma_bm4d=sigma_bm4d
    )
    val_dataset = ValidateDataset(patch_shape)

    # Read segmentation path lookup (if applicable)
    if segmentation_prefixes_path:
        segmentation_paths = util.read_json(segmentation_prefixes_path)
    else:
        segmentation_paths = dict()

    # Load data
    for brain_id in tqdm(brain_ids, desc="Load Data"):
        # Set image path
        img_path = get_img_prefix(brain_id, img_paths_json) + str(0)

        # Set segmentation path
        if brain_id in segmentation_paths:
            segmentation_path = segmentation_paths[brain_id]
        else:
            segmentation_path = None

        # Set SWC pointer
        if swc_pointers:
            swc_pointer = deepcopy(swc_pointers)
            swc_pointer["path"] += f"/{brain_id}/world"
        else:
            swc_pointer = None

        # Ingest data
        val_dataset.ingest_brain(brain_id, img_path)
        train_dataset.ingest_brain(
            brain_id, img_path, segmentation_path, swc_pointer
        )

    # Generate validation examples
    for _ in range(n_validate_examples):
        brain_id = train_dataset.sample_brain()
        voxel = train_dataset.sample_voxel(brain_id)
        val_dataset.ingest_example(brain_id, voxel)
    return train_dataset, val_dataset


def to_tensor(arr):
    """
    Converts the given NumPy array to a torch tensor.

    Parameters
    ----------
    arr : numpy.ndarray
        Array to be converted.

    Returns
    -------
    torch.Tensor
        Array converted to a torch tensor.
    """
    return torch.tensor(arr, dtype=torch.float)