Merge pull request #3 from AllenNeuralDynamics/refactor-update-train

refactor: updated doc

Author: anna-grim

src/aind_exaspim_image_compression/machine_learning/data_handling.py

@@ -26,21 +26,29 @@
 
 # --- 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 by using the voxel
+    coordinates from SWC files that represent neuron tracings.
+    """
 
     def __init__(
         self,
         patch_shape,
         anisotropy=(0.748, 0.748, 1.0),
         boundary_buffer=4000,
-        foreground_sampling_rate=0.25,
+        foreground_sampling_rate=0.5,
+        sigma=50,
     ):
         # Call parent class
         super(TrainDataset, self).__init__()
 
         # Class attributes
         self.anisotropy = anisotropy
         self.boundary_buffer = boundary_buffer
-        self.denoise_bm4d = BM4D()
+        self.denoise_bm4d = BM4D(sigma=sigma)
         self.foreground_sampling_rate = foreground_sampling_rate
         self.patch_shape = patch_shape
         self.swc_reader = Reader()
@@ -50,11 +58,11 @@ def __init__(
         self.imgs = dict()
 
     # --- Ingest data ---
-    def ingest_img(self, brain_id, img_path, swc_pointer):
-        self.foreground[brain_id] = self.ingest_swcs(swc_pointer)
+    def ingest_brain(self, brain_id, img_path, swc_pointer):
+        self.foreground[brain_id] = self.load_swcs(swc_pointer)
         self.imgs[brain_id] = img_util.read(img_path)
 
-    def ingest_swcs(self, swc_pointer):
+    def load_swcs(self, swc_pointer):
         if swc_pointer:
             # Initializations
             swc_dicts = self.swc_reader.read(swc_pointer)
@@ -146,7 +154,7 @@ def __len__(self):
         """
         return len(self.ids)
 
-    def ingest_img(self, brain_id, img_path):
+    def ingest_brain(self, brain_id, img_path):
         self.imgs[brain_id] = img_util.read(img_path)
 
     def ingest_example(self, brain_id, voxel):
@@ -190,10 +198,6 @@ def __init__(self, dataset, batch_size=16):
             Dataset to iterated over.
         batch_size : int
             Number of examples in each batch.
-
-        Returns
-        -------
-        None
         """
         # Instance attributes
         self.dataset = dataset
@@ -337,8 +341,8 @@ def init_datasets(
             swc_pointer = None
 
         # Ingest data
-        train_dataset.ingest_img(brain_id, img_path, swc_pointer)
-        val_dataset.ingest_img(brain_id, img_path)
+        train_dataset.ingest_brain(brain_id, img_path, swc_pointer)
+        val_dataset.ingest_brain(brain_id, img_path)
 
     # Generate validation examples
     for _ in range(n_validate_examples):

src/aind_exaspim_image_compression/utils/img_util.py

@@ -191,12 +191,12 @@ def get_patch(img, voxel, shape, is_center=True):
     numpy.ndarray
         Patch extracted from the given image.
     """
-    # Get image patch coordiantes
+    # Get patch coordinates
     start, end = get_start_end(voxel, shape, is_center=is_center)
     valid_start = any([s >= 0 for s in start])
     valid_end = any([e < img.shape[i + 2] for i, e in enumerate(end)])
 
-    # Get image patch
+    # Read patch
     if valid_start and valid_end:
         return img[
             0, 0, start[0]: end[0], start[1]: end[1], start[2]: end[2]
@@ -383,7 +383,7 @@ def compute_cratio(img, codec, patch_shape=(64, 64, 64)):
     Parameters
     ----------
     img : np.ndarray
-        Image to compute compression ratio of.
+        Image to compute the compression ratio of.
     codec : blosc.Blosc
         Blosc codec used to compress each chunk.
     patch_shape : Tuple[int]