Multi anchor crop #41
Conversation
add functionality for randomizing crop anchor
WalkthroughThe updates enhance the Changes
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (invoked as PR comments)
Additionally, you can add CodeRabbit Configration File (
|
…biogtr into aadi/multi-anchor-crop
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 20
Out of diff range and nitpick comments (9)
tests/test_training.py (2)
Line range hint
9-10: Remove unused imports to clean up the code and avoid confusion.- from biogtr.models.global_tracking_transformer import GlobalTrackingTransformer - from omegaconf.DictConfig import DictConfig
Line range hint
46-46: The local variablefeatsis assigned but never used in the test function. Consider removing it if it's not necessary for the test logic.- feats = 128tests/test_datasets.py (1)
Line range hint
52-107: Add assertions to check the values of the crops in addition to their shapes for more comprehensive testing.+ assert torch.allclose(expected_crops, instances[0].get_crops()), "Crops do not match expected values."tests/test_models.py (2)
Line range hint
363-363: Remove the unused variableimg_shape.- img_shape = (1, 100, 100)This change removes the unused variable
img_shape, which is not needed in the function and could potentially lead to confusion.
Line range hint
435-435: Remove the unused variablecfg.- cfg = {"resnet18", "ResNet18_Weights.DEFAULT"}This change removes the unused variable
cfg, which is not needed in the function and could potentially lead to confusion.biogtr/datasets/sleap_dataset.py (1)
Line range hint
9-9: Remove unused import ofwarnings.- import warningsbiogtr/data_structures.py (3)
335-343: Clarify the return type ofanchorproperty.The
anchorproperty's docstring could be clearer about its return type. It currently suggests a list of strings but ends with a return of an empty string if no centroids are defined. This could lead to type inconsistencies.
Line range hint
546-586: Ensure consistent handling of default values inFrameconstructor.The constructor of
Frameclass does not consistently handle default values for parameters likeimg_shapeandinstances. This could lead to unexpected behavior or errors when these values are accessed later in the code.
652-652: Document thetrack_lookupparameter into_slpmethod.The
track_lookupparameter in theto_slpmethod of theFrameclass is not documented in the method's docstring. Adding this documentation would improve code readability and maintainability.
Review Details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (22)
- .gitignore (1 hunks)
- biogtr/data_structures.py (8 hunks)
- biogtr/datasets/sleap_dataset.py (5 hunks)
- biogtr/inference/boxes.py (6 hunks)
- biogtr/models/init.py (1 hunks)
- biogtr/models/attention_head.py (2 hunks)
- biogtr/models/embedding.py (9 hunks)
- biogtr/models/global_tracking_transformer.py (3 hunks)
- biogtr/models/mlp.py (1 hunks)
- biogtr/models/model_utils.py (1 hunks)
- biogtr/models/transformer.py (2 hunks)
- biogtr/models/visual_encoder.py (3 hunks)
- biogtr/training/configs/base.yaml (1 hunks)
- biogtr/training/losses.py (1 hunks)
- environment.yml (1 hunks)
- environment_cpu.yml (1 hunks)
- environment_osx-arm64.yml (1 hunks)
- tests/configs/base.yaml (5 hunks)
- tests/test_data_structures.py (1 hunks)
- tests/test_datasets.py (2 hunks)
- tests/test_models.py (6 hunks)
- tests/test_training.py (2 hunks)
Files skipped from review due to trivial changes (3)
- .gitignore
- environment.yml
- environment_cpu.yml
Additional Context Used
Ruff (14)
biogtr/datasets/sleap_dataset.py (1)
9-9:
warningsimported but unusedbiogtr/models/__init__.py (4)
3-3:
.attention_head.ATTWeightHeadimported but unused; consider removing, adding to__all__, or using a redundant alias
5-5:
.embedding.Embeddingimported but unused; consider removing, adding to__all__, or using a redundant alias
6-6:
.transformer.Transformerimported but unused; consider removing, adding to__all__, or using a redundant alias
7-7:
.visual_encoder.VisualEncoderimported but unused; consider removing, adding to__all__, or using a redundant aliasbiogtr/models/visual_encoder.py (1)
3-3:
typing.Tupleimported but unusedtests/test_datasets.py (1)
8-8:
biogtr.models.model_utils.get_deviceimported but unusedtests/test_models.py (4)
243-243: Local variable
d_modelis assigned to but never used
250-250: Local variable
timesis assigned to but never used
363-363: Local variable
img_shapeis assigned to but never used
435-435: Local variable
cfgis assigned to but never usedtests/test_training.py (3)
9-9:
biogtr.models.global_tracking_transformer.GlobalTrackingTransformerimported but unused
10-10:
omegaconf.DictConfigimported but unused
46-46: Local variable
featsis assigned to but never used
Additional comments not posted (15)
environment_osx-arm64.yml (1)
20-21: Ensure the addition oftimmaligns with the project's dependency management strategy. Iftimmis a crucial dependency for the new features, this change is justified.biogtr/models/attention_head.py (1)
27-28: The use of theMLPclass for projections in the attention mechanism is a good design choice, promoting modularity and reuse of the MLP functionality.biogtr/models/mlp.py (1)
46-60: The forward method of theMLPclass is correctly implemented, ensuring that each layer is followed by a ReLU activation and dropout, except for the last layer. This is a standard practice for MLPs.biogtr/models/visual_encoder.py (1)
20-53: TheVisualEncoderclass has been effectively updated to utilize thetimmlibrary for model creation. This aligns with the PR's objectives to enhance model flexibility and performance.tests/configs/base.yaml (2)
Line range hint
3-22: The updates to the model configuration in the YAML file are consistent with the changes in the model's codebase, ensuring that the configuration aligns with the new functionalities.
Line range hint
68-86: The dataset configuration has been updated to include new parameters likeanchors, which is crucial for the new multi-anchor cropping feature. This change is necessary and well-implemented.biogtr/training/configs/base.yaml (1)
17-20: The model configuration updates in this YAML file are consistent with the intended changes to the model's architecture and functionalities. This ensures that the configuration is in sync with the code changes.biogtr/models/global_tracking_transformer.py (2)
Line range hint
16-64: Ensure validation forencoder_cfgto prevent runtime errors.
61-61: Verify thatFrameobjects havehas_instances()andhas_features()methods.biogtr/inference/boxes.py (4)
36-36: LGTM! Good use of assertion to ensure tensor dimensionality.
66-66: LGTM! Correct handling of 3D tensors in area calculation.
80-83: LGTM! Correct implementation of clipping with 3D tensor support.
100-101: LGTM! Correct identification of non-empty boxes with 3D tensor support.biogtr/training/losses.py (1)
55-55: Ensure NaN values are expected and handled correctly throughout the model's pipeline.Verification successful
Handling of NaN values is evident across various parts of the model's pipeline, including data structures, visualization, training losses, inference metrics, and dataset processing. This indicates a comprehensive approach to managing NaN values, which aligns with the use of
torch.nanmeanin theAssoLossclass to handle NaN values effectively.Here are the specific files and contexts where NaN handling is implemented:
- Data Structures: Usage of
nanmeanfor bounding box calculations.- Visualization: Interpolation and masking strategies to handle NaN values in visualization data.
- Training Losses: Direct use of
torch.nanmeanin theforwardmethod ofAssoLoss.- Inference Metrics: Initialization of evaluation matrices with NaN values and handling during metric calculations.
- Datasets: Various checks and handling mechanisms for NaN values in dataset processing functions.
This comprehensive handling suggests that the model's pipeline is designed to expect and correctly manage NaN values, supporting the changes made in the
AssoLossclass.Scripts Executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for handling of NaN values in the model pipeline. # Test: Search for NaN handling in the model pipeline. Expect: Handling logic for NaN values. rg --type python 'nan'Length of output: 53
Script:
#!/bin/bash # Description: Check for handling of NaN values in the model pipeline. # Test: Search for NaN handling in the model pipeline. Expect: Handling logic for NaN values. rg --type py 'nan'Length of output: 2247
biogtr/models/transformer.py (1)
174-174: Proper handling of NaN values inpred_box.
| def test_multianchor_embedding(): | ||
| frames = 32 | ||
| objects = 10 | ||
| d_model = 256 |
There was a problem hiding this comment.
Remove the unused variable d_model.
- d_model = 256This change removes the unused variable d_model, which is not needed in the function and could potentially lead to confusion.
Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| d_model = 256 |
biogtr/datasets/sleap_dataset.py
Outdated
| anchors: One of: | ||
| * a string indicating a single node to center crops around | ||
| * a list of skeleton node names to be used as the center of crops | ||
| * an int indicating the number of anchors to randomly select | ||
| anchors: One of: | ||
| * a string indicating a single node to center crops around | ||
| * a list of skeleton node names to be used as the center of crops | ||
| * an int indicating the number of anchors to randomly select |
There was a problem hiding this comment.
Remove duplicate documentation and code for anchors parameter.
- anchors: One of:
- * a string indicating a single node to center crops around
- * a list of skeleton node names to be used as the center of crops
- * an int indicating the number of anchors to randomly selectCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| anchors: One of: | |
| * a string indicating a single node to center crops around | |
| * a list of skeleton node names to be used as the center of crops | |
| * an int indicating the number of anchors to randomly select | |
| anchors: One of: | |
| * a string indicating a single node to center crops around | |
| * a list of skeleton node names to be used as the center of crops | |
| * an int indicating the number of anchors to randomly select | |
| anchors: One of: | |
| * a string indicating a single node to center crops around | |
| * a list of skeleton node names to be used as the center of crops | |
| * an int indicating the number of anchors to randomly select |
| if not torch.nn.parameter.is_lazy(p) and p.dim() > 1: | ||
| try: | ||
| nn.init.xavier_uniform_(p) | ||
| except ValueError as e: | ||
| print(f"Failed Trying to initialize {p}") | ||
| raise (e) |
There was a problem hiding this comment.
Handle exceptions more gracefully during parameter initialization.
- except ValueError as e:
- print(f"Failed Trying to initialize {p}")
- raise (e)
+ except ValueError as e:
+ logging.error(f"Failed to initialize parameter {p}: {e}")
+ raiseCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if not torch.nn.parameter.is_lazy(p) and p.dim() > 1: | |
| try: | |
| nn.init.xavier_uniform_(p) | |
| except ValueError as e: | |
| print(f"Failed Trying to initialize {p}") | |
| raise (e) | |
| if not torch.nn.parameter.is_lazy(p) and p.dim() > 1: | |
| try: | |
| nn.init.xavier_uniform_(p) | |
| except ValueError as e: | |
| logging.error(f"Failed to initialize parameter {p}: {e}") | |
| raise |
| @@ -60,7 +58,10 @@ def __init__( | |||
| """ | |||
| super().__init__() | |||
|
|
|||
| self.visual_encoder = VisualEncoder(encoder_model, encoder_cfg, d_model) | |||
| if encoder_cfg is not None: | |||
There was a problem hiding this comment.
Consider refactoring the lambda function for better readability.
- for frame in filter(lambda f: f.has_instances() and not f.has_features(), frames):
+ def frame_filter(f):
+ return f.has_instances() and not f.has_features()
+ for frame in filter(frame_filter, frames):Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if encoder_cfg is not None: | |
| if encoder_cfg is not None: | |
| def frame_filter(f): | |
| return f.has_instances() and not f.has_features() | |
| for frame in filter(frame_filter, frames): |
| bbox[:, :, [0, 2]] /= w | ||
| bbox[:, :, [1, 3]] /= h | ||
|
|
||
| boxes.append(bbox) | ||
| times.append(torch.full((bbox.shape[0],), fidx)) | ||
|
|
||
| boxes = torch.cat(boxes, dim=0) # N x 4 | ||
| boxes = torch.cat(boxes, dim=0) # N, n_anchors, 4 |
There was a problem hiding this comment.
Ensure all frames have consistent dimensions before normalization.
+ if not all(f.img_shape == frames[0].img_shape for f in frames):
+ raise ValueError("All frames must have the same dimensions.")
_, h, w = frames[0].img_shape.flatten()Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| bbox[:, :, [0, 2]] /= w | |
| bbox[:, :, [1, 3]] /= h | |
| boxes.append(bbox) | |
| times.append(torch.full((bbox.shape[0],), fidx)) | |
| boxes = torch.cat(boxes, dim=0) # N x 4 | |
| boxes = torch.cat(boxes, dim=0) # N, n_anchors, 4 | |
| if not all(f.img_shape == frames[0].img_shape for f in frames): | |
| raise ValueError("All frames must have the same dimensions.") | |
| _, h, w = frames[0].img_shape.flatten() | |
| bbox[:, :, [0, 2]] /= w | |
| bbox[:, :, [1, 3]] /= h | |
| boxes.append(bbox) | |
| times.append(torch.full((bbox.shape[0],), fidx)) | |
| boxes = torch.cat(boxes, dim=0) # N, n_anchors, 4 |
| @@ -10,7 +10,7 @@ def test_instance(): | |||
|
|
|||
| gt_track_id = 0 | |||
| pred_track_id = 0 | |||
| bbox = torch.randn((1, 4)) | |||
| bbox = torch.randn((1, 1, 4)) | |||
There was a problem hiding this comment.
Refactor the test to use a testing framework for better structure and output.
+ import unittest
+ class TestInstance(unittest.TestCase):
+ def test_instance(self):
- def test_instance():Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| bbox = torch.randn((1, 1, 4)) | |
| import unittest | |
| class TestInstance(unittest.TestCase): | |
| def test_instance(self): | |
| bbox = torch.randn((1, 1, 4)) |
| return Boxes(self.tensor[item]) | ||
| b = self.tensor[item] | ||
| assert ( | ||
| b.dim() == 2 | ||
| b.dim() == 3 |
There was a problem hiding this comment.
Add more comments to explain the indexing logic for better clarity.
+ # Indexing can be done with an int, slice, or a BoolTensor.
+ # This ensures that the result is always a valid `Boxes` object.Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| return Boxes(self.tensor[item]) | |
| b = self.tensor[item] | |
| assert ( | |
| b.dim() == 2 | |
| b.dim() == 3 | |
| return Boxes(self.tensor[item]) | |
| # Indexing can be done with an int, slice, or a BoolTensor. | |
| # This ensures that the result is always a valid `Boxes` object. | |
| b = self.tensor[item] | |
| assert ( | |
| b.dim() == 3 |
There was a problem hiding this comment.
Actionable comments posted: 0
Out of diff range and nitpick comments (1)
biogtr/inference/post_processing.py (1)
57-60: Consider optimizing the tensor operations for intersection calculation.The current implementation involves multiple tensor operations which might be optimized further. Consider using built-in PyTorch functions that could potentially handle these operations more efficiently if available.
Review Details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (1)
- biogtr/inference/post_processing.py (2 hunks)
Additional comments not posted (4)
biogtr/inference/post_processing.py (4)
Line range hint
25-38: Good use of assertions to ensure necessary variables are present.The function
weight_decay_timecorrectly uses assertions to ensure thatreid_features,T, andkare not None whendecay_timeis greater than 0. This prevents potential runtime errors from missing data.
85-89: Robust handling of empty boxes in IoU calculation.The function
_pairwise_ioueffectively handles cases where the intersection might be zero by assigning NaN and then usingnanmeanto compute the average while ignoring NaN values. This ensures that the function remains robust against empty input boxes.
Line range hint
100-117: Effective error handling and method flexibility in IoU weighting.The function
weight_iouprovides flexibility in how the IoU is used to weight the association matrix and includes robust error handling to ensure thatlast_iousis provided when necessary. It also correctly raises a ValueError for unsupported methods, guiding the user towards correct usage.
Line range hint
130-152: Well-implemented logic for filtering based on maximum center distance.The function
filter_max_center_disteffectively uses tensor operations to compute distances between box centers and filters the association matrix based on these distances. It also uses assertions to ensure all necessary inputs are provided, preventing runtime errors.
There was a problem hiding this comment.
Actionable comments posted: 4
Out of diff range and nitpick comments (1)
tests/test_inference.py (1)
Line range hint
88-88: Remove the unused variableNto clean up the code.- N = N_t * T
Review Details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (9)
- .gitignore (1 hunks)
- biogtr/inference/post_processing.py (3 hunks)
- biogtr/models/attention_head.py (1 hunks)
- biogtr/models/embedding.py (9 hunks)
- biogtr/models/global_tracking_transformer.py (3 hunks)
- biogtr/models/mlp.py (1 hunks)
- biogtr/models/visual_encoder.py (3 hunks)
- tests/test_inference.py (1 hunks)
- tests/test_models.py (6 hunks)
Files skipped from review as they are similar to previous changes (5)
- .gitignore
- biogtr/inference/post_processing.py
- biogtr/models/attention_head.py
- biogtr/models/global_tracking_transformer.py
- biogtr/models/mlp.py
Additional Context Used
Ruff (5)
biogtr/models/visual_encoder.py (1)
3-3:
typing.Tupleimported but unusedtests/test_inference.py (2)
4-4:
pytestimported but unused
88-88: Local variable
Nis assigned to but never usedtests/test_models.py (2)
250-250: Local variable
timesis assigned to but never used
373-373: Local variable
img_shapeis assigned to but never used
Additional comments not posted (17)
biogtr/models/visual_encoder.py (1)
Line range hint
70-89: The forward method is implemented correctly and handles grayscale images appropriately.tests/test_inference.py (2)
46-46: The configuration for the encoder in theGlobalTrackingTransformeris correctly set up for the test.
Line range hint
88-88: The metrics testing is comprehensive and covers various scenarios effectively.tests/test_models.py (10)
6-6: The MLP model is correctly imported and used in the tests.
7-7: The ATTWeightHead model is correctly imported and used in the tests.
52-52: The VisualEncoder is correctly set up with various channel configurations for testing.
243-243: Remove the unused variabled_modelto clean up the code.
69-69: The Embedding class is tested thoroughly for various configurations to ensure exceptions are raised as expected.
118-118: The basic functionality of the Embedding class is correctly tested with various configurations.
200-200: The Embedding class is tested with different keyword arguments to ensure correct handling.
282-282: The TransformerEncoderLayer is correctly tested with and without positional embeddings.
Line range hint
373-373: The TransformerDecoderLayer is correctly tested with and without positional embeddings.
Line range hint
373-373: The full transformer logic is correctly tested and the output shapes are asserted as expected.biogtr/models/embedding.py (4)
149-158: The forward method correctly computes the embeddings and checks the output dimensions.
Line range hint
178-215: The_sine_box_embeddingmethod correctly computes sine positional embeddings for boxes, handling normalization and scaling appropriately.
215-215: The_sine_temp_embeddingmethod correctly computes sine temporal embeddings, handling the computation of denominators and lookup table correctly.
264-306: The_learned_pos_embeddingmethod correctly computes learned positional embeddings for boxes, handling the computation of weights and lookup table correctly.
| if in_chans == 1: | ||
| self.in_chans = 3 | ||
| else: | ||
| self.in_chans = in_chans |
There was a problem hiding this comment.
Review the logic for handling in_chans. Currently, if in_chans is 1, it is forcibly set to 3, which might not be intended for all use cases.
| N = frames * objects | ||
|
|
||
| boxes = torch.rand(size=(N, n_anchors, 4)) | ||
| times = torch.rand(size=(N,)) |
There was a problem hiding this comment.
Remove the unused variable times to clean up the code.
- times = torch.rand(size=(N,))Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| times = torch.rand(size=(N,)) |
biogtr/models/embedding.py
Outdated
| if self.emb_type == "pos" and mlp_cfg is not None: | ||
| if self.mode == "fixed": | ||
| self.mlp = MLP( | ||
| input_dim=n_points * self.features, | ||
| output_dim=self.features, | ||
| **mlp_cfg, | ||
| ) | ||
| else: | ||
| in_dim = (self.features // (4 * n_points)) * (4 * n_points) | ||
| self.mlp = MLP( | ||
| input_dim=in_dim, | ||
| output_dim=self.features, | ||
| **mlp_cfg, | ||
| ) | ||
| else: | ||
| self.mlp = torch.nn.Identity() |
There was a problem hiding this comment.
Review the initialization logic for mlp_cfg. Ensure that it is correctly initialized based on the mode and type of embedding.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #41 +/- ##
==========================================
+ Coverage 69.60% 69.81% +0.21%
==========================================
Files 28 29 +1
Lines 2385 2508 +123
==========================================
+ Hits 1660 1751 +91
- Misses 725 757 +32 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Actionable comments posted: 4
Out of diff range and nitpick comments (2)
biogtr/models/embedding.py (1)
Line range hint
6-103: Constructor ofEmbeddingis well-implemented, handling various configurations and initializing the embeddings correctly.- self.mlp = MLP(...)Ensure the MLP setup line is covered by tests to verify its functionality.
biogtr/datasets/sleap_dataset.py (1)
Line range hint
9-9: Remove unused import.- import warnings
Review Details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (10)
- .gitignore (1 hunks)
- biogtr/datasets/sleap_dataset.py (5 hunks)
- biogtr/inference/post_processing.py (3 hunks)
- biogtr/models/attention_head.py (1 hunks)
- biogtr/models/embedding.py (9 hunks)
- biogtr/models/global_tracking_transformer.py (3 hunks)
- biogtr/models/mlp.py (1 hunks)
- biogtr/models/visual_encoder.py (3 hunks)
- tests/test_inference.py (1 hunks)
- tests/test_models.py (6 hunks)
Files skipped from review as they are similar to previous changes (4)
- .gitignore
- biogtr/inference/post_processing.py
- biogtr/models/attention_head.py
- biogtr/models/global_tracking_transformer.py
Additional Context Used
Ruff (6)
biogtr/datasets/sleap_dataset.py (1)
9-9:
warningsimported but unusedbiogtr/models/visual_encoder.py (1)
3-3:
typing.Tupleimported but unusedtests/test_inference.py (2)
4-4:
pytestimported but unused
88-88: Local variable
Nis assigned to but never usedtests/test_models.py (2)
250-250: Local variable
timesis assigned to but never used
373-373: Local variable
img_shapeis assigned to but never used
GitHub Check Runs (1)
codecov/patch success (22)
biogtr/datasets/sleap_dataset.py: [warning] 95-95: biogtr/datasets/sleap_dataset.py#L95
Added line #L95 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 184-184: biogtr/datasets/sleap_dataset.py#L184
Added line #L184 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 188-188: biogtr/datasets/sleap_dataset.py#L188
Added line #L188 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 294-294: biogtr/datasets/sleap_dataset.py#L294
Added line #L294 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 300-303: biogtr/datasets/sleap_dataset.py#L300-L303
Added lines #L300 - L303 were not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 305-307: biogtr/datasets/sleap_dataset.py#L305-L307
Added lines #L305 - L307 were not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 310-310: biogtr/datasets/sleap_dataset.py#L310
Added line #L310 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 313-313: biogtr/datasets/sleap_dataset.py#L313
Added line #L313 was not covered by tests
biogtr/datasets/sleap_dataset.py: [warning] 322-322: biogtr/datasets/sleap_dataset.py#L322
Added line #L322 was not covered by tests
biogtr/models/embedding.py: [warning] 152-152: biogtr/models/embedding.py#L152
Added line #L152 was not covered by tests
biogtr/models/mlp.py: [warning] 45-45: biogtr/models/mlp.py#L45
Added line #L45 was not covered by tests
Additional comments not posted (18)
biogtr/models/mlp.py (2)
7-43: Constructor setup for MLP looks robust and handles various configurations well.
47-61: Forward method implementation in MLP is correctly applying activations and handling dropout.biogtr/models/visual_encoder.py (1)
Line range hint
70-86: Forward method inVisualEncodercorrectly handles image inputs, including grayscale, and properly uses the feature extractor and output layer.tests/test_inference.py (3)
46-46: Thetest_trackerfunction correctly sets up and tests the tracker module with various configurations.
Line range hint
88-88: Thetest_post_processingfunction effectively tests different post-processing methods under various conditions.
Line range hint
88-88: Thetest_metricsfunction correctly tests the metrics calculation for tracking, ensuring accurate match detection and switch counting.tests/test_models.py (8)
6-6: Thetest_mlpfunction correctly tests the MLP logic, ensuring it processes inputs as expected and produces outputs of the correct shape.
7-7: Thetest_att_weight_headfunction correctly tests the self-attention head logic, ensuring it processes inputs as expected and produces outputs of the correct shape.
52-76: Thetest_encoderfunction correctly tests the feature extractor logic, ensuring it handles different channel configurations and produces outputs of the correct shape.
69-69: Thetest_embedding_validityfunction correctly tests the embedding usage, ensuring exceptions are raised under invalid configurations.
118-118: Thetest_embedding_basicfunction correctly tests the basic embedding logic, ensuring it processes inputs as expected and produces outputs of the correct shape.
200-200: Thetest_embedding_kwargsfunction correctly tests the embedding configuration logic, ensuring it handles different keyword arguments and produces outputs of the correct shape.
240-240: Thetest_multianchor_embeddingfunction correctly tests the multi-anchor embedding logic, ensuring it handles multiple anchors and produces outputs of the correct shape.
282-282: Thetest_transformer_encoderfunction correctly tests the transformer encoder layer logic, ensuring it handles positional embeddings and produces outputs of the correct shape.biogtr/models/embedding.py (3)
143-158: Forward method inEmbeddingcorrectly processes inputs through the appropriate embedding function and checks the output shape.
Line range hint
178-215: The_sine_box_embeddingmethod correctly computes sine positional embeddings for boxes, handling different configurations such as normalization and scaling.
264-306: The_learned_pos_embeddingmethod correctly computes learned positional embeddings for boxes, handling different configurations such as over_boxes.biogtr/datasets/sleap_dataset.py (1)
42-46: Remove duplicate documentation foranchorsparameter.
| if isinstance(anchors, int): | ||
| self.anchors = anchors | ||
| elif isinstance(anchors, str): | ||
| self.anchors = [anchors.lower()] | ||
| else: | ||
| self.anchors = [anchor.lower() for anchor in anchors] | ||
|
|
||
| if ( | ||
| isinstance(self.anchors, list) and len(self.anchors) == 0 | ||
| ) or self.anchors == 0: | ||
| raise ValueError(f"Must provide at least one anchor but got {self.anchors}") |
There was a problem hiding this comment.
Refactor the handling of anchors to simplify the logic.
- if isinstance(anchors, int):
- self.anchors = anchors
- elif isinstance(anchors, str):
- self.anchors = [anchors.lower()]
- else:
- self.anchors = [anchor.lower() for anchor in anchors]
- if (
- isinstance(self.anchors, list) and len(self.anchors) == 0
- ) or self.anchors == 0:
- raise ValueError(f"Must provide at least one anchor but got {self.anchors}")
+ # Simplified handling of anchors
+ self.anchors = [anchors.lower()] if isinstance(anchors, str) else anchors
+ if not self.anchors:
+ raise ValueError("Must provide at least one anchor.")Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if isinstance(anchors, int): | |
| self.anchors = anchors | |
| elif isinstance(anchors, str): | |
| self.anchors = [anchors.lower()] | |
| else: | |
| self.anchors = [anchor.lower() for anchor in anchors] | |
| if ( | |
| isinstance(self.anchors, list) and len(self.anchors) == 0 | |
| ) or self.anchors == 0: | |
| raise ValueError(f"Must provide at least one anchor but got {self.anchors}") | |
| # Simplified handling of anchors | |
| self.anchors = [anchors.lower()] if isinstance(anchors, str) else anchors | |
| if not self.anchors: | |
| raise ValueError("Must provide at least one anchor.") |
| if c == 1: | ||
| img = np.concatenate( | ||
| [img, img, img], axis=-1 | ||
| ) # convert to grayscale to rgb | ||
|
|
There was a problem hiding this comment.
Optimize the conversion of grayscale images to RGB.
- img = np.concatenate(
- [img, img, img], axis=-1
- ) # convert to grayscale to rgb
+ img = np.stack([img] * 3, axis=-1) # convert grayscale to RGBCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if c == 1: | |
| img = np.concatenate( | |
| [img, img, img], axis=-1 | |
| ) # convert to grayscale to rgb | |
| if c == 1: | |
| img = np.stack([img] * 3, axis=-1) # convert grayscale to RGB |
| if isinstance(self.anchors, int): | ||
| anchors_to_choose = list(pose.keys()) + ["midpoint"] | ||
| anchors = np.random.choice(anchors_to_choose, self.anchors) | ||
| else: | ||
| if self.verbose: | ||
| warnings.warn( | ||
| f"{self.anchor} not in {[key for key in pose.keys()]}! Defaulting to midpoint" | ||
| ) | ||
| anchor = "midpoint" | ||
| anchors = self.anchors | ||
|
|
||
| if anchor != "midpoint": | ||
| centroid = pose[anchor] | ||
| for anchor in anchors: | ||
| if anchor == "midpoint" or anchor == "centroid": | ||
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | ||
|
|
||
| elif anchor in pose: | ||
| centroid = np.array(pose[anchor]) | ||
| if np.isnan(centroid).any(): | ||
| centroid = np.array([np.nan, np.nan]) | ||
|
|
||
| if np.isnan(centroid).any(): | ||
| elif anchor not in pose and len(anchors) == 1: | ||
| anchor = "midpoint" | ||
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | ||
| else: | ||
| # print(f'{self.anchor} not an available option amongst {pose.keys()}. Using midpoint') | ||
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | ||
|
|
||
| bbox = data_utils.pad_bbox( | ||
| data_utils.get_bbox(centroid, self.crop_size), | ||
| padding=self.padding, | ||
| ) | ||
| elif anchor in pose: | ||
| centroid = np.array(pose[anchor]) | ||
| if np.isnan(centroid).any(): | ||
| centroid = np.array([np.nan, np.nan]) | ||
|
|
||
| elif anchor not in pose and len(anchors) == 1: | ||
| anchor = "midpoint" | ||
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | ||
|
|
||
| else: | ||
| centroid = np.array([np.nan, np.nan]) | ||
|
|
||
| if np.isnan(centroid).all(): | ||
| bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan]) | ||
|
|
||
| else: | ||
| bbox = data_utils.pad_bbox( | ||
| data_utils.get_bbox(centroid, self.crop_size), | ||
| padding=self.padding, | ||
| ) | ||
|
|
||
| if bbox.isnan().all(): | ||
| crop = torch.zeros( | ||
| c, | ||
| self.crop_size + 2 * self.padding, | ||
| self.crop_size + 2 * self.padding, | ||
| dtype=img.dtype, | ||
| ) | ||
| else: | ||
| crop = data_utils.crop_bbox(img, bbox) | ||
|
|
||
| crops.append(crop) | ||
| centroids[anchor] = centroid | ||
| boxes.append(bbox) | ||
|
|
||
| if len(crops) > 0: | ||
| crops = torch.concat(crops, dim=0) | ||
|
|
||
| crop = data_utils.crop_bbox(img, bbox) | ||
| if len(boxes) > 0: | ||
| boxes = torch.stack(boxes, dim=0) |
There was a problem hiding this comment.
Refactor the handling of anchor points to simplify the logic.
- if isinstance(self.anchors, int):
- anchors_to_choose = list(pose.keys()) + ["midpoint"]
- anchors = np.random.choice(anchors_to_choose, self.anchors)
- else:
- anchors = self.anchors
- for anchor in anchors:
- if anchor == "midpoint" or anchor == "centroid":
- centroid = np.nanmean(np.array(list(pose.values())), axis=0)
- elif anchor in pose:
- centroid = np.array(pose[anchor])
- if np.isnan(centroid).any():
- centroid = np.array([np.nan, np.nan])
- elif anchor not in pose and len(anchors) == 1:
- anchor = "midpoint"
- centroid = np.nanmean(np.array(list(pose.values())), axis=0)
- else:
- centroid = np.array([np.nan, np.nan])
- if np.isnan(centroid).all():
- bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan])
- else:
- bbox = data_utils.pad_bbox(
- data_utils.get_bbox(centroid, self.crop_size),
- padding=self.padding,
- )
- if bbox.isnan().all():
- crop = torch.zeros(
- c,
- self.crop_size + 2 * self.padding,
- self.crop_size + 2 * self.padding,
- dtype=img.dtype,
- )
- else:
- crop = data_utils.crop_bbox(img, bbox)
- crops.append(crop)
- centroids[anchor] = centroid
- boxes.append(bbox)
- if len(crops) > 0:
- crops = torch.concat(crops, dim=0)
- if len(boxes) > 0:
- boxes = torch.stack(boxes, dim=0)
+ # Proposed refactoring code hereCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if isinstance(self.anchors, int): | |
| anchors_to_choose = list(pose.keys()) + ["midpoint"] | |
| anchors = np.random.choice(anchors_to_choose, self.anchors) | |
| else: | |
| if self.verbose: | |
| warnings.warn( | |
| f"{self.anchor} not in {[key for key in pose.keys()]}! Defaulting to midpoint" | |
| ) | |
| anchor = "midpoint" | |
| anchors = self.anchors | |
| if anchor != "midpoint": | |
| centroid = pose[anchor] | |
| for anchor in anchors: | |
| if anchor == "midpoint" or anchor == "centroid": | |
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | |
| elif anchor in pose: | |
| centroid = np.array(pose[anchor]) | |
| if np.isnan(centroid).any(): | |
| centroid = np.array([np.nan, np.nan]) | |
| if np.isnan(centroid).any(): | |
| elif anchor not in pose and len(anchors) == 1: | |
| anchor = "midpoint" | |
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | |
| else: | |
| # print(f'{self.anchor} not an available option amongst {pose.keys()}. Using midpoint') | |
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | |
| bbox = data_utils.pad_bbox( | |
| data_utils.get_bbox(centroid, self.crop_size), | |
| padding=self.padding, | |
| ) | |
| elif anchor in pose: | |
| centroid = np.array(pose[anchor]) | |
| if np.isnan(centroid).any(): | |
| centroid = np.array([np.nan, np.nan]) | |
| elif anchor not in pose and len(anchors) == 1: | |
| anchor = "midpoint" | |
| centroid = np.nanmean(np.array(list(pose.values())), axis=0) | |
| else: | |
| centroid = np.array([np.nan, np.nan]) | |
| if np.isnan(centroid).all(): | |
| bbox = torch.tensor([np.nan, np.nan, np.nan, np.nan]) | |
| else: | |
| bbox = data_utils.pad_bbox( | |
| data_utils.get_bbox(centroid, self.crop_size), | |
| padding=self.padding, | |
| ) | |
| if bbox.isnan().all(): | |
| crop = torch.zeros( | |
| c, | |
| self.crop_size + 2 * self.padding, | |
| self.crop_size + 2 * self.padding, | |
| dtype=img.dtype, | |
| ) | |
| else: | |
| crop = data_utils.crop_bbox(img, bbox) | |
| crops.append(crop) | |
| centroids[anchor] = centroid | |
| boxes.append(bbox) | |
| if len(crops) > 0: | |
| crops = torch.concat(crops, dim=0) | |
| crop = data_utils.crop_bbox(img, bbox) | |
| if len(boxes) > 0: | |
| boxes = torch.stack(boxes, dim=0) | |
| # Proposed refactoring code here |
biogtr/models/visual_encoder.py
Outdated
| def __init__( | ||
| self, | ||
| model_name: str = "resnet18", | ||
| d_model: int = 512, | ||
| in_chans: int = 3, | ||
| pretrained: bool = False, | ||
| **kwargs, | ||
| ): | ||
| """Initialize Visual Encoder. | ||
|
|
||
| Args: | ||
| model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). | ||
| cfg (dict): Dictionary of arguments to pass to the CNN constructor, | ||
| e.g: `cfg = {"weights": "ResNet18_Weights.DEFAULT"}` | ||
| d_model (int): Output embedding dimension. | ||
| in_chans: the number of input channels of the image. | ||
| pretrained: whether or not to use pretrained weights from hugging_face | ||
| kwargs: see `timm.create_model` for kwargs. | ||
| """ | ||
| super().__init__() | ||
|
|
||
| self.model_name = model_name | ||
| self.model_name = model_name.lower() | ||
| self.d_model = d_model | ||
|
|
||
| self.feature_extractor, out_dim = self.select_feature_extractor(model_name, cfg) | ||
|
|
||
| self.feature_extractor = torch.nn.Sequential( | ||
| *list(self.feature_extractor.children())[:-1] | ||
| if in_chans == 1: | ||
| self.in_chans = 3 | ||
| else: | ||
| self.in_chans = in_chans | ||
|
|
||
| self.feature_extractor = timm.create_model( | ||
| model_name=self.model_name, | ||
| in_chans=self.in_chans, | ||
| pretrained=pretrained, | ||
| num_classes=0, | ||
| **kwargs, |
There was a problem hiding this comment.
Constructor of VisualEncoder is well-implemented, handling different configurations and initializing the model correctly.
- from typing import TupleRemove the unused import to clean up the code.
Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| def __init__( | |
| self, | |
| model_name: str = "resnet18", | |
| d_model: int = 512, | |
| in_chans: int = 3, | |
| pretrained: bool = False, | |
| **kwargs, | |
| ): | |
| """Initialize Visual Encoder. | |
| Args: | |
| model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). | |
| cfg (dict): Dictionary of arguments to pass to the CNN constructor, | |
| e.g: `cfg = {"weights": "ResNet18_Weights.DEFAULT"}` | |
| d_model (int): Output embedding dimension. | |
| in_chans: the number of input channels of the image. | |
| pretrained: whether or not to use pretrained weights from hugging_face | |
| kwargs: see `timm.create_model` for kwargs. | |
| """ | |
| super().__init__() | |
| self.model_name = model_name | |
| self.model_name = model_name.lower() | |
| self.d_model = d_model | |
| self.feature_extractor, out_dim = self.select_feature_extractor(model_name, cfg) | |
| self.feature_extractor = torch.nn.Sequential( | |
| *list(self.feature_extractor.children())[:-1] | |
| if in_chans == 1: | |
| self.in_chans = 3 | |
| else: | |
| self.in_chans = in_chans | |
| self.feature_extractor = timm.create_model( | |
| model_name=self.model_name, | |
| in_chans=self.in_chans, | |
| pretrained=pretrained, | |
| num_classes=0, | |
| **kwargs, | |
| def __init__( | |
| self, | |
| model_name: str = "resnet18", | |
| d_model: int = 512, | |
| in_chans: int = 3, | |
| pretrained: bool = False, | |
| **kwargs, | |
| ): | |
| """Initialize Visual Encoder. | |
| Args: | |
| model_name (str): Name of the CNN architecture to use (e.g. "resnet18", "resnet50"). | |
| d_model (int): Output embedding dimension. | |
| in_chans: the number of input channels of the image. | |
| pretrained: whether or not to use pretrained weights from hugging_face | |
| kwargs: see `timm.create_model` for kwargs. | |
| """ | |
| super().__init__() | |
| self.model_name = model_name.lower() | |
| self.d_model = d_model | |
| if in_chans == 1: | |
| self.in_chans = 3 | |
| else: | |
| self.in_chans = in_chans | |
| self.feature_extractor = timm.create_model( | |
| model_name=self.model_name, | |
| in_chans=self.in_chans, | |
| pretrained=pretrained, | |
| num_classes=0, | |
| **kwargs, |
There was a problem hiding this comment.
Actionable comments posted: 3
Review Details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (3)
- biogtr/models/embedding.py (9 hunks)
- biogtr/models/visual_encoder.py (3 hunks)
- tests/test_models.py (6 hunks)
Additional Context Used
Ruff (3)
biogtr/models/visual_encoder.py (1)
3-3:
typing.Tupleimported but unusedtests/test_models.py (2)
302-302: Local variable
timesis assigned to but never used
425-425: Local variable
img_shapeis assigned to but never used
Additional comments not posted (12)
biogtr/models/visual_encoder.py (4)
43-46: Review the logic for handlingin_chans. Currently, ifin_chansis 1, it is forcibly set to 3, which might not be intended for all use cases.
60-112: Theselect_feature_extractormethod is well-implemented, handling different configurations and initializing the model correctly.
114-126: Theencoder_dimmethod correctly computes the embedding dimension based on the model.
Line range hint
128-142: Theforwardmethod correctly handles grayscale images by repeating the channels and is well-implemented.biogtr/models/embedding.py (4)
Line range hint
34-92: The__init__method correctly initializes the MLP based onn_pointsandmlp_cfgparameters.
Line range hint
178-215: The_sine_box_embeddingmethod correctly handles multiple anchors and includes the necessary MLP transformations for positional embeddings.
264-306: The_learned_pos_embeddingmethod correctly handles multiple anchors and includes the necessary exception handling for incorrectn_points.
143-158: Theforwardmethod correctly includes the necessary exception handling for embedding dimension mismatches and logic for correct dimension projection.tests/test_models.py (4)
45-85: Thetest_encoder_timmfunction correctly tests theVisualEncoderwith different input channels and verifies the output shape.
88-128: Thetest_encoder_torchfunction correctly tests theVisualEncoderwith different input channels and verifies the output shape.
Line range hint
170-218: Thetest_embedding_basicfunction correctly tests theEmbeddingclass with different configurations and verifies the output shape.
292-333: Thetest_multianchor_embeddingfunction correctly tests theEmbeddingclass with multiple anchors and verifies the output shape.
| @@ -3,6 +3,7 @@ | |||
| from typing import Tuple | |||
There was a problem hiding this comment.
Remove the unused import of Tuple to clean up the code.
- from typing import TupleCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| from typing import Tuple |
| @@ -3,6 +3,7 @@ | |||
| from typing import Tuple, Optional | |||
| import math | |||
| import torch | |||
| from biogtr.models.mlp import MLP | |||
There was a problem hiding this comment.
Remove the unused import of Tuple to clean up the code.
- from typing import TupleCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| from biogtr.models.mlp import MLP | |
| from biogtr.models.mlp import MLP |
| @@ -3,7 +3,8 @@ | |||
| import pytest | |||
| import torch | |||
| from biogtr.data_structures import Frame, Instance | |||
| from biogtr.models.attention_head import MLP, ATTWeightHead | |||
| from biogtr.models.mlp import MLP | |||
There was a problem hiding this comment.
Remove the unused import of ATTWeightHead to clean up the code.
- from biogtr.models.attention_head import ATTWeightHeadCommittable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| from biogtr.models.mlp import MLP | |
| from biogtr.models.mlp import MLP |
As per #32, the major failure mode we're seeing in the model while animal tracking is when nodes disappear causing a shift in visual features and a switch to occur. We've tried randomly selecting available node, using the pose centroid as the crop, and reducing the crop size but none have worked. Instead here, we add functionality to stack crops from as many nodes along the channel dimension and use that as input into the model.
Summary by CodeRabbit
New Features
Bug Fixes
Enhancements
timmfor better model support.torch.nn.LazyLinearfor output layer initialization.Tests