Fix augmentation in TopdownConfmaps pipeline (#78)

* Fix topdown aug

* Implement tracker module (#70)

* Refactor preprocessing config

* Merge train and val data configs

* Remove pipeline name

* Modify backbone_config

* Modify ckpts

* Fix inference tests

* Fix device for inference

* Fix scale in inference

* Fix Predictor

* Modify `bottom_up` to `bottomup`

* Fix bottomup inference

* Fix scale in augmentation

* Add tracker

* Fix tracker queue

* Add local queues

* Modify local queues\

* Add features

* Add optical flow

* Add Optical flow

* Add tracking score

* Refactor candidate update

* Integrate with Predictors

* Fix lint

* Fix tracks

* Resume training and automatically compute crop size for TopDownConfmaps pipeline (#79)

* Add option to automatically compute crop size

* Move find_crop_size to Trainer

* Fix skeleton name

* Add crop size to config

* Add resumable training option

* Add tests fore resuming training

* Fix tests

* Fix test for wandb folder

* LitData Refactor PR1: Get individual functions for data pipelines (#81)

* Add functions for data pipelines

* Add test cases

* Format file

* Add more test cases

* Fix augmentation test

docs/config.md

@@ -26,7 +26,8 @@ The config file has three main sections:
         - `max_width`: (int) Maximum width the image should be padded to. If not provided, the
         original image size will be retained. Default: None.
         - `scale`: (float or List[float]) Factor to resize the image dimensions by, specified as either a float scalar or as a 2-tuple of [scale_x, scale_y]. If a scalar is provided, both dimensions are resized by the same factor. 
-        - `crop_hw`: (List[int]) Crop height and width of each instance (h, w) for centered-instance model.
+        - `crop_hw`: (Tuple[int]) Crop height and width of each instance (h, w) for centered-instance model. If `None`, this would be automatically computed based on the largest instance in the `sio.Labels` file.
+        - `min_crop_size`: (int) Minimum crop size to be used if `crop_hw` is `None`.
     - `use_augmentations_train`: (bool) True if the data augmentation should be applied to the training data, else False. 
     - `augmentation_config`: (only if `use_augmentations` is `True`)
         - `random crop`: (Optional) (Dict[float]) {"random_crop_p": None, "crop_height": None. "crop_width": None}, where *random_crop_p* is the probability of applying random crop and *crop_height* and *crop_width* are the desired output size (out_h, out_w) of the crop. 
@@ -156,12 +157,13 @@ The config file has three main sections:
     - `use_wandb`: (bool) True to enable wandb logging.
     - `save_ckpt`: (bool) True to enable checkpointing. 
     - `save_ckpt_path`: (str) Directory path to save the training config and checkpoint files. *Default*: "./"
+    - `resume_ckpt_path`: (str) Path to `.ckpt` file from which training is resumed. *Default*: `None`.
     - `wandb`: (Only if `use_wandb` is `True`, else skip this)
-        - `entity`: (str) Entity of wandb project.
         - `project`: (str) Project name for the wandb project.
         - `name`: (str) Name of the current run.
         - `api_key`: (str) API key. The API key is masked when saved to config files.
         - `wandb_mode`: (str) "offline" if only local logging is required. Default: "None".
+        - `prv_runid`: (str) Previous run ID if training should be resumed from a previous ckpt. *Default*: `None`.
         - `log_params`: (List[str]) List of config parameters to save it in wandb logs. For example, to save learning rate from trainer config section, use "trainer_config.optimizer.lr" (provide the full path to the specific config parameter).
     - `optimizer_name`: (str) Optimizer to be used. One of ["Adam", "AdamW"].
     - `optimizer`

docs/config_bottomup.yaml

@@ -96,6 +96,7 @@ trainer_config:
   use_wandb: false
   save_ckpt: true
   save_ckpt_path: min_inst_bottomup1
+  resume_ckpt_path:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

docs/config_centroid.yaml

@@ -113,12 +113,13 @@ trainer_config:
   use_wandb: false
   save_ckpt: true
   save_ckpt_path: 'min_inst_centroid'
+  resume_ckpt_path:
   wandb: # sample wandb config
-    entity: 
     project: 'test_centroid_centered'
     name: 'fly_unet_centered'
     wandb_mode: ''
     api_key: ''
+    prv_runid:
     log_params:
     - trainer_config.optimizer_name
     - trainer_config.optimizer.amsgrad

docs/config_topdown_centered_instance.yaml

@@ -94,6 +94,7 @@ trainer_config:
   use_wandb: false
   save_ckpt: true
   save_ckpt_path: 'min_inst_centered'
+  resume_ckpt_path:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

sleap_nn/data/augmentation.py

@@ -1,7 +1,6 @@
 """This module implements data pipeline blocks for augmentation operations."""
 
 from typing import Any, Dict, Optional, Tuple, Union, Iterator
-
 import kornia as K
 import torch
 from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D
@@ -11,6 +10,227 @@
 from torch.utils.data.datapipes.datapipe import IterDataPipe
 
 
+def apply_intensity_augmentation(
+    image: torch.Tensor,
+    instances: torch.Tensor,
+    uniform_noise_min: Optional[float] = 0.0,
+    uniform_noise_max: Optional[float] = 0.04,
+    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.0,
+    contrast_min: Optional[float] = 0.5,
+    contrast_max: Optional[float] = 2.0,
+    contrast_p: float = 0.0,
+    brightness: Optional[float] = 0.0,
+    brightness_p: float = 0.0,
+) -> Tuple[torch.Tensor]:
+    """Apply kornia intensity augmentation on image and instances.
+
+    Args:
+        image: Input image. Shape: (n_samples, C, H, W)
+        instances: Input keypoints. (n_samples, n_instances, n_nodes, 2) or (n_samples, n_nodes, 2)
+        uniform_noise_min: Minimum value for uniform noise (uniform_noise_min >=0).
+        uniform_noise_max: Maximum value for uniform noise (uniform_noise_max <=1).
+        uniform_noise_p: Probability of applying random uniform noise.
+        gaussian_noise_mean: The mean of the gaussian distribution.
+        gaussian_noise_std: The standard deviation of the gaussian distribution.
+        gaussian_noise_p: Probability of applying random gaussian noise.
+        contrast_min: Minimum contrast factor to apply. Default: 0.5.
+        contrast_max: Maximum contrast factor to apply. Default: 2.0.
+        contrast_p: Probability of applying random contrast.
+        brightness: The brightness factor to apply Default: 0.0.
+        brightness_p: Probability of applying random brightness.
+
+    Returns:
+        Returns tuple: (image, instances) with augmentation applied.
+    """
+    aug_stack = []
+    if uniform_noise_p > 0:
+        aug_stack.append(
+            RandomUniformNoise(
+                noise=(uniform_noise_min, uniform_noise_max),
+                p=uniform_noise_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if gaussian_noise_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomGaussianNoise(
+                mean=gaussian_noise_mean,
+                std=gaussian_noise_std,
+                p=gaussian_noise_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if contrast_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomContrast(
+                contrast=(contrast_min, contrast_max),
+                p=contrast_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if brightness_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomBrightness(
+                brightness=brightness,
+                p=brightness_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+
+    augmenter = AugmentationSequential(
+        *aug_stack,
+        data_keys=["input", "keypoints"],
+        keepdim=True,
+        same_on_batch=True,
+    )
+
+    inst_shape = instances.shape
+    # Before (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # Before (cropped image): (B=1, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    instances = instances.reshape(inst_shape[0], -1, 2)
+    # (n_samples, C, H, W), (n_samples, n_instances * n_nodes, 2) OR (n_samples, n_nodes, 2)
+
+    aug_image, aug_instances = augmenter(image, instances)
+
+    # After (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # After (cropped image): (n_samples, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    return aug_image, aug_instances.reshape(*inst_shape)
+
+
+def apply_geometric_augmentation(
+    image: torch.Tensor,
+    instances: torch.Tensor,
+    rotation: Optional[float] = 15.0,
+    scale: Union[
+        Optional[float], Tuple[float, float], Tuple[float, float, float, float]
+    ] = None,
+    translate_width: Optional[float] = 0.02,
+    translate_height: Optional[float] = 0.02,
+    affine_p: float = 0.0,
+    erase_scale_min: Optional[float] = 0.0001,
+    erase_scale_max: Optional[float] = 0.01,
+    erase_ratio_min: Optional[float] = 1,
+    erase_ratio_max: Optional[float] = 1,
+    erase_p: float = 0.0,
+    mixup_lambda: Union[Optional[float], Tuple[float, float], None] = None,
+    mixup_p: float = 0.0,
+    random_crop_height: int = 0,
+    random_crop_width: int = 0,
+    random_crop_p: float = 0.0,
+) -> Tuple[torch.Tensor]:
+    """Apply kornia geometric augmentation on image and instances.
+
+    Args:
+        image: Input image. Shape: (n_samples, C, H, W)
+        instances: Input keypoints. (n_samples, n_instances, n_nodes, 2) or (n_samples, n_nodes, 2)
+        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.
+        scale: scaling factor interval. If (a, b) represents isotropic scaling, the scale
+            is randomly sampled from the range a <= scale <= b. If (a, b, c, d), the scale
+            is randomly sampled from the range a <= scale_x <= b, c <= scale_y <= d.
+            Default: None.
+        translate_width: Maximum absolute fraction for horizontal translation. For example,
+            if translate_width=a, then horizontal shift is randomly sampled in the range
+            -img_width * a < dx < img_width * a. Will not translate by default.
+        translate_height: Maximum absolute fraction for vertical translation. For example,
+            if translate_height=a, then vertical shift is randomly sampled in the range
+            -img_height * a < dy < img_height * a. Will not translate by default.
+        affine_p: Probability of applying random affine transformations.
+        erase_scale_min: Minimum value of range of proportion of erased area against input image. Default: 0.0001.
+        erase_scale_max: Maximum value of range of proportion of erased area against input image. Default: 0.01.
+        erase_ratio_min: Minimum value of range of aspect ratio of erased area. Default: 1.
+        erase_ratio_max: Maximum value of range of aspect ratio of erased area. Default: 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_height: Desired output height of the crop. Must be int.
+        random_crop_width: Desired output width of the crop. Must be int.
+        random_crop_p: Probability of applying random crop.
+
+
+    Returns:
+        Returns tuple: (image, instances) with augmentation applied.
+    """
+    if isinstance(scale, float):
+        scale = (scale, scale)
+    aug_stack = []
+    if affine_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomAffine(
+                degrees=rotation,
+                translate=(translate_width, translate_height),
+                scale=scale,
+                p=affine_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+
+    if erase_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomErasing(
+                scale=(erase_scale_min, erase_scale_max),
+                ratio=(erase_ratio_min, erase_ratio_max),
+                p=erase_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if mixup_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomMixUpV2(
+                lambda_val=mixup_lambda,
+                p=mixup_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if random_crop_p > 0:
+        if random_crop_height > 0 and random_crop_width > 0:
+            aug_stack.append(
+                K.augmentation.RandomCrop(
+                    size=(random_crop_height, random_crop_width),
+                    pad_if_needed=True,
+                    p=random_crop_p,
+                    keepdim=True,
+                    same_on_batch=True,
+                )
+            )
+        else:
+            raise ValueError(f"random_crop_hw height and width must be greater than 0.")
+
+    augmenter = AugmentationSequential(
+        *aug_stack,
+        data_keys=["input", "keypoints"],
+        keepdim=True,
+        same_on_batch=True,
+    )
+
+    inst_shape = instances.shape
+    # Before (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # Before (cropped image): (B=1, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    instances = instances.reshape(inst_shape[0], -1, 2)
+    # (n_samples, C, H, W), (n_samples, n_instances * n_nodes, 2) OR (n_samples, n_nodes, 2)
+
+    aug_image, aug_instances = augmenter(image, instances)
+
+    # After (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # After (cropped image): (n_samples, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    return aug_image, aug_instances.reshape(*inst_shape)
+
+
 class RandomUniformNoise(IntensityAugmentationBase2D):
     """Data transformer for applying random uniform noise to input images.
 

sleap_nn/data/confidence_maps.py

@@ -1,13 +1,99 @@
 """Generate confidence maps."""
 
-from typing import Dict, Iterator
+from typing import Dict, Iterator, Tuple
 
 import torch
 from torch.utils.data.datapipes.datapipe import IterDataPipe
 
 from sleap_nn.data.utils import make_grid_vectors
 
 
+def generate_confmaps(
+    instance: torch.Tensor,
+    img_hw: Tuple[int],
+    sigma: float = 1.5,
+    output_stride: int = 2,
+) -> torch.Tensor:
+    """Generate Confidence maps.
+
+    Args:
+        instance: Input keypoints. (n_samples, n_instances, n_nodes, 2) or
+            (n_samples, n_nodes, 2).
+        img_hw: Image size as tuple (height, width).
+        sigma: The standard deviation of the Gaussian distribution that is used to
+            generate confidence maps. Default: 1.5.
+        output_stride: The relative stride to use when generating confidence maps.
+            A larger stride will generate smaller confidence maps. Default: 2.
+
+    Returns:
+        Confidence maps for the input keypoints.
+    """
+    if instance.ndim != 3:
+        instance = instance.view(instance.shape[0], -1, 2)
+        # instances: (n_samples, n_nodes, 2)
+
+    height, width = img_hw
+
+    xv, yv = make_grid_vectors(height, width, output_stride)
+
+    confidence_maps = make_confmaps(
+        instance,
+        xv,
+        yv,
+        sigma * output_stride,
+    )  # (n_samples, n_nodes, height/ output_stride, width/ output_stride)
+
+    return confidence_maps
+
+
+def generate_multiconfmaps(
+    instances: torch.Tensor,
+    img_hw: Tuple[int],
+    num_instances: int,
+    sigma: float = 1.5,
+    output_stride: int = 2,
+    is_centroids: bool = False,
+) -> torch.Tensor:
+    """Generate multi-instance confidence maps.
+
+    Args:
+        instances: Input keypoints. (n_samples, n_instances, n_nodes, 2) or
+            for centroids - (n_samples, n_instances, 2)
+        img_hw: Image size as tuple (height, width).
+        num_instances: Original number of instances in the frame.
+        sigma: The standard deviation of the Gaussian distribution that is used to
+            generate confidence maps. Default: 1.5.
+        output_stride: The relative stride to use when generating confidence maps.
+            A larger stride will generate smaller confidence maps. Default: 2.
+        is_centroids: True if confidence maps should be generates for centroids else False.
+            Default: False.
+
+    Returns:
+        Confidence maps for the input keypoints.
+    """
+    if is_centroids:
+        points = instances[:, :num_instances, :].unsqueeze(dim=-2)
+        # (n_samples, n_instances, 1, 2)
+    else:
+        points = instances[
+            :, :num_instances, :, :
+        ]  # (n_samples, n_instances, n_nodes, 2)
+
+    height, width = img_hw
+
+    xv, yv = make_grid_vectors(height, width, output_stride)
+
+    confidence_maps = make_multi_confmaps(
+        points,
+        xv,
+        yv,
+        sigma * output_stride,
+    )  # (n_samples, n_nodes, height/ output_stride, width/ output_stride).
+    # If `is_centroids`, (n_samples, 1, height/ output_stride, width/ output_stride).
+
+    return confidence_maps
+
+
 def make_confmaps(
     points_batch: torch.Tensor, xv: torch.Tensor, yv: torch.Tensor, sigma: float
 ) -> torch.Tensor: