Added SingleInstanceConfmapsPipeline

[alckasoc]

Summary by CodeRabbit

• New Features
  • Added a search functionality to the app.

[coderabbitai]

CodeRabbit review skipped

By default, CodeRabbit only reviews PRs on the default branch. If you wish to have PRs reviewed on additional branches, you can configure that under the repository settings in the UI.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (45f5b42) 99.78% compared to head (bd09be8) 99.78%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #23   +/-   ##
=======================================
  Coverage   99.78%   99.78%           
=======================================
  Files          19       19           
  Lines         936      949   +13     
=======================================
+ Hits          934      947   +13     
  Misses          2        2           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coderabbitai]

Walkthrough

The changes encompass a comprehensive refactoring and enhancement of the sleap_nn codebase. This involves the addition of new classes and methods, removal of unused code, and modification of existing code to improve functionality and readability. The changes span across various modules such as architectures, data handling, and testing.

Changes

Files	Summary
.gitignore	Addition of .ruff_cache/ to the ignore list and removal of the ruff directory.
sleap_nn/architectures/...	Refactoring and enhancement of architecture code, including addition, removal, and modification of methods and classes.
sleap_nn/data/...	Changes in data handling and augmentation code, including addition of new parameters and modification of function signatures.
sleap_nn/evaluation.py	Minor changes to function documentation.
tests/architectures/...	Addition of new test cases and modification of existing ones to align with architecture code changes.
tests/data/...	Changes in test cases to align with modifications in data handling code, including addition of new test cases and modification of existing ones.

🍂 As November begins, the leaves fall and the code refines, 🍁

🐇 With each hop and skip, the rabbit aligns. 🐾

🎉 Celebrating the changes, big and small, 🎊

🌟 Making the codebase shine, standing tall. 🌠

🎈 So here's to the season of change, crisp and bright, 🍎

🌙 And to the code that's now more light. 🌕

---

Tips

Chat with CodeRabbit Bot (@coderabbitai)

• If you reply to a review comment from CodeRabbit, the bot will automatically respond.
• To engage with CodeRabbit bot directly around the specific lines of code in the PR, mention @coderabbitai in your review comment
• Note: Review comments are made on code diffs or files, not on the PR overview.
• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger a review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai help to get help.

Note: For conversation with the bot, please use the review comments on code diffs or files.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• The JSON schema for the configuration file is available here.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/coderabbit-overrides.v2.json

[coderabbitai]

Review Status

Actionable comments generated: 14

Commits

Files that changed from the base of the PR and between 0dff84f and 3efac3c.

Files selected for processing (19)

• .gitignore (1 hunks)
• sleap_nn/architectures/common.py (1 hunks)
• sleap_nn/architectures/encoder_decoder.py (3 hunks)
• sleap_nn/architectures/heads.py (1 hunks)
• sleap_nn/architectures/model.py (1 hunks)
• sleap_nn/architectures/unet.py (3 hunks)
• sleap_nn/architectures/utils.py (1 hunks)
• sleap_nn/data/augmentation.py (1 hunks)
• sleap_nn/data/confidence_maps.py (1 hunks)
• sleap_nn/data/general.py (1 hunks)
• sleap_nn/data/instance_cropping.py (3 hunks)
• sleap_nn/data/pipelines.py (3 hunks)
• tests/architectures/test_architecture_utils.py (1 hunks)
• tests/architectures/test_common.py (2 hunks)
• tests/architectures/test_encoder_decoder.py (3 hunks)
• tests/architectures/test_heads.py (1 hunks)
• tests/architectures/test_model.py (1 hunks)
• tests/architectures/test_unet.py (3 hunks)
• tests/data/test_pipelines.py (3 hunks)

Files skipped from review due to trivial changes (7)

• .gitignore
• sleap_nn/architectures/common.py
• sleap_nn/data/confidence_maps.py
• sleap_nn/data/general.py
• sleap_nn/data/instance_cropping.py
• tests/architectures/test_common.py
• tests/architectures/test_encoder_decoder.py

Additional comments (Suppressed): 29

tests/architectures/test_architecture_utils.py (2)

• 7-13: The test test_get_act_fn is well written and covers the expected behavior of the get_act_fn function. It tests for an invalid input and checks the type of the returned activation function for valid inputs.

• 16-25: The test test_get_children_layers is also well written. It creates a model with nested layers and checks if the get_children_layers function correctly retrieves the child layers. It also checks the input and output features of the retrieved layers.

tests/data/test_pipelines.py (3)

• 13-16: The function name has been changed from test_sleap_dataset to test_key_filter. Ensure that this change is reflected in all calls to this function throughout the codebase.

• 49-76: The structure of the augmentations dictionary has been changed. It now has two nested dictionaries: intensity and geometric. This change improves the organization of the code by grouping related augmentation parameters together. However, ensure that all parts of the code that use this dictionary have been updated to match the new structure.

• 102-129: Similar to the previous comment, the structure of the augmentations dictionary has been changed in this part of the code as well. Ensure that all parts of the code that use this dictionary have been updated to match the new structure.

sleap_nn/data/augmentation.py (1)

• 1-12: The import statements have been cleaned up, removing unused imports such as Text, GeometricAugmentationBase2D, Resample, SamplePadding, and warp_affine. This is a good practice as it helps to keep the codebase clean and maintainable. However, please verify that these removed imports are not used elsewhere in the code.

tests/architectures/test_unet.py (3)

• 1-6: The import statement for get_children_layers has been updated to import from sleap_nn.architectures.utils instead of sleap_nn.architectures.common. Ensure that the function is correctly implemented in the new module and that all other references to this function in the codebase have been updated accordingly.

• 56-80: The forward method of the UNet model now returns a dictionary with keys "outputs" and "strides". Ensure that all calls to this method throughout the codebase have been updated to handle the new return type. Also, the output tensor is now accessed through y["outputs"][-1] instead of y. Make sure that this change is reflected in all parts of the code that use the output tensor. The convolution operation on line 78 now uses y["outputs"][-1] as input instead of y. Verify that this change does not affect the expected results of the operation.

• 90-96: The eval method of the Encoder model is now called without assignment to _. This is a minor change and does not affect the functionality of the code.

tests/architectures/test_model.py (3)

• 1-32: The test_get_backbone() function is well written and covers the necessary test cases. It tests the get_backbone() function with a valid configuration and asserts that the returned object is an instance of torch.nn.Module. It also tests the function with an invalid input and asserts that a KeyError is raised.

• 34-51: The test_get_head() function is also well written and covers the necessary test cases. It tests the get_head() function with a valid configuration and asserts that the returned object is an instance of Head. It also tests the function with an invalid input and asserts that a KeyError is raised.

• 54-115: The test_unet_model() function tests the Model class and its methods. It creates a Model object with a backbone configuration and a list of head configurations, and asserts that the backbone configuration and head configurations are correctly set. It then performs inference on a random input tensor and asserts the properties of the output. The function also tests the Model.from_config() method and asserts the properties of the output. The function is well written and covers the necessary test cases.

Overall, the new hunks for tests/architectures/test_model.py are well written and cover the necessary test cases. The code is clear, concise, and easy to understand. The use of the pytest testing framework and the OmegaConf library for configuration management is appropriate and effective. The tests ensure the functionality and correctness of the get_backbone(), get_head(), and Model classes.

sleap_nn/data/pipelines.py (4)

• 3-16: The import statement for DictConfig has been updated to reflect the correct import path. The rest of the import statements remain the same. No issues found.

• 40-51: The pipeline now includes a conditional check for data augmentation. If use_augmentations is set to True in the augmentation_config, the KorniaAugmenter is applied to the data pipeline with intensity augmentations. This is a new feature that was not present in the old hunk.

• 58-72: Similar to the previous hunk, a conditional check for data augmentation is performed. If use_augmentations is set to True, the KorniaAugmenter is applied to the data pipeline with geometric augmentations. This is a new feature that was not present in the old hunk. The ConfidenceMapGenerator and KeyFilter are applied to the data pipeline as before.

• 75-121: A new class SingleInstanceConfmapsPipeline is introduced. This class has a similar structure to the TopdownConfmapsPipeline class. It includes a make_training_pipeline method that creates a training pipeline with input data only. The method includes conditional checks for data augmentation and random cropping. If use_augmentations is set to True, the KorniaAugmenter is applied to the data pipeline with intensity and geometric augmentations. If random_crop_p is set to a non-zero value, the KorniaAugmenter is applied to the data pipeline with random cropping. The ConfidenceMapGenerator is applied to the data pipeline as before.

sleap_nn/architectures/unet.py (3)

