Top-down Centered-instance Pipeline (#16)

* added make_centered_bboxes & normalize_bboxes

* added make_centered_bboxes & normalize_bboxes

* created test_instance_cropping.py

* added test normalize bboxes; added find_global_peaks_rough

* black formatted

* black formatted peak_finding

* added make_grid_vectors, normalize_bboxes, integral_regression, added docstring to make_centered_bboxes, fixed find_global_peaks_rough; added crop_bboxes

* finished find_global_peaks with integral regression over centroid crops!

* reformatted with pydocstyle & black

* moved make_grid_vectors to data/utils

* removed normalize_bboxes

* added tests docstrings

* sorted imports with isort

* remove unused imports

* updated test cases for instance cropping

* added minimal_cms.pt fixture + unit tests

* added minimal_bboxes fixture; added unit tests for crop_bboxes & integral_regression

* added find_global_peaks unit tests

* finished find_local_peaks_rough!

* finished find_local_peaks!

* added unit tests for find_local_peaks and find_local_peaks_rough

* updated test cases

* added more test cases for find_local_peaks

* updated test cases

* added architectures folder

* added maxpool2d same padding, get_act_fn; added simpleconvblock, simpleupsamplingblock, encoder, decoder; added unet

* added test_unet_reference

* black formatted common.py & test_unet.py

* deleted tmp nb

* _calc_same_pad returns int

* fixed test case

* added simpleconvblock tests

* added tests

* added tests for simple upsampling block

* updated test_unet

* removed unnecessary variables

* updated augmentation random erase default values

* created data/pipelines.py

* added base config in config/data; temporary till config system settled

* updated variable defaults to 0 and edited variable names in augmentation

* updated parameter names in data/instance_cropping

* added data/pipelines topdown pipeline make_base_pipeline

* added test_pipelines

* removed configs

* updated augmentation class

* modified test

* updated pipelines docstring

* removed make_base_pipeline and updated tests

* removed empty_cache in SleapDataset

* updated test_pipelines

* updated sleapdataset to return a dict

* added key filter transformer block, removed sleap dataset, added type hinting

* updated type hints

* added coderabbit suggestions

* fixed small squeeze issue

sleap_nn/architectures/common.py

@@ -1,4 +1,6 @@
 """Common utilities for architecture and model building."""
+from typing import List
+
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -134,7 +136,7 @@ def get_act_fn(activation: str) -> nn.Module:
  return activations[activation]
 
 
-def get_children_layers(model: torch.nn.Module):
+def get_children_layers(model: torch.nn.Module) -> List[nn.Module]:
  """Recursively retrieves a flattened list of all children modules and submodules within the given model.
 
  Args:

sleap_nn/data/augmentation.py

@@ -1,5 +1,5 @@
 """This module implements data pipeline blocks for augmentation operations."""
-from typing import Any, Dict, Optional, Text, Tuple, Union
+from typing import Any, Dict, Iterator, Optional, Text, Tuple, Union
 
 import kornia as K
 import torch
@@ -100,9 +100,6 @@ class KorniaAugmenter(IterDataPipe):
  Attributes:
  source_dp: The input `IterDataPipe` with examples that contain `"instances"` and
  `"image"` keys.
- crop_hw: Desired output size (out_h, out_w) of the crop. Must be Tuple[int, int],
- then out_h = size[0], out_w = size[1].
- crop_p: Probability of applying random crop.
  rotation: Angles in degrees as a scalar float of the amount of rotation. A
  random angle in `(-rotation, rotation)` will be sampled and applied to both
  images and keypoints. Set to 0 to disable rotation augmentation.
@@ -125,11 +122,14 @@ class KorniaAugmenter(IterDataPipe):
  contrast_p: Probability of applying random contrast.
  brightness: The brightness factor to apply Default: `(1.0, 1.0)`.
  brightness_p: Probability of applying random brightness.
- erase_scale: Range of proportion of erased area against input image. Default: `(0.02, 0.33)`.
- erase_ratio: Range of aspect ratio of erased area. Default: `(0.3, 3.3)`.
+ erase_scale: Range of proportion of erased area against input image. Default: `(0.0001, 0.01)`.
+ erase_ratio: Range of aspect ratio of erased area. Default: `(1, 1)`.
  erase_p: Probability of applying random erase.
  mixup_lambda: min-max value of mixup strength. Default is 0-1. Default: `None`.
  mixup_p: Probability of applying random mixup v2.
+ random_crop_hw: Desired output size (out_h, out_w) of the crop. Must be Tuple[int, int],
+ then out_h = size[0], out_w = size[1].
+ random_crop_p: Probability of applying random crop.
 
  Notes:
  This block expects the "image" and "instances" keys to be present in the input
@@ -150,24 +150,24 @@ def __init__(
  rotation: Optional[float] = 15.0,
  scale: Optional[float] = 0.05,
  translate: Optional[Tuple[float, float]] = (0.02, 0.02),
- affine_p: float = 0.5,
+ affine_p: float = 0.0,
  uniform_noise: Optional[Tuple[float, float]] = (0.0, 0.04),
- uniform_noise_p: float = 0.5,
+ uniform_noise_p: float = 0.0,
  gaussian_noise_mean: Optional[float] = 0.02,
  gaussian_noise_std: Optional[float] = 0.004,
- gaussian_noise_p: float = 0.5,
+ gaussian_noise_p: float = 0.0,
  contrast: Optional[Tuple[float, float]] = (0.5, 2.0),
- contrast_p: float = 0.5,
+ contrast_p: float = 0.0,
  brightness: Optional[float] = 0.0,
- brightness_p: float = 0.5,
+ brightness_p: float = 0.0,
  erase_scale: Optional[Tuple[float, float]] = (0.0001, 0.01),
  erase_ratio: Optional[Tuple[float, float]] = (1, 1),
- erase_p: float = 0.5,
+ erase_p: float = 0.0,
  mixup_lambda: Union[Optional[float], Tuple[float, float], None] = None,
- mixup_p: float = 0.5,
- crop_hw: Tuple[int, int] = (0, 0),
- crop_p: float = 0.0,
- ):
+ mixup_p: float = 0.0,
+ random_crop_hw: Tuple[int, int] = (0, 0),
+ random_crop_p: float = 0.0,
+ ) -> None:
  """Initialize the block and the augmentation pipeline."""
  self.source_dp = source_dp
  self.rotation = rotation
@@ -188,8 +188,8 @@ def __init__(
  self.erase_p = erase_p
  self.mixup_lambda = mixup_lambda
  self.mixup_p = mixup_p
- self.crop_hw = crop_hw
- self.crop_p = crop_p
+ self.random_crop_hw = random_crop_hw
+ self.random_crop_p = random_crop_p
 
  aug_stack = []
  if self.affine_p > 0:
@@ -259,19 +259,21 @@ def __init__(
  same_on_batch=True,
  )
  )
- if self.crop_p > 0:
- if self.crop_hw[0] > 0 and self.crop_hw[1] > 0:
+ if self.random_crop_p > 0:
+ if self.random_crop_hw[0] > 0 and self.random_crop_hw[1] > 0:
  aug_stack.append(
  K.augmentation.RandomCrop(
- size=self.crop_hw,
+ size=self.random_crop_hw,
  pad_if_needed=True,
- p=self.crop_p,
+ p=self.random_crop_p,
  keepdim=True,
  same_on_batch=True,
  )
  )
  else:
- raise ValueError(f"crop_hw height and width must be greater than 0.")
+ raise ValueError(
+ f"random_crop_hw height and width must be greater than 0."
+ )
 
  self.augmenter = AugmentationSequential(
  *aug_stack,
@@ -280,12 +282,31 @@ def __init__(
  same_on_batch=True,
  )
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
  """Return an example dictionary with the augmented image and instances."""
  for ex in self.source_dp:
- inst_shape = ex["instances"].shape # (B, num_instances, num_nodes, 2)
- image, instances = ex["image"], ex["instances"].reshape(
- inst_shape[0], -1, 2
- )
- aug_image, aug_instances = self.augmenter(image, instances)
- yield {"image": aug_image, "instances": aug_instances.reshape(*inst_shape)}
+ if "instance_image" in ex and "instance" in ex:
+ inst_shape = ex["instance"].shape
+ # (B, channels, height, width), (1, num_nodes, 2)
+ image, instances = ex["instance_image"], ex["instance"].unsqueeze(0)
+ aug_image, aug_instances = self.augmenter(image, instances)
+ ex.update(
+ {
+ "instance_image": aug_image,
+ "instance": aug_instances.reshape(*inst_shape),
+ }
+ )
+ elif "image" in ex and "instances" in ex:
+ inst_shape = ex["instances"].shape # (B, num_instances, num_nodes, 2)
+ image, instances = ex["image"], ex["instances"].reshape(
+ inst_shape[0], -1, 2
+ ) # (B, channels, height, width), (B, num_instances x num_nodes, 2)
+
+ aug_image, aug_instances = self.augmenter(image, instances)
+ ex.update(
+ {
+ "image": aug_image,
+ "instances": aug_instances.reshape(*inst_shape),
+ }
+ )
+ yield ex

sleap_nn/data/confidence_maps.py

@@ -1,5 +1,5 @@
 """Generate confidence maps."""
-from typing import Optional
+from typing import Dict, Iterator, Optional
 
 import sleap_io as sio
 import torch
@@ -10,7 +10,7 @@
 
 def make_confmaps(
  points: torch.Tensor, xv: torch.Tensor, yv: torch.Tensor, sigma: float
-):
+) -> torch.Tensor:
  """Make confidence maps from a set of points from a single instance.
 
  Args:
@@ -70,15 +70,15 @@ def __init__(
  output_stride: int = 1,
  instance_key: str = "instance",
  image_key: str = "instance_image",
- ):
+ ) -> None:
  """Initialize ConfidenceMapGenerator with input `DataPipe`, sigma, and output stride."""
  self.source_dp = source_dp
  self.sigma = sigma
  self.output_stride = output_stride
  self.instance_key = instance_key
  self.image_key = image_key
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
  """Generate confidence maps for each example."""
  for example in self.source_dp:
  instance = example[self.instance_key]

