Merge pull request #8 from AllenNeuralDynamics/refactor-training

feat: brightness biased sampling

Author: anna-grim

pyproject.toml

@@ -29,6 +29,7 @@ dependencies = [
     'scikit-learn',
     'scipy',
     'tensorboard',
+    'tensorstore',
     'torch',
     'torchvision',
     'tqdm',

src/aind_exaspim_image_compression/inference.py

@@ -123,7 +123,7 @@ def predict_patch(patch, model, normalization_percentiles=[5, 99.9]):
         Denoised 3D patch with the same shape as input patch.
     """
     # Run model
-    assert len(normalization_percentiles) == 2, "Must provide two percentiles" 
+    assert len(normalization_percentiles) == 2, "Must provide two percentiles"
     mn, mx = np.percentile(patch, normalization_percentiles)
     patch = to_tensor((patch - mn) / max(mx, 1))
     with torch.no_grad():
@@ -134,13 +134,20 @@ def predict_patch(patch, model, normalization_percentiles=[5, 99.9]):
     return np.abs(pred[0, 0, ...] * mx + mn).astype(np.uint16)
 
 
-def _predict_batch(img, model, starts, patch_size, trim=5):
+def _predict_batch(
+    img,
+    model,
+    starts,
+    patch_size,
+    normalization_percentiles=[5, 99.9],
+    trim=5,
+):
     # Subroutine
     def read_patch(i):
         start = starts[i]
         end = [min(s + patch_size, d) for s, d in zip(start, (D, H, W))]
         patch = img[0, 0, start[0]:end[0], start[1]:end[1], start[2]:end[2]]
-        mn, mx = np.percentile(patch, [5, 99.9])
+        mn, mx = np.percentile(patch, normalization_percentiles)
         patch = add_padding((patch - mn) / max(mx, 1), patch_size)
         return i, patch.astype(np.float32), (mn, mx)
 

src/aind_exaspim_image_compression/machine_learning/data_handling.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import random
+import tensorstore as ts
 import torch
 
 from aind_exaspim_image_compression.utils import img_util, util
@@ -37,10 +38,12 @@ def __init__(
         self,
         patch_shape,
         anisotropy=(0.748, 0.748, 1.0),
-        boundary_buffer=4000,
+        boundary_buffer=5000,
         foreground_sampling_rate=0.2,
-        n_examples_per_epoch=200,
-        sigma_bm4d=50,
+        min_brightness=200,
+        n_examples_per_epoch=300,
+        normalization_percentiles=[0.5, 99.9],
+        sigma_bm4d=30,
     ):
         # Call parent class
         super(TrainDataset, self).__init__()
@@ -49,51 +52,110 @@ def __init__(
         self.anisotropy = anisotropy
         self.boundary_buffer = boundary_buffer
         self.foreground_sampling_rate = foreground_sampling_rate
+        self.min_brightness = min_brightness
         self.n_examples_per_epoch = n_examples_per_epoch
+        self.normalization_percentiles = normalization_percentiles
         self.patch_shape = patch_shape
         self.sigma_bm4d = sigma_bm4d
         self.swc_reader = Reader()
 
         # Data structures
-        self.foreground = dict()
+        self.segmentations = dict()
+        self.skeletons = dict()
         self.imgs = dict()
 
     # --- Ingest data ---
-    def ingest_brain(self, brain_id, img_path, swc_pointer):
-        self.foreground[brain_id] = self.load_swcs(swc_pointer)
+    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 : hashable
+            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.segmentations[brain_id] = self.load_segmentation(segmentation_path)
         self.imgs[brain_id] = img_util.read(img_path)
+        self.skeletons[brain_id] = self.load_swcs(swc_pointer)
+
+    def load_segmentation(self, segmentation_path):
+        """
+        Reads a segmentation mask generated by Google Applied Sciences (GAS).
+
+        Parameters
+        ----------
+        segmentation_path : str
+            Path to segmentation.
+
+        Returns
+        -------
+        ...
+        """
+        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]]
+            return label_mask
+        else:
+            return None
 
     def load_swcs(self, 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 foreground voxels
+            # Extract skeleton voxels
             if n_points > 0:
                 start = 0
-                foreground = np.zeros((n_points, 3), dtype=np.int32)
+                skeletons = np.zeros((n_points, 3), dtype=np.int32)
                 for swc_dict in swc_dicts:
                     end = start + len(swc_dict["xyz"])
-                    foreground[start:end] = self.to_voxels(swc_dict["xyz"])
+                    skeletons[start:end] = self.to_voxels(swc_dict["xyz"])
                     start = end
-                return foreground
-        return set()
+                return skeletons
+        return None
 
     # --- Core Routines ---
     def __getitem__(self, dummy_input):
         # Sample image patch
         brain_id = self.sample_brain()
         voxel = self.sample_voxel(brain_id)
-        noise = self.get_patch(brain_id, voxel)
-        mn, mx = np.percentile(noise, 5), np.percentile(noise, 99.9)
+        noise = self.read_patch(brain_id, voxel)
+        mn, mx = np.percentile(noise, self.normalization_percentiles)
 
         # Denoise image patch
         denoised = bm4d(noise, self.sigma_bm4d)
 
         # Normalize image patches
-        noise = (noise - mn) / max(mx, 1)
-        denoised = (denoised - mn) / max(mx, 1)
+        noise = (noise - mn) / max(mx - mn, 1)
+        denoised = (denoised - mn) / max(mx - mn, 1)
         return noise, denoised, (mn, mx)
 
     def sample_brain(self):
@@ -114,12 +176,47 @@ def sample_voxel(self, brain_id):
             return self.sample_interior_voxel(brain_id)
 
     def sample_foreground_voxel(self, brain_id):
-        if len(self.foreground[brain_id]) > 0:
-            idx = random.randint(0, len(self.foreground[brain_id]) - 1)
-            shift = np.random.randint(0, 16, size=3)
-            return tuple(self.foreground[brain_id][idx] + shift)
+        if self.skeletons[brain_id] is not None and np.random.random() > 0.5:
+            return self.sample_skeleton_voxel(brain_id)
+        #elif self.segmentations[brain_id] is not None:
+        #    return self.sample_segmentation_voxel(brain_id)
         else:
-            return self.sample_interior_voxel(brain_id)
+            return self.sample_bright_voxel(brain_id)
+
+    def sample_skeleton_voxel(self, brain_id):
+        idx = random.randint(0, len(self.foreground[brain_id]) - 1)
+        shift = np.random.randint(0, 16, size=3)
+        return tuple(self.foreground[brain_id][idx] + shift)
+
+    def sample_segmentation_voxel(self, brain_id):
+        cnt = 0
+        while cnt < 32:
+            # Read random image patch
+            voxel = self.sample_interior_voxel(brain_id)
+            labels_patch = self.read_precomputed_patch(brain_id, voxel)
+
+            # Check if labels patch has large enough object
+            # --> call fastremap
+            # --> find largest object
+        return voxel
+
+    def sample_bright_voxel(self, brain_id):
+        cnt = 0
+        brightest_voxel, max_brightness = None, 0
+        while cnt < 32:
+            # Read random image patch
+            voxel = self.sample_interior_voxel(brain_id)
+            img_patch = self.read_patch(brain_id, voxel)
+
+            # Check if image patch is bright enough
+            brightness = np.max(img_patch)
+            if brightness >= self.min_brightness:
+                return voxel
+            elif brightness > max_brightness:
+                brightest_voxel = voxel
+                max_brightness = brightness
+            cnt += 1
+        return brightest_voxel
 
     def sample_interior_voxel(self, brain_id):
         voxel = list()
@@ -140,36 +237,54 @@ def __len__(self):
         """
         return self.n_examples_per_epoch
 
-    def get_patch(self, brain_id, voxel):
-        s, e = img_util.get_start_end(voxel, self.patch_shape)
-        return self.imgs[brain_id][0, 0, s[0]: e[0], s[1]: e[1], s[2]: e[2]]
+    def read_patch(self, brain_id, center):
+        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
+        ----------
+        ...
+        """
+        s = img_util.get_slices(center, self.patch_shape)
+        return self.segmentations[brain_id][(0, 0, *s)].read().result()
 
     def to_voxels(self, xyz_arr):
         for i in range(3):
             xyz_arr[:, i] = xyz_arr[:, i] / self.anisotropy[i]
         return np.flip(xyz_arr, axis=1).astype(int)
 
-    def update_foreground_sampling_rate(self, foreground_sampling_rate):
-        self.foreground_sampling_rate = foreground_sampling_rate
-
 
 class ValidateDataset(Dataset):
 
-    def __init__(self, patch_shape, sigma_bm4d=50):
+    def __init__(
+        self,
+        patch_shape,
+        normalization_percentiles=[0.5, 99.9],
+        sigma_bm4d=30,
+    ):
         """
         Instantiates a ValidateDataset object.
 
         Parameters
         ----------
         patch_shape : Tuple[int]
             Shape of image patches to be extracted.
-        sigma_bm4d : float
-            Smoothing parameter used in the BM4D denoising algorithm.
+        normalization_percentiles : List[float], optional
+            Upper and lower percentiles used to normalize the input image.
+            Default is [0.5, 99.5].
+        sigma_bm4d : float, optional
+            Smoothing parameter used in the BM4D denoising algorithm. Default
+            is 30.
         """
         # Call parent class
         super(ValidateDataset, self).__init__()
 
         # Instance attributes
+        self.normalization_percentiles = normalization_percentiles
         self.patch_shape = patch_shape
         self.sigma_bm4d = sigma_bm4d
 
@@ -217,13 +332,13 @@ def ingest_example(self, brain_id, voxel):
             Voxel coordinates of the patch center in the brain volume.
         """
         # Get image patches
-        noise = self.get_patch(brain_id, voxel)
-        mn, mx = np.percentile(noise, 5), np.percentile(noise, 99.9)
+        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 = (noise - mn) / max(mx, 1)
-        denoised = (denoised - mn) / max(mx, 1)
+        noise = (noise - mn) / max(mx - mn, 1)
+        denoised = (denoised - mn) / max(mx - mn, 1)
 
         # Store results
         self.example_ids.append((brain_id, voxel))
@@ -251,9 +366,9 @@ def __getitem__(self, idx):
         """
         return self.noise[idx], self.denoised[idx], self.mn_mxs[idx]
 
-    def get_patch(self, brain_id, voxel):
-        s, e = img_util.get_start_end(voxel, self.patch_shape)
-        return self.imgs[brain_id][0, 0, s[0]: e[0], s[1]: e[1], s[2]: e[2]]
+    def read_patch(self, brain_id, center):
+        slices = img_util.get_slices(center, self.patch_shape)
+        return self.imgs[brain_id][(0, 0, *slices)]
 
 
 # --- Custom Dataloader ---
@@ -326,8 +441,9 @@ def init_datasets(
     foreground_sampling_rate=0.5,
     n_train_examples_per_epoch=100,
     n_validate_examples=0,
-    sigma_bm4d=50,
-    swc_dict=None
+    segmentation_prefixes_path=None,
+    sigma_bm4d=30,
+    swc_pointers=None
 ):
     # Initializations
     train_dataset = TrainDataset(
@@ -338,19 +454,35 @@ def init_datasets(
     )
     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 paths
+        # Set image path
         img_path = get_img_prefix(brain_id, img_paths_json) + str(0)
-        if swc_dict:
-            swc_pointer = deepcopy(swc_dict)
+
+        # 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
-        train_dataset.ingest_brain(brain_id, img_path, swc_pointer)
         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):