• 2-12: The import statement for math and F from torch.nn has been removed. Ensure that these are not used elsewhere in the code or in the functions that are not shown in the hunks. If they are used, they should be imported back.

• 44-53: The constructor now initializes several new instance variables. Ensure that these are used appropriately in the rest of the class and that their values are validated where necessary.

• 94-105: The forward method now returns a tuple instead of a single tensor. This is a breaking change and all calls to this method should be updated to handle the new return type. Also, the docstring should be updated to reflect this change.

-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> Tuple[List[torch.Tensor], List]:

sleap_nn/architectures/utils.py (1)

• 6-42: The get_act_fn function is well implemented. It provides a clear and concise way to get an instance of a PyTorch activation function module based on the provided name. The function also handles the case where an unsupported activation function name is provided by raising a KeyError. This is a good practice as it helps to prevent silent failures and makes debugging easier.

sleap_nn/architectures/encoder_decoder.py (3)

• 28-34: The import statement for get_act_fn has been moved from sleap_nn.architectures.common to sleap_nn.architectures.utils. Ensure that the function is correctly defined in the new module and that all references to it in the codebase have been updated.

• 430-437: The current_strides attribute has been added to the Decoder class. This attribute is used to keep track of the stride values for each block in the decoder stack. This is a good practice as it allows for dynamic adjustment of the stride values based on the specific requirements of each block.

• 458-485: The forward method of the Decoder class has been updated to return a dictionary containing a list of output tensors and the current strides from the decoder blocks. This change provides more information about the operations performed by the decoder, which can be useful for debugging and analysis. However, ensure that all calls to this method throughout the codebase have been updated to handle the new return type.

-    def forward(self, x: torch.Tensor, features: List[torch.Tensor]) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor, features: List[torch.Tensor]
+    ) -> Tuple[List[torch.Tensor], List]:

sleap_nn/architectures/model.py (5)

• 1-6: The introduction of the module docstring provides a clear explanation of the purpose of the module and its functionality. This is a good practice as it helps other developers understand the purpose and usage of the module.

• 27-52: The get_backbone function is well implemented. It checks if the provided backbone name is supported and raises a KeyError with a helpful error message if it is not. This is a good practice as it helps prevent errors due to unsupported backbones and provides clear feedback to the user.

• 55-97: The get_head function is similar to the get_backbone function and is also well implemented. It checks if the provided head name is supported and raises a KeyError with a helpful error message if it is not. This is a good practice as it helps prevent errors due to unsupported heads and provides clear feedback to the user.

• 138-144: The from_config class method is a good practice as it provides a convenient way to create a Model instance from a configuration dictionary.

• 145-154: The forward method performs a forward pass through the model and returns the outputs of the heads. It correctly uses the output stride of each head to select the corresponding output from the backbone. This is a good practice as it ensures that the correct output is used for each head.

tests/architectures/test_heads.py (1)

• 17-27: The test for the base Head class is well written and covers the expected behavior of the class. It checks for the correct initialization of attributes and the raising of NotImplementedError for the channels property, which should be implemented in subclasses.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 3efac3c and abe7f97.

Files selected for processing (7)

• sleap_nn/data/augmentation.py (5 hunks)
• sleap_nn/data/confidence_maps.py (2 hunks)
• sleap_nn/data/instance_centroids.py (1 hunks)
• sleap_nn/data/instance_cropping.py (3 hunks)
• sleap_nn/data/pipelines.py (3 hunks)
• tests/data/test_instance_cropping.py (1 hunks)
• tests/data/test_pipelines.py (4 hunks)

Files skipped from review due to trivial changes (2)

• sleap_nn/data/confidence_maps.py
• sleap_nn/data/instance_centroids.py

Additional comments: 16

tests/data/test_instance_cropping.py (1)

• 46-46: The shape of sample["instance_image"] has changed from (1, 100, 100) to (1, 1, 100, 100). Ensure that this change is intentional and that it doesn't break any functionality that depends on the shape of this array.

tests/data/test_pipelines.py (4)

• 3-13: The import statements have been reordered and sleap_io has been added. The SingleInstanceConfmapsPipeline has been imported from sleap_nn.data.pipelines. Ensure that this new pipeline is implemented correctly and that it is compatible with the rest of the codebase.

• 56-80: The augmentation parameters have been restructured into two categories: intensity and geometric. This change could improve readability and organization of the code. However, ensure that the code that uses these parameters has been updated to reflect this new structure.