sleap_nn/data/general.py

@@ -0,0 +1,39 @@
+"""General purpose transformers for common pipeline processing tasks."""
+from typing import Callable, Dict, Iterator, List, Text
+
+import torch
+from torch.utils.data.datapipes.datapipe import IterDataPipe
+
+
+class KeyFilter(IterDataPipe):
+ """Transformer for filtering example keys."""
+
+ def __init__(self, source_dp: IterDataPipe, keep_keys: List[Text] = None) -> None:
+ """Initialize KeyFilter with the source `DataPipe."""
+ self.dp = source_dp
+ self.keep_keys = set(keep_keys) if keep_keys else None
+
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
+ """Return a dictionary filtered for the relevant outputs.
+
+ The input dictionary includes:
+ - image: the full frame image
+ - instances: all keypoints of all instances in the frame image
+ - centroids: all centroids of all instances in the frame image
+ - instance: the individual instance's keypoints
+ - instance_bbox: the individual instance's bbox
+ - instance_image: the individual instance's cropped image
+ - confidence_maps: the individual instance's heatmap
+ """
+ for example in self.dp:
+ if self.keep_keys is None:
+ # If keep_keys is not provided, yield the entire example.
+ yield example
+ else:
+ # Filter the example dictionary based on keep_keys.
+ filtered_example = {
+ key: value
+ for key, value in example.items()
+ if key in self.keep_keys
+ }
+ yield filtered_example

