Fix ID models (#345)

gitttt-1234 · web-flow · commit df7ae20a9495 · 2025-10-06T18:15:22.000-07:00
This PR fixes a few minor issues with TopDown and BottomUp ID models. 
- The ID models dataset classes were re-computing the tracks from the
labels file. However, they should just grab it from the head config
`classes` parameter.
- Fix shape mismatch issue with BottomUp ID models
diff --git a/docs/config.md b/docs/config.md
@@ -689,9 +689,9 @@ The trainer configuration section controls the training process, including data
 ### Data Loader Settings
 - `train_data_loader`:
     - `batch_size`: (int) Number of samples per batch or batch size for training data. **Default**: `1`
-    - `shuffle`: (bool) True to have the data reshuffled at every epoch. **Default**: `False`
+    - `shuffle`: (bool) True to have the data reshuffled at every epoch. **Default**: `True`
     - `num_workers`: (int) Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. **Default**: `0`
-- `val_data_loader`: (Similar to `train_data_loader`)
+- `val_data_loader`: (Similar to `train_data_loader`, but `shuffle` is set to `False` by default)
 
 **Example Data Loader configurations:**
 
diff --git a/docs/sample_configs/config_bottomup_unet_large_rf.yaml b/docs/sample_configs/config_bottomup_unet_large_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_bottomup_unet_medium_rf.yaml b/docs/sample_configs/config_bottomup_unet_medium_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_centroid_unet.yaml b/docs/sample_configs/config_centroid_unet.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_single_instance_unet_large_rf.yaml b/docs/sample_configs/config_single_instance_unet_large_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_single_instance_unet_medium_rf.yaml b/docs/sample_configs/config_single_instance_unet_medium_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_topdown_centered_instance_unet_large_rf.yaml b/docs/sample_configs/config_topdown_centered_instance_unet_large_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_topdown_centered_instance_unet_medium_rf.yaml b/docs/sample_configs/config_topdown_centered_instance_unet_medium_rf.yaml
@@ -48,7 +48,7 @@ data_config:
       mixup_lambda_min: 0.01
       mixup_lambda_max: 0.05
       mixup_p: 0.0
-  skeletons: []
+  skeletons:
 model_config:
   init_weights: default
   pretrained_backbone_weights: null
diff --git a/docs/sample_configs/config_topdown_multi_class_centered_instance_unet.yaml b/docs/sample_configs/config_topdown_multi_class_centered_instance_unet.yaml
@@ -59,7 +59,7 @@ model_config:
       kernel_size: 3
       filters: 16
       filters_rate: 1.5
-      max_stride: 8
+      max_stride: 16
       stem_stride: null
       middle_block: true
       up_interpolate: true
diff --git a/sleap_nn/config/data_config.py b/sleap_nn/config/data_config.py
@@ -466,6 +466,10 @@ def data_mapper(legacy_config: dict) -> DataConfig:
     data_cfg_args["use_augmentations_train"] = (
         True if any(intensity_args.values()) or any(geometric_args.values()) else False
     )
-    data_cfg_args["skeletons"] = skeletons_list
+    data_cfg_args["skeletons"] = (
+        skeletons_list
+        if skeletons_list is not None and len(skeletons_list) > 0
+        else None
+    )
 
     return DataConfig(**data_cfg_args)
diff --git a/sleap_nn/config/get_config.py b/sleap_nn/config/get_config.py
@@ -37,7 +37,8 @@
 from sleap_nn.config.data_config import DataConfig, PreprocessingConfig
 from sleap_nn.config.model_config import ModelConfig
 from sleap_nn.config.trainer_config import (
-    DataLoaderConfig,
+    TrainDataLoaderConfig,
+    ValDataLoaderConfig,
     LRSchedulerConfig,
     StepLRConfig,
     EarlyStoppingConfig,
@@ -647,7 +648,7 @@ def get_model_config(
 
 def get_trainer_config(
     batch_size: int = 1,
-    shuffle_train: bool = False,
+    shuffle_train: bool = True,
     num_workers: int = 0,
     ckpt_save_top_k: int = 1,
     ckpt_save_last: Optional[bool] = None,
@@ -697,7 +698,7 @@ def get_trainer_config(
 
     Args:
         batch_size: Number of samples per batch or batch size for training data. Default: 4.
-        shuffle_train: True to have the train data reshuffled at every epoch. Default: False.
+        shuffle_train: True to have the train data reshuffled at every epoch. Default: True.
         num_workers: Number of subprocesses to use for data loading. 0 means that the data
             will be loaded in the main process. Default: 0.
         ckpt_save_top_k: If save_top_k == k, the best k models according to the quantity
@@ -785,10 +786,10 @@ def get_trainer_config(
             should wait to receive a response and should be set to a small value to minimize the impact on training speed.
     """
     # constrict trainer config
-    train_dataloader_cfg = DataLoaderConfig(
+    train_dataloader_cfg = TrainDataLoaderConfig(
         batch_size=batch_size, shuffle=shuffle_train, num_workers=num_workers
     )
-    val_dataloader_cfg = DataLoaderConfig(
+    val_dataloader_cfg = ValDataLoaderConfig(
         batch_size=batch_size, shuffle=False, num_workers=num_workers
     )
 
diff --git a/sleap_nn/config/trainer_config.py b/sleap_nn/config/trainer_config.py
@@ -14,9 +14,7 @@
 
 @define
 class DataLoaderConfig:
-    """Train and val DataLoaderConfig.
-
-    Any parameters from Torch's DataLoader could be used.
+    """Train DataLoaderConfig.
 
     Attributes:
         batch_size: (int) Number of samples per batch or batch size for training/validation data. *Default*: `1`.
@@ -29,6 +27,32 @@ class DataLoaderConfig:
     num_workers: int = 0
 
 
+@define
+class TrainDataLoaderConfig(DataLoaderConfig):
+    """Train DataLoaderConfig.
+
+    Attributes:
+        batch_size: (int) Number of samples per batch or batch size for training/validation data. *Default*: `1`.
+        shuffle: (bool) True to have the data reshuffled at every epoch. *Default*: `True`.
+        num_workers: (int) Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. *Default*: `0`.
+    """
+
+    shuffle: bool = True
+
+
+@define
+class ValDataLoaderConfig(DataLoaderConfig):
+    """Validation DataLoaderConfig.
+
+    Attributes:
+        batch_size: (int) Number of samples per batch or batch size for training/validation data. *Default*: `1`.
+        shuffle: (bool) True to have the data reshuffled at every epoch. *Default*: `False`.
+        num_workers: (int) Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. *Default*: `0`.
+    """
+
+    shuffle: bool = False
+
+
 @define
 class ModelCkptConfig:
     """Configuration for model checkpoint.
@@ -232,8 +256,8 @@ class TrainerConfig:
         zmq: Zmq config with publish and controller port addresses.
     """
 
-    train_data_loader: DataLoaderConfig = field(factory=DataLoaderConfig)
-    val_data_loader: DataLoaderConfig = field(factory=DataLoaderConfig)
+    train_data_loader: TrainDataLoaderConfig = field(factory=TrainDataLoaderConfig)
+    val_data_loader: ValDataLoaderConfig = field(factory=ValDataLoaderConfig)
     model_ckpt: ModelCkptConfig = field(factory=ModelCkptConfig)
     trainer_devices: Optional[Any] = field(
         default=None,
@@ -346,7 +370,7 @@ def trainer_mapper(legacy_config: dict) -> TrainerConfig:
             "num_workers"
         ]
 
-    trainer_cfg_args["train_data_loader"] = DataLoaderConfig(
+    trainer_cfg_args["train_data_loader"] = TrainDataLoaderConfig(
         **train_dataloader_cfg_args
     )
 
@@ -359,7 +383,7 @@ def trainer_mapper(legacy_config: dict) -> TrainerConfig:
             "num_workers"
         ]
 
-    trainer_cfg_args["val_data_loader"] = DataLoaderConfig(**val_dataloader_cfg_args)
+    trainer_cfg_args["val_data_loader"] = ValDataLoaderConfig(**val_dataloader_cfg_args)
 
     # model ckpt
     if (
diff --git a/sleap_nn/data/custom_datasets.py b/sleap_nn/data/custom_datasets.py
@@ -499,10 +499,8 @@ def __init__(
         )
         self.confmap_head_config = confmap_head_config
         self.class_maps_head_config = class_maps_head_config
-        self.tracks = []
-        for train_label in self.labels:
-            self.tracks.extend([x.name for x in train_label.tracks if x is not None])
-        self.tracks = list(set(self.tracks))
+
+        self.class_names = self.class_maps_head_config.classes
         self.class_map_threshold = class_map_threshold
 
     def __getitem__(self, index) -> Dict:
@@ -537,19 +535,18 @@ def __getitem__(self, index) -> Dict:
             max_instances=self.max_instances,
             user_instances_only=self.user_instances_only,
         )
-
-        sample["track_ids"] = torch.Tensor(
+        track_ids = torch.Tensor(
             [
                 (
-                    self.tracks.index(lf.instances[idx].track.name)
+                    self.class_names.index(lf.instances[idx].track.name)
                     if lf.instances[idx].track is not None
                     else -1
                 )
                 for idx in range(sample["num_instances"])
             ]
         ).to(torch.int32)
 
-        sample["num_tracks"] = torch.tensor(len(self.tracks), dtype=torch.int32)
+        sample["num_tracks"] = torch.tensor(len(self.class_names), dtype=torch.int32)
 
         # apply normalization
         sample["image"] = apply_normalization(sample["image"])
@@ -613,7 +610,7 @@ def __getitem__(self, index) -> Dict:
             instances=sample["instances"],
             img_hw=img_hw,
             num_instances=sample["num_instances"],
-            class_inds=sample["track_ids"],
+            class_inds=track_ids,
             num_tracks=sample["num_tracks"],
             class_map_threshold=self.class_map_threshold,
             sigma=self.class_maps_head_config.sigma,
@@ -669,6 +666,8 @@ class CenteredInstanceDataset(BaseDataset):
         crop_size: Crop size of each instance for centered-instance model.
         rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
             disk occurs only once across all workers.
+        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
+            (required keys: `sigma`, `output_stride`, `part_names` and `anchor_part` depending on the model type ).
 
     Note: If scale is provided for centered-instance model, the images are cropped out
     from the scaled image with the given crop size.
@@ -823,7 +822,9 @@ def __getitem__(self, index) -> Dict:
         sample["frame_idx"] = torch.tensor(lf.frame_idx, dtype=torch.int32)
         sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
         sample["num_instances"] = num_instances
-        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width])
+        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
+            0
+        )
         sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)
 
         # apply augmentation
@@ -928,6 +929,10 @@ class TopDownCenteredInstanceMultiClassDataset(CenteredInstanceDataset):
         crop_size: Crop size of each instance for centered-instance model.
         rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
             disk occurs only once across all workers.
+        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
+            (required keys: `sigma`, `output_stride`, `part_names` and `anchor_part` depending on the model type ).
+        class_vectors_head_config: DictConfig object with all the keys in the `head_config` section.
+            (required keys: `classes`, `num_fc_layers`, `num_fc_units`, `output_stride`, `loss_weight`).
 
     Note: If scale is provided for centered-instance model, the images are cropped out
     from the scaled image with the given crop size.
@@ -938,6 +943,7 @@ def __init__(
         labels: List[sio.Labels],
         crop_size: int,
         confmap_head_config: DictConfig,
+        class_vectors_head_config: DictConfig,
         max_stride: int,
         anchor_ind: Optional[int] = None,
         user_instances_only: bool = True,
@@ -973,10 +979,8 @@ def __init__(
             use_existing_imgs=use_existing_imgs,
             rank=rank,
         )
-        self.tracks = []
-        for train_label in self.labels:
-            self.tracks.extend([x.name for x in train_label.tracks if x is not None])
-        self.tracks = list(set(self.tracks))
+        self.class_vectors_head_config = class_vectors_head_config
+        self.class_names = self.class_vectors_head_config.classes
 
     def __getitem__(self, index) -> Dict:
         """Return dict with cropped image and confmaps of instance for given index."""
@@ -1055,7 +1059,7 @@ def __getitem__(self, index) -> Dict:
         track_ids = torch.Tensor(
             [
                 (
-                    self.tracks.index(lf.instances[idx].track.name)
+                    self.class_names.index(lf.instances[idx].track.name)
                     if lf.instances[idx].track is not None
                     else -1
                 )
@@ -1064,7 +1068,7 @@ def __getitem__(self, index) -> Dict:
         ).to(torch.int32)
         class_vectors = make_class_vectors(
             class_inds=track_ids,
-            n_classes=torch.tensor(len(self.tracks), dtype=torch.int32),
+            n_classes=torch.tensor(len(self.class_names), dtype=torch.int32),
         )
 
         # get the centroids based on the anchor idx
@@ -1082,7 +1086,9 @@ def __getitem__(self, index) -> Dict:
         sample["frame_idx"] = torch.tensor(lf.frame_idx, dtype=torch.int32)
         sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
         sample["num_instances"] = num_instances
-        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width])
+        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
+            0
+        )
         sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)
 
         # apply augmentation
@@ -1140,7 +1146,6 @@ def __getitem__(self, index) -> Dict:
             output_stride=self.confmap_head_config.output_stride,
         )
 
-        sample["track_id"] = track_ids[inst_idx]
         sample["class_vectors"] = class_vectors[inst_idx].to(torch.float32)
 
         sample["confidence_maps"] = confidence_maps
@@ -1370,7 +1375,7 @@ class SingleInstanceDataset(BaseDataset):
         cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
         use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
         confmap_head_config: DictConfig object with all the keys in the `head_config` section.
-        (required keys: `sigma`, `output_stride` and `anchor_part` depending on the model type ).
+        (required keys: `sigma`, `output_stride` and `part_names` depending on the model type ).
         rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
             disk occurs only once across all workers.
     """
@@ -1834,6 +1839,7 @@ def get_train_val_datasets(
         train_dataset = TopDownCenteredInstanceMultiClassDataset(
             labels=train_labels,
             confmap_head_config=config.model_config.head_configs.multi_class_topdown.confmaps,
+            class_vectors_head_config=config.model_config.head_configs.multi_class_topdown.class_vectors,
             max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                 "max_stride"
             ],
@@ -1866,6 +1872,7 @@ def get_train_val_datasets(
         val_dataset = TopDownCenteredInstanceMultiClassDataset(
             labels=val_labels,
             confmap_head_config=config.model_config.head_configs.multi_class_topdown.confmaps,
+            class_vectors_head_config=config.model_config.head_configs.multi_class_topdown.class_vectors,
             max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                 "max_stride"
             ],
diff --git a/sleap_nn/data/providers.py b/sleap_nn/data/providers.py
@@ -93,7 +93,7 @@ def process_lf(
         "instances": instances,
         "video_idx": torch.tensor(video_idx, dtype=torch.int32),
         "frame_idx": torch.tensor(lf.frame_idx, dtype=torch.int32),
-        "orig_size": torch.Tensor([img_height, img_width]),
+        "orig_size": torch.Tensor([img_height, img_width]).unsqueeze(0),
         "num_instances": num_instances,
     }
 
@@ -187,7 +187,7 @@ def run(self):
                         "image": torch.from_numpy(img.copy()),
                         "frame_idx": torch.tensor(idx, dtype=torch.int32),
                         "video_idx": torch.tensor(0, dtype=torch.int32),
-                        "orig_size": torch.Tensor(img.shape[-2:]),
+                        "orig_size": torch.Tensor(img.shape[-2:]).unsqueeze(0),
                     }
                 )
 
@@ -320,7 +320,7 @@ def run(self):
                     "image": torch.from_numpy(img.copy()),
                     "frame_idx": torch.tensor(lf.frame_idx, dtype=torch.int32),
                     "video_idx": torch.tensor(video_idx, dtype=torch.int32),
-                    "orig_size": torch.Tensor(img.shape[-2:]),
+                    "orig_size": torch.Tensor(img.shape[-2:]).unsqueeze(0),
                 }
 
                 if self.instances_key:
diff --git a/sleap_nn/train.py b/sleap_nn/train.py
diff --git a/tests/config/test_trainer_config.py b/tests/config/test_trainer_config.py
diff --git a/tests/data/test_custom_datasets.py b/tests/data/test_custom_datasets.py
diff --git a/tests/training/test_lightning_modules.py b/tests/training/test_lightning_modules.py