• 108-132: The same restructuring of augmentation parameters is applied here as well. Again, verify that the code that uses these parameters has been updated to reflect this new structure.

• 158-218: A new test function test_singleinstanceconfmapspipeline has been added. This function tests the SingleInstanceConfmapsPipeline with a single instance example. The test seems to be well-structured and covers the key aspects of the pipeline. However, ensure that the expected shapes of the image and confidence_maps are correct.

sleap_nn/data/instance_cropping.py (3)

• 2-5: The import statements have been updated to remove unused imports and add necessary ones. The removed imports include numpy, sleap_io, and Optional from typing. The numpy and sleap_io libraries were not used in the old code, and the Optional type hint was not used in the function signatures. The new code does not add any new imports. This change improves the code by removing unnecessary imports, which can improve performance and readability.

• 70-81: The comments describing the shapes of the tensors image, instances, and centroids have been updated. The old comments indicated that the batch size (B) could be greater than 1, while the new comments specify that B=1. This change could potentially impact the functionality of the code if it is expected to handle batch sizes greater than 1. Please verify that all calls to this function throughout the codebase are updated to match the new requirement of B=1.

• 94-96: The instance_image tensor is no longer being squeezed along the 0th dimension. This change could potentially impact the functionality of the code if the downstream code expects the instance_image tensor to have a certain shape. Please verify that all uses of the instance_image tensor throughout the codebase are updated to match the new shape.

sleap_nn/data/augmentation.py (4)

• 2-9: The import of GeometricAugmentationBase2D, Resample, SamplePadding, and warp_affine has been removed. Ensure that these are not used elsewhere in the code or that their functionality has been replaced.

• 130-132: The input_key parameter has been added to the KorniaAugmenter class. This allows the user to specify whether the input is an image or instance. This is a good addition as it provides flexibility to the user.

• 167-173: > Note: This review was outside of the patch, so it was mapped to the patch with the greatest overlap. Original lines [170-194]

The input_key parameter is being stored as an instance variable. This is consistent with the other parameters and allows it to be used in other methods of the class.

• 290-310: The __iter__ method has been updated to handle both image and instance inputs based on the input_key parameter. This is a good change as it allows the augmenter to handle both types of inputs. However, it's important to ensure that the reshaping of the instances tensor is done correctly for both cases.

sleap_nn/data/pipelines.py (4)

• 6-13: The import statement for DictConfig has been changed from from omegaconf.dictconfig import DictConfig to from omegaconf.omegaconf import DictConfig. Please verify if this change is intentional and if it doesn't break the code.

• 27-45: The make_training_pipeline method no longer takes a filename parameter and instead directly uses the data_provider argument. This change implies that the data provider is already initialized with the data source, which is a shift from the previous design where the data source was specified at the pipeline level. This change could impact how the pipeline is used in the rest of the codebase, so please verify that all calls to this method have been updated accordingly. Also, the KorniaAugmenter is now being used with the intensity augmentations from the augmentation_config. This is a new addition and should be verified for correctness.

• 54-70: The KorniaAugmenter is now being used with the geometric augmentations from the augmentation_config. This is a new addition and should be verified for correctness. Also, the ConfidenceMapGenerator and KeyFilter are being used in the same way as before, so no issues here.

• 74-116: A new class SingleInstanceConfmapsPipeline has been added. This class seems to be a variant of the TopdownConfmapsPipeline class for single-instance confidence map models. The make_training_pipeline method in this class is similar to the one in TopdownConfmapsPipeline, but it applies both intensity and geometric augmentations at the same time if use_augmentations is True. Also, it uses the ConfidenceMapGenerator with instance_key and image_key parameters, which is different from the TopdownConfmapsPipeline. Please verify that this new class and its method are working as expected.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between abe7f97 and 89781fb.

Files selected for processing (8)

• sleap_nn/data/augmentation.py (4 hunks)
• sleap_nn/data/confidence_maps.py (3 hunks)
• sleap_nn/data/pipelines.py (2 hunks)
• sleap_nn/data/providers.py (2 hunks)
• tests/data/test_confmaps.py (2 hunks)
• tests/data/test_instance_cropping.py (2 hunks)
• tests/data/test_pipelines.py (7 hunks)
• tests/data/test_providers.py (1 hunks)

Files skipped from review due to trivial changes (4)

• sleap_nn/data/augmentation.py
• sleap_nn/data/providers.py
• tests/data/test_instance_cropping.py
• tests/data/test_providers.py