sleap_nn/data/instance_centroids.py

@@ -1,5 +1,5 @@
 """Handle calculation of instance centroids."""
-from typing import Optional
+from typing import Dict, Iterator, Optional
 
 import torch
 from torch.utils.data.datapipes.datapipe import IterDataPipe
@@ -79,15 +79,15 @@ def __init__(
  self,
  source_dp: IterDataPipe,
  anchor_ind: Optional[int] = None,
- ):
+ ) -> None:
  """Initialize InstanceCentroidFinder with the source `DataPipe."""
  self.source_dp = source_dp
  self.anchor_ind = anchor_ind
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
  """Add `"centroids"` key to example."""
- for example in self.source_dp:
- example["centroids"] = find_centroids(
- example["instances"], anchor_ind=self.anchor_ind
+ for ex in self.source_dp:
+ ex["centroids"] = find_centroids(
+ ex["instances"], anchor_ind=self.anchor_ind
  )
- yield example
+ yield ex

sleap_nn/data/instance_cropping.py

@@ -1,5 +1,5 @@
 """Handle cropping of instances."""
-from typing import Optional
+from typing import Dict, Iterator, Optional, Tuple
 
 import numpy as np
 import sleap_io as sio
@@ -58,50 +58,46 @@ class InstanceCropper(IterDataPipe):
 
  Attributes:
  source_dp: The previous `DataPipe` with samples that contain an `instances` key.
- crop_width: Width of the crop in pixels
- crop_height: Height of the crop in pixels
+ crop_hw: Height and Width of the crop in pixels
  """
 
- def __init__(
- self,
- source_dp: IterDataPipe,
- crop_width: int,
- crop_height: int,
- ):
+ def __init__(self, source_dp: IterDataPipe, crop_hw: Tuple[int, int]) -> None:
  """Initialize InstanceCropper with the source `DataPipe."""
  self.source_dp = source_dp
- self.crop_width = crop_width
- self.crop_height = crop_height
+ self.crop_hw = crop_hw
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
  """Generate instance cropped examples."""
- for example in self.source_dp:
- image = example["image"] # (frames, channels, height, width)
- instances = example["instances"] # (frames, n_instances, n_nodes, 2)
- centroids = example["centroids"] # (frames, n_instances, 2)
+ for ex in self.source_dp:
+ image = ex["image"] # (B, channels, height, width)
+ instances = ex["instances"] # (B, n_instances, num_nodes, 2)
+ centroids = ex["centroids"] # (B, n_instances, 2)
  for instance, centroid in zip(instances[0], centroids[0]):
  # Generate bounding boxes from centroid.
- bbox = torch.unsqueeze(
- make_centered_bboxes(centroid, self.crop_height, self.crop_width), 0
- ) # (frames, 4, 2)
+ instance_bbox = torch.unsqueeze(
+ make_centered_bboxes(centroid, self.crop_hw[0], self.crop_hw[1]), 0
+ ) # (B, 4, 2)
 
- box_size = (self.crop_height, self.crop_width)
+ box_size = (self.crop_hw[0], self.crop_hw[1])
 
- # Generate cropped image of shape (frames, channels, crop_height, crop_width)
+ # Generate cropped image of shape (B, channels, crop_height, crop_width)
  instance_image = crop_and_resize(
  image,
- boxes=bbox,
+ boxes=instance_bbox,
  size=box_size,
  )
 
  # Access top left point (x,y) of bounding box and subtract this offset from
  # position of nodes.
- point = bbox[0][0]
+ point = instance_bbox[0][0]
  center_instance = instance - point
 
  instance_example = {
- "instance_image": instance_image, # (frames, channels, crop_height, crop_width)
- "bbox": bbox, # (frames, 4, 2)
- "instance": center_instance, # (n_instances, 2)
+ "instance_image": instance_image.squeeze(
+ 0
+ ), # (B=1, channels, crop_height, crop_width)
+ "instance_bbox": instance_bbox, # (B, 4, 2)
+ "instance": center_instance, # (num_nodes, 2)
  }
- yield instance_example
+ ex.update(instance_example)
+ yield ex

sleap_nn/data/normalization.py

@@ -1,4 +1,6 @@
 """This module implements data pipeline blocks for normalization operations."""
+from typing import Dict, Iterator
+
 import torch
 from torch.utils.data.datapipes.datapipe import IterDataPipe
 
@@ -16,11 +18,11 @@ class Normalizer(IterDataPipe):
  def __init__(
  self,
  source_dp: IterDataPipe,
- ):
+ ) -> None:
  """Initialize the block."""
  self.source_dp = source_dp
 
- def __iter__(self):
+ def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
  """Return an example dictionary with the augmented image and instance."""
  for ex in self.source_dp:
  if not torch.is_floating_point(ex["image"]):