Additional comments: 21

sleap_nn/data/confidence_maps.py (4)

• 10-18: The function signature of make_confmaps has been updated to handle a batch of points for multiple instances. Ensure that all calls to this function throughout the codebase have been updated to match the new signature. Also, verify that the input tensor points_batch is correctly formatted as (batch_size, n_nodes, 2).

• 32-46: The logic for generating confidence maps has been updated to handle a batch of points for multiple instances. The reshaping of tensors and the calculation of the confidence maps have been updated accordingly. The changes look good and should work as expected.

• 68-81: The ConfidenceMapGenerator class has been updated with new parameters image_key and instance_key. These parameters allow users to specify keys for accessing image and instance data. Ensure that these keys are correctly set in the upstream code where the ConfidenceMapGenerator class is used.

• 83-99: The logic for generating confidence maps in the ConfidenceMapGenerator class has been updated to handle the new image_key and instance_key parameters. The changes look good and should work as expected.

tests/data/test_confmaps.py (4)

• 16-22: The ConfidenceMapGenerator now takes two additional parameters: image_key and instance_key. Ensure that these keys are correctly set in the upstream code and that they correspond to the correct data in the pipeline.

• 28-34: The ConfidenceMapGenerator is being instantiated with different parameters. Ensure that the changes in sigma and output_stride are intentional and that they do not negatively impact the performance or accuracy of the model.

• 25-26: The shape of the confidence maps has changed from (2, 100, 100) to (1, 2, 100, 100). This change suggests that an additional dimension has been added to the confidence maps. Ensure that this change is propagated throughout the codebase and that it does not break any functionality that expects the old shape.

• 37-38: The shape of the confidence maps has changed from (2, 50, 50) to (1, 2, 50, 50). As with the previous comment, ensure that this change is propagated throughout the codebase and that it does not break any functionality that expects the old shape.

41:
The points tensor is now being unsqueezed before being passed to make_confmaps. Ensure that the make_confmaps function can handle this change and that it does not break any functionality that expects the old tensor shape.

56:
The ground truth tensor is now being unsqueezed before being compared to the confidence maps. Ensure that this change is propagated throughout the codebase and that it does not break any functionality that expects the old tensor shape.

tests/data/test_pipelines.py (8)

• 10-13: The import of SingleInstanceConfmapsPipeline is new and it's not used in the old code. Ensure that this new pipeline is compatible with the existing codebase and that it meets the requirements of the tasks it will be used for.

• 22-28: The ConfidenceMapGenerator now takes additional parameters image_key and instance_key. Ensure that these keys are correctly set in the data that is passed to the ConfidenceMapGenerator.

• 46-47: The shape of instance_image and confidence_maps has changed. This could potentially break downstream code that expects the old shape. Please verify that all code that uses these outputs has been updated to handle the new shape.

• 91-92: The way LabelsReader is used has changed. Instead of passing the filename directly, now the labels are loaded first and then passed to LabelsReader. Ensure that this change does not affect the performance or functionality of the code.

• 101-102: Similar to the previous comment, the shape of instance_image and confidence_maps has changed. This could potentially break downstream code that expects the old shape. Please verify that all code that uses these outputs has been updated to handle the new shape.

• 144-145: The way LabelsReader is used has changed. Instead of passing the filename directly, now the labels are loaded first and then passed to LabelsReader. Ensure that this change does not affect the performance or functionality of the code.

• 162-163: The shape of instance_image and confidence_maps has changed. This could potentially break downstream code that expects the old shape. Please verify that all code that uses these outputs has been updated to handle the new shape.

• 166-226: A new test function test_singleinstanceconfmapspipeline has been added. This function tests the SingleInstanceConfmapsPipeline with a single instance example. The test seems to be well-structured and covers the necessary checks for the keys and shapes of the output. However, ensure that this test covers all the necessary aspects of the SingleInstanceConfmapsPipeline.

sleap_nn/data/pipelines.py (5)

• 27-45: The method signature for make_training_pipeline has been changed to remove the filename parameter. Ensure that all calls to this function throughout the codebase have been updated to match the new signature. Also, verify that the data_provider argument now directly provides the data without needing to call from_filename.

• 40-45: The KorniaAugmenter now takes two additional parameters: image_key and instance_key. These parameters allow users to specify keys for accessing image and instance data. This change enhances the flexibility of the data augmentation process by allowing different keys to be used for different stages of the pipeline. However, ensure that the keys "image" and "instances" are present in the data provided by data_provider.

• 58-68: The KorniaAugmenter is now applied twice if use_augmentations is True, once for intensity augmentations and once for geometric augmentations. This change allows for more granular control over the augmentation process. However, ensure that the keys "instance_image" and "instance" are present in the data after the InstanceCropper is applied.

• 70-76: The ConfidenceMapGenerator now takes two additional parameters: image_key and instance_key. These parameters allow users to specify keys for accessing image and instance data. This change enhances the flexibility of the confidence map generation process by allowing different keys to be used for different stages of the pipeline. However, ensure that the keys "instance_image" and "instance" are present in the data after the previous steps are applied.

• 82-130: A new class SingleInstanceConfmapsPipeline has been added. This class is similar to TopdownConfmapsPipeline but is designed for single-instance confidence map models. The make_training_pipeline method in this class creates a training pipeline in a similar way to TopdownConfmapsPipeline, but with some differences in the keys used for the KorniaAugmenter and ConfidenceMapGenerator. Ensure that this new class is integrated properly with the rest of the codebase and that the keys "image" and "instances" are present in the data provided by data_provider.

[coderabbitai]

Review Status

Actionable comments generated: 1

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 89781fb and 2a8ac09.

Files selected for processing (2)

• sleap_nn/data/pipelines.py (2} hunks)
• tests/data/test_pipelines.py (6} hunks)

Additional comments: 14

sleap_nn/data/pipelines.py (5)

• 27-39: The function signature for make_training_pipeline has been changed to remove the filename parameter. Ensure that all calls to this function throughout the codebase have been updated to match the new signature. Also, verify that the data_provider argument now directly provides the data without needing to call from_filename.

• 40-45: The KorniaAugmenter now takes two additional parameters: image_key and instance_key. These keys provide more flexibility in specifying which parts of the data to augment. Ensure that the keys correspond to the correct fields in the data.

• 58-68: The KorniaAugmenter is now applied twice if use_augmentations is True, once for intensity augmentations and once for geometric augmentations. This change allows for more granular control over the augmentation process. However, ensure that the order of these operations is correct and that applying these augmentations twice does not lead to over-augmentation.

• 70-76: The ConfidenceMapGenerator now takes two additional parameters: image_key and instance_key. These keys provide more flexibility in specifying which parts of the data to generate confidence maps for. Ensure that the keys correspond to the correct fields in the data.

• 82-131: A new class SingleInstanceConfmapsPipeline has been introduced. This class appears to be a variant of the TopdownConfmapsPipeline that operates on single instances rather than batches. The structure and methods of this class are similar to those of TopdownConfmapsPipeline, but there are some differences in the details of the pipeline construction. Ensure that this new class is integrated properly with the rest of the codebase and that its behavior is as expected.

tests/data/test_pipelines.py (9)

• 3-13: The import of sleap_io as sio is new. Ensure that this package is installed and available in the environment where the tests are run. Also, the SingleInstanceConfmapsPipeline is imported, which is a new addition to the codebase. Make sure that the implementation of this class is correct and it is working as expected.

• 22-28: The ConfidenceMapGenerator now takes two additional parameters: image_key and instance_key. Ensure that these keys are correctly used in the ConfidenceMapGenerator class and that they match the keys used in the data pipeline.

• 46-48: The shape of instance_image and confidence_maps in the sample has an additional dimension of size 1 at the second position. This might be due to the batched inputs. Ensure that this change is intentional and that it doesn't break any downstream code that uses these outputs.

• 90-92: The LabelsReader is now initialized with a labels parameter instead of a filename parameter. Make sure that the LabelsReader class has been updated to handle this change and that sio.load_slp(minimal_instance) returns the expected data structure.

• 107-108: Similar to the previous comment, the shape of instance_image and confidence_maps in the sample has an additional dimension of size 1 at the second position. Ensure that this change is intentional and that it doesn't break any downstream code that uses these outputs.

• 149-150: The LabelsReader is now initialized with a labels parameter instead of a filename parameter. Make sure that the LabelsReader class has been updated to handle this change and that sio.load_slp(minimal_instance) returns the expected data structure.

• 167-168: Similar to the previous comments, the shape of instance_image and confidence_maps in the sample has an additional dimension of size 1 at the second position. Ensure that this change is intentional and that it doesn't break any downstream code that uses these outputs.

• 171-212: A new test function test_singleinstanceconfmapspipeline is added. This function tests the SingleInstanceConfmapsPipeline with a single instance example. The test includes a comprehensive set of augmentations in the data configuration. Ensure that all these augmentations are implemented correctly in the pipeline.

• 214-231: The SingleInstanceConfmapsPipeline is tested with the LabelsReader initialized with labels parameter. The shapes of image and confidence_maps in the sample are checked. Make sure that the SingleInstanceConfmapsPipeline and LabelsReader are working as expected and the shapes of the outputs are correct.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 2a8ac09 and 4d6e355.

Files selected for processing (2)

• sleap_nn/data/providers.py (3} hunks)
• sleap_nn/evaluation.py (2} hunks)

Files skipped from review due to trivial changes (1)

• sleap_nn/evaluation.py

Additional comments: 3

sleap_nn/data/providers.py (3)

• 1-5: The import statements are unchanged and are correctly placed at the top of the file. No issues found here.

• 11-29: The LabelsReader class has been modified. The user_instances_only attribute and its related logic in the __init__ method have been removed. This change simplifies the class but also removes the ability to filter user instances. Ensure that this change is intentional and that it doesn't break any dependent code. If the filtering logic is needed elsewhere, consider moving it to a separate function or method.

• 38-52: The __iter__ method remains largely unchanged, except for the removal of the filtering logic for user instances. The method still correctly transposes the image from HWC to CHW format, stacks the instances into a numpy array, adds a singleton time dimension to the image and instances, and yields a dictionary containing the image and instances as torch tensors. No issues found here.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 4d6e355 and c68de13.

Files ignored due to filter (1)

• pyproject.toml

Files selected for processing (3)

• sleap_nn/data/pipelines.py (2 hunks)
• sleap_nn/data/providers.py (2 hunks)
• tests/data/test_pipelines.py (6 hunks)

Additional comments: 12

sleap_nn/data/providers.py (2)

• 18-24: The new parameter user_instances_only has been added to the __init__ method of the LabelsReader class. Ensure that all calls to this method throughout the codebase have been updated to match the new signature.

• 34-39: The skeleton attribute in the Labels object has been changed to skeletons. This change could potentially break other parts of the code that are using the skeleton attribute. Please verify that all references to skeleton have been updated to skeletons.

- skeletons=[labels.skeletons],
+ skeletons=labels.skeletons,

sleap_nn/data/pipelines.py (4)

• 24-46: The make_training_pipeline method in the TopdownConfmapsPipeline class has been updated to include data augmentation using the KorniaAugmenter class. The image_key and instance_key arguments are set to "image" and "instances", respectively. Ensure that these keys are consistent with the keys used in the data provider.

• 55-68: The make_training_pipeline method has been further updated to include random cropping and geometric augmentations. The image_key and instance_key arguments are set to "instance_image" and "instance", respectively. Ensure that these keys are consistent with the keys used in the data provider.

• 70-88: The ConfidenceMapGenerator and KeyFilter classes are used to generate confidence maps and filter keys, respectively. The image_key and instance_key arguments are set to "instance_image" and "instance", respectively. Ensure that these keys are consistent with the keys used in the data provider.

• 104-146: The make_training_pipeline method in the SingleInstanceConfmapsPipeline class has been updated to include data augmentation, random cropping, confidence map generation, and key filtering. The image_key and instance_key arguments are set to "image" and "instances", respectively. Ensure that these keys are consistent with the keys used in the data provider.

tests/data/test_pipelines.py (6)

• 1-16: The import statements look fine. The new import for sleap_io is used later in the test functions.

• 19-31: The ConfidenceMapGenerator class is now initialized with additional arguments image_key and instance_key. Ensure that these keys are present in the data passed to this class.

• 43-50: The test_topdownconfmapspipeline function is updated to use the LabelsReader class with the labels argument instead of filename. This change is reflected in line 90. Ensure that the labels argument is correctly passed to the LabelsReader class.

• 87-108: The test_topdownconfmapspipeline function is updated to include additional keys in the gt_sample_keys list. Ensure that these keys are present in the data passed to this function.

• 171-232: A new test function test_singleinstanceconfmapspipeline is added to test the SingleInstanceConfmapsPipeline class. The function seems to be correctly implemented, but ensure that the labels argument is correctly passed to the LabelsReader class (line 217) and that the keys in gt_sample_keys are present in the data passed to this function.

• 233-286: The test_singleinstanceconfmapspipeline function is repeated with different configurations. This is a good practice to test the function under different conditions.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between c68de13 and 84222d1.

Files ignored due to filter (1)

• pyproject.toml

Files selected for processing (3)

• sleap_nn/data/providers.py (1 hunks)
• tests/data/test_instance_cropping.py (2 hunks)
• tests/data/test_pipelines.py (6 hunks)

Files skipped from review due to trivial changes (1)

• sleap_nn/data/providers.py

Additional comments: 8

tests/data/test_instance_cropping.py (2)

• 44-50:
  The updated shape assertions in the test indicate a change in the data structure of instance and instance_image. Verify that these changes are intentional and correctly reflect the expected output shapes of the InstanceCropper.

• 56-59:
  The assertion for centered_instance now compares with gt.unsqueeze(0), indicating that the output has an additional batch dimension. Ensure that this change is consistent with the expected behavior of the InstanceCropper after the recent updates.

tests/data/test_pipelines.py (6)

• 3-3: The import of sleap_io as sio is added, but it's not clear if the old import was removed. Verify that any old import of sleap_io has been removed to avoid redundancy.

• 22-27: The ConfidenceMapGenerator is now initialized with two new parameters image_key and instance_key. Ensure that these keys are correctly used in the downstream code and that their introduction does not break any existing functionality.

• 31-31: The gt_sample_keys list has been updated. Verify that all tests and downstream code that rely on the structure of the samples have been updated accordingly.

• 91-105: The "general" key has been removed from the base_topdown_data_config. Ensure that this change is reflected in all places where base_topdown_data_config is used and that it does not break any existing functionality.

• 177-237: A new test function test_singleinstanceconfmapspipeline has been added. Ensure that this test covers all the necessary cases and that the assertions are correct and meaningful.

• 239-291: The base_singleinstance_data_config is redefined with different augmentation settings. Verify that this redefinition is intentional and that the test is correctly assessing the impact of these changes.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 84222d1 and 2080316.

Files selected for processing (1)

• tests/data/test_pipelines.py (6 hunks)

Additional comments: 4

tests/data/test_pipelines.py (4)

• 3-3: The import of sleap_io as sio is added, which is consistent with the summary stating that a new import statement for the sleap_io module has been added.

• 22-27: The ConfidenceMapGenerator instantiation now includes additional parameters image_key and instance_key. This change aligns with the summary, which mentions that the ConfidenceMapGenerator instantiation has been modified to include these parameters.

• 93-95: The TopdownConfmapsPipeline class has been updated to accept a data_provider parameter instead of a filename, and the test has been modified accordingly. This change is consistent with the summary.

• 175-290: A new test function test_singleinstanceconfmapspipeline has been added for the SingleInstanceConfmapsPipeline class. This addition is consistent with the summary, which mentions the introduction of new test functions and modifications to existing ones related to the new SingleInstanceConfmapsPipeline class.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 2080316 and bd09be8.

Files selected for processing (1)

• tests/data/test_pipelines.py (6 hunks)

Additional comments: 6

tests/data/test_pipelines.py (6)

• 3-3: The import of sleap_io as sio is added as per the summary. Ensure that this module is used wherever necessary in the test cases.

• 22-27: The ConfidenceMapGenerator instantiation now includes additional parameters image_key and instance_key. Verify that these keys are correctly used in the rest of the codebase and that the corresponding data is available in the expected format.

• 93-95: The TopdownConfmapsPipeline now accepts a data_provider parameter instead of a filename. Ensure that all instances of TopdownConfmapsPipeline creation have been updated to use a data_provider and that the LabelsReader is correctly initialized with the necessary data.

• 218-222: A new test function test_singleinstanceconfmapspipeline has been added for SingleInstanceConfmapsPipeline. Ensure that the LabelsReader is correctly initialized with the modified labels and that the make_training_pipeline method is correctly implemented to handle the single instance case.

• 234-235: Verify that the shape of the image and confidence_maps in the sample from SingleInstanceConfmapsPipeline matches the expected output dimensions. This is important to ensure that the pipeline is correctly processing the data.

• 289-290: Similar to the previous comment, verify that the shape of the image and confidence_maps in the sample from SingleInstanceConfmapsPipeline with augmented data matches the expected output dimensions.