sleap_nn/config/model_config.py (2)

1-2: Consider adding explicit attribute definitions.

The ModelConfig class is correctly using the @attr.s decorator with auto_attribs=True. However, the class body is empty, which is unusual. Consider adding explicit attribute definitions to improve code clarity and enable IDE autocompletion.

Here's an example of how you could define the attributes:

@attr.s(auto_attribs=True)
class ModelConfig:
    backbone: dict
    heads: dict
    base_checkpoint: str

This will make the class structure more explicit and easier to understand at a glance.

🧰 Tools

🪛 Ruff

1-1: Undefined name attr

(F821)

---

3-9: Docstring looks good, minor formatting suggestion.

The docstring provides clear and informative descriptions of the class and its attributes. Well done!

Consider adding a period at the end of the last line for consistency:

-        base_checkpoint: Path to model folder for loading a checkpoint. Should contain the .h5 file
+        base_checkpoint: Path to model folder for loading a checkpoint. Should contain the .h5 file.

sleap_nn/config/training_job.py (1)

1-42: Overall assessment and recommendations

The TrainingJobConfig class provides a well-structured and documented foundation for managing training job configurations. However, there are a few improvements needed to make it complete and error-free:

1. Add the missing imports for attr, DataConfig, and other required classes.
2. Implement the remaining attributes mentioned in the class docstring.
3. Consider adding type hints for all attributes to improve code readability and catch potential type-related issues early.
4. If not already present in your project, consider adding a requirements.txt or setup.py file to specify the attrs library as a dependency.

Once these changes are made, the TrainingJobConfig class will be a robust and well-documented configuration management solution for your training jobs.

🧰 Tools

🪛 Ruff

27-27: Undefined name attr

(F821)

---

42-42: Undefined name DataConfig

(F821)

---

42-42: Undefined name attr

(F821)

---

42-42: Undefined name DataConfig

(F821)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (2)

1-3: Remove unused import

The List type from typing is imported but not used in the file. Consider removing it to keep the imports clean.

Apply this diff to remove the unused import:

-from typing import Optional, Tuple, List, Dict
+from typing import Optional, Tuple, Dict

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

---

38-51: Approve IntensityConfig class and suggest comment update

The IntensityConfig class is well-structured with appropriate attributes for intensity-related augmentation options. However, there's a comment indicating that this part is still under review.

Once the review of this section is complete, consider removing or updating the following comment:

# still looking over this part and below (function to check if is greater than 0)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (1)

1-3: Remove unused import

The List type is imported but not used in the file. Consider removing it to keep the imports clean.

Apply this diff to remove the unused import:

-from typing import Optional, Tuple, List, Dict
+from typing import Optional, Tuple, Dict

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (1)

1-3: Remove unused import

The List type is imported from typing but not used in the file. Consider removing it to keep the imports clean.

Apply this diff to remove the unused import:

-from typing import Optional, Tuple, List, Dict
+from typing import Optional, Tuple, Dict

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (2)

1-3: Optimize imports

Please make the following changes to the import statements:

1. Remove the unused import of List.
2. Add Union to the imports from typing.

Apply this diff to optimize the imports:

 import attrs
 from omegaconf import MISSING
-from typing import Optional, Tuple, List, Dict
+from typing import Optional, Tuple, Dict, Union

This change will remove the unused import and add the necessary Union type.

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

---

94-127: Approve GeometricConfig class and fix typo

The GeometricConfig class is well-structured with appropriate attributes and default values. The docstring provides clear explanations for each attribute.

There's a small typo in the docstring for the erase_ratio_min attribute. Please fix it as follows:

-        erase_ration_min: (float) Minimum value of range of aspect ratio of erased area. Default: 1.
+        erase_ratio_min: (float) Minimum value of range of aspect ratio of erased area. Default: 1.

This change will correct the attribute name in the docstring.

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (4)

1-3: Remove unused import

The List type is imported from typing but not used in the file. To keep the imports clean, you can remove it.

Apply this change:

- from typing import Optional, Tuple, List, Dict
+ from typing import Optional, Tuple, Dict

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

---

12-35: Approve DataConfig class and suggest forward references

The DataConfig class is well-structured with clear attribute definitions and a comprehensive docstring. However, there are undefined names for PreprocessingConfig and AugmentationConfig.

To resolve the undefined names and prevent potential circular imports, use forward references for these classes. Apply this diff:

+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sleap_nn.config.data_config import PreprocessingConfig, AugmentationConfig

 @attrs.define
 class DataConfig:
     # ... (existing code)
-    preprocessing: PreprocessingConfig = attrs.field(factory=PreprocessingConfig)
-    augmentation_config: Optional[AugmentationConfig] = None
+    preprocessing: 'PreprocessingConfig' = attrs.field(factory=lambda: PreprocessingConfig())
+    augmentation_config: Optional['AugmentationConfig'] = None

This change will resolve the undefined names issue and prevent potential circular imports.

🧰 Tools

🪛 Ruff

33-33: Undefined name PreprocessingConfig

(F821)

---

33-33: Undefined name PreprocessingConfig

(F821)

---

35-35: Undefined name AugmentationConfig

(F821)

---

38-56: Approve PreprocessingConfig class and add missing import

The PreprocessingConfig class is well-structured with appropriate attributes for preprocessing options. However, the Union type is used but not imported.

Add the following import at the top of the file:

- from typing import Optional, Tuple, Dict
+ from typing import Optional, Tuple, Dict, Union

This will resolve the undefined Union type issue.

🧰 Tools

🪛 Ruff

54-54: Undefined name Union

(F821)

---

58-70: Approve AugmentationConfig class and use forward references

The AugmentationConfig class is well-structured, but there are undefined names for IntensityConfig and GeometricConfig.

To resolve the undefined names and prevent potential circular imports, use forward references for these classes. Apply this diff:

 @attrs.define
 class AugmentationConfig:
     # ... (existing code)
     random_crop: Optional[Dict[str, Optional[float]]] = None
-    intensity: Optional[IntensityConfig] = attrs.field(default=None)
-    geometric: Optional[GeometricConfig] = attrs.field(default=None)
+    intensity: Optional['IntensityConfig'] = attrs.field(default=None)
+    geometric: Optional['GeometricConfig'] = attrs.field(default=None)

This change will resolve the undefined names issue and prevent potential circular imports.

🧰 Tools

🪛 Ruff

69-69: Undefined name IntensityConfig

(F821)

---

70-70: Undefined name GeometricConfig

(F821)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/data_config.py (3)

1-3: Approve imports with minor suggestion

The imports look good overall. However, there's an unused import that can be removed.

Remove the unused List import:

-from typing import Optional, Tuple, List, Dict
+from typing import Optional, Tuple, Dict

🧰 Tools

🪛 Ruff

3-3: typing.List imported but unused

Remove unused import: typing.List

(F401)

---

13-36: Approve DataConfig class and suggest forward references

The DataConfig class is well-structured with clear documentation. However, there are undefined names for PreprocessingConfig and AugmentationConfig.

To avoid potential circular imports, use forward references for these classes. Add the following import at the top of the file and modify the class attributes:

from __future__ import annotations

Then update the attributes:

-    preprocessing: PreprocessingConfig = attrs.field(factory=PreprocessingConfig)
-    augmentation_config: Optional[AugmentationConfig] = None
+    preprocessing: 'PreprocessingConfig' = attrs.field(factory=lambda: PreprocessingConfig())
+    augmentation_config: Optional['AugmentationConfig'] = None

This change will resolve the undefined names issue and prevent potential circular imports.

🧰 Tools

🪛 Ruff

34-34: Undefined name PreprocessingConfig

(F821)

---

34-34: Undefined name PreprocessingConfig

(F821)

---

36-36: Undefined name AugmentationConfig

(F821)

---

60-72: Approve AugmentationConfig class and suggest forward references

The AugmentationConfig class is well-structured with clear documentation. However, there are undefined names for IntensityConfig and GeometricConfig.

To avoid potential circular imports, use forward references for these classes. Modify the class attributes as follows:

-    intensity: Optional[IntensityConfig] = attrs.field(default=None)
-    geometric: Optional[GeometricConfig] = attrs.field(default=None)
+    intensity: Optional['IntensityConfig'] = attrs.field(default=None)
+    geometric: Optional['GeometricConfig'] = attrs.field(default=None)

This change will resolve the undefined names issue and prevent potential circular imports.

🧰 Tools

🪛 Ruff

71-71: Undefined name IntensityConfig

(F821)

---

72-72: Undefined name GeometricConfig

(F821)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/model_config.py (1)

1-9: Consider adding type hints to class attributes.

While the attributes are well-described in the docstring, adding type hints would improve code clarity and enable better static type checking. Consider updating the class definition to include type information for each attribute. For example:

@attr.s(auto_attribs=True)
class ModelConfig:
    backbone: dict
    heads: dict
    base_checkpoint: str

Replace dict and str with more specific types if applicable. This change will enhance code readability and make it easier for developers to understand the expected types of these configuration parameters.

🧰 Tools

🪛 Ruff

1-1: Undefined name attr

(F821)

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/model_config.py (1)

23-23: Consider using consistent decorators across the file.

The UNetConfig class uses @attrs.define, while the main ModelConfig class uses @attr.s. For consistency, consider using the same decorator style throughout the file. If you're using a newer version of attrs, you might want to update all decorators to @attrs.define. Otherwise, change this to @attr.s(auto_attribs=True) to match the main class.

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/model_config.py (3)

20-46: LGTM! Consider refactoring error messages.

The methods in the ModelConfig class are well-implemented and provide good initialization and validation logic.

To reduce code duplication in error messages, consider creating a helper method for generating error messages:

def _get_weight_error_message(self, backbone_type, valid_weights):
    return f"Invalid pre-trained weights for {backbone_type}. Must be one of {valid_weights}"

# Then use it in validate_pre_trained_weights:
if self.backbone_type == BackboneType.CONVNEXT:
    if self.pre_trained_weights not in convnext_weights:
        raise ValueError(self._get_weight_error_message("ConvNext", convnext_weights))
elif self.backbone_type == BackboneType.SWINT:
    if self.pre_trained_weights not in swint_weights:
        raise ValueError(self._get_weight_error_message("SwinT", swint_weights))

This refactoring will make the code more maintainable and reduce the risk of inconsistencies in error messages.

🧰 Tools

🪛 Ruff

25-25: Undefined name BackboneType

(F821)

---

26-26: Undefined name UNetConfig

(F821)

---

27-27: Undefined name BackboneType

(F821)

---

28-28: Undefined name ConvNextConfig

(F821)

---

29-29: Undefined name BackboneType

(F821)

---

30-30: Undefined name SwinTConfig

(F821)

---

38-38: Undefined name BackboneType

(F821)

---

41-41: Undefined name BackboneType

(F821)

---

44-44: Undefined name BackboneType

(F821)

---

53-65: LGTM! Consider adding validation for None values.

The UNetConfig class is well-defined with appropriate attributes for UNet configuration.

Consider adding validation for attributes that have None as default value, such as max_stride and stem_stride. You could do this in a post-init method:

@attrs.define
class UNetConfig:
    # ... existing attributes ...

    def __attrs_post_init__(self):
        if self.max_stride is None:
            # Set a default value or raise an error if it's required
            self.max_stride = 16  # Example default value
        if self.stem_stride is None:
            # Set a default value or raise an error if it's required
            self.stem_stride = 1  # Example default value

This ensures that these critical parameters always have valid values.

---

66-89: LGTM! Consider reordering attributes for consistency.

The ConvNextConfig and SwinTConfig classes are well-defined with appropriate attributes for their respective architectures.

For better consistency between the two classes, consider reordering the attributes in SwinTConfig to match the order in ConvNextConfig as closely as possible. For example:

@attrs.define
class SwinTConfig:
    model_type: str = "tiny"
    arch: dict = attrs.field(factory=lambda: {'embed': 96, 'depths': [2, 2, 6, 2], 'channels': [3, 6, 12, 24]})
    in_channels: int = 1
    patch_size: list = attrs.field(factory=lambda: [4, 4])
    stem_patch_stride: int = 2
    window_size: list = attrs.field(factory=lambda: [7, 7])
    kernel_size: int = 3
    filters_rate: float = 1.5
    convs_per_block: int = 2
    up_interpolate: bool = True

This reordering makes it easier to compare the two configurations at a glance.

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/model_config.py (2)

55-67: Consider adding type hints to UNetConfig attributes.

To improve code clarity and maintainability, consider adding type hints to the attributes in the UNetConfig class. For example:

@attrs.define
class UNetConfig:
    in_channels: int = 1
    kernel_size: int = 3
    filters: int = 32
    filters_rate: float = 1.5
    max_stride: Optional[int] = None
    stem_stride: Optional[int] = None
    middle_block: bool = True
    up_interpolate: bool = True
    stacks: int = 3
    convs_per_block: int = 2

Don't forget to import Optional from typing if you use it for max_stride and stem_stride.

---

68-91: Enhance ConvNextConfig and SwinTConfig with type hints and consider using attrs.Factory.

1. Add type hints to improve code clarity:

@attrs.define
class ConvNextConfig:
    model_type: str = "tiny"
    arch: Dict[str, Union[List[int], List[int]]] = attrs.field(factory=lambda: {'depths': [3, 3, 9, 3], 'channels': [96, 192, 384, 768]})
    stem_patch_kernel: int = 4
    stem_patch_stride: int = 2
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 1.5
    convs_per_block: int = 2
    up_interpolate: bool = True

@attrs.define
class SwinTConfig:
    model_type: str = "tiny"
    arch: Dict[str, Union[int, List[int]]] = attrs.field(factory=lambda: {'embed': 96, 'depths': [2, 2, 6, 2], 'channels': [3, 6, 12, 24]})
    patch_size: List[int] = attrs.field(factory=lambda: [4, 4])
    stem_patch_stride: int = 2
    window_size: List[int] = attrs.field(factory=lambda: [7, 7])
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 1.5
    convs_per_block: int = 2
    up_interpolate: bool = True

2. Consider using attrs.Factory instead of lambda functions for better readability:

from attrs import Factory

# In ConvNextConfig
arch: Dict[str, Union[List[int], List[int]]] = attrs.field(factory=Factory(lambda: {'depths': [3, 3, 9, 3], 'channels': [96, 192, 384, 768]}))

# In SwinTConfig
arch: Dict[str, Union[int, List[int]]] = attrs.field(factory=Factory(lambda: {'embed': 96, 'depths': [2, 2, 6, 2], 'channels': [3, 6, 12, 24]}))
patch_size: List[int] = attrs.field(factory=Factory(lambda: [4, 4]))
window_size: List[int] = attrs.field(factory=Factory(lambda: [7, 7]))

These changes will improve type checking and make the code more explicit about the expected types for each attribute.

Configuration used: CodeRabbit UI
Review profile: CHILL

sleap_nn/config/model_config.py (3)

1-2: Add missing imports for type hints.

To improve type checking and code clarity, please add the following imports at the beginning of the file:

from typing import Dict, List, Optional, Union

These imports are necessary for the type hints used throughout the file.

---

56-92: Enhance documentation and type hinting for nested configuration classes.

For the UNetConfig, ConvNextConfig, and SwinTConfig classes:

1. Add docstrings to explain the purpose of each class and its attributes.
2. Consider adding type hints to all attributes for better code clarity and type checking.

For example, for the UNetConfig class:

@attrs.define
class UNetConfig:
    """Configuration for UNet backbone.

    Attributes:
        in_channels (int): Number of input channels.
        kernel_size (int): Size of the convolutional kernel.
        filters (int): Number of filters in the first layer.
        filters_rate (float): Rate at which the number of filters increases.
        max_stride (Optional[int]): Maximum stride in the network.
        stem_stride (Optional[int]): Stride in the stem of the network.
        middle_block (bool): Whether to include a middle block.
        up_interpolate (bool): Whether to use interpolation for upsampling.
        stacks (int): Number of encoder/decoder stacks.
        convs_per_block (int): Number of convolutions per block.
    """
    in_channels: int = 1
    kernel_size: int = 3
    filters: int = 32
    filters_rate: float = 1.5
    max_stride: Optional[int] = None
    stem_stride: Optional[int] = None
    middle_block: bool = True
    up_interpolate: bool = True
    stacks: int = 3
    convs_per_block: int = 2

Apply similar improvements to ConvNextConfig and SwinTConfig classes.

---

94-120: Improve consistency and documentation for additional configuration classes.

For the HeadConfig, SingleInstanceConfig, and ConfMapsConfig classes:

1. Add docstrings to HeadConfig and SingleInstanceConfig explaining their purpose and attributes, similar to ConfMapsConfig.
2. Consider using a consistent style for optional attributes. For example, in ConfMapsConfig, you could use attrs.field(default=None) instead of Optional[Type] = None for consistency with other classes.

Example for HeadConfig:

@attrs.define
class HeadConfig:
    """Configuration for model heads.

    Attributes:
        head_configs (Dict[str, Optional[Dict]]): A dictionary of head configurations.
            Keys represent head types, and values are their respective configurations.
    """
    head_configs: Dict[str, Optional[Dict]] = attrs.field(
        factory=lambda: {
            "single_instance": None,
            "centroid": None,
            "centered_instance": None,
            "bottomup": None
        }
    )

Apply similar improvements to SingleInstanceConfig.

🧰 Tools

🪛 Ruff

96-96: Undefined name Dict

(F821)

---

96-96: Undefined name Optional

(F821)

---

96-96: Undefined name Dict

(F821)

---

107-107: Undefined name Optional

(F821)

---

107-107: Undefined name ConfMapsConfig

(F821)

---

118-118: Undefined name Optional

(F821)

---

118-118: Undefined name List

(F821)

---

119-119: Undefined name Optional

(F821)

---

120-120: Undefined name Optional

(F821)

Configuration used: CodeRabbit UI
Review profile: CHILL

Configuration used: CodeRabbit UI
Review profile: CHILL

tests/config/test_training_job_config.py (1)

171-171: 🛠️ Refactor suggestion

Add tests for error cases.

The test suite should include tests for error cases such as invalid file paths and malformed YAML.

def test_from_yaml_errors():
    """Test error handling in from_yaml."""
    with pytest.raises(ValueError, match="Failed to parse YAML"):
        TrainingJobConfig.from_yaml("invalid: yaml: :")

def test_load_yaml_errors():
    """Test error handling in load_yaml."""
    with pytest.raises(FileNotFoundError):
        TrainingJobConfig.load_yaml("nonexistent.yaml")

    with tempfile.NamedTemporaryFile(mode='w') as f:
        f.write("invalid: yaml: :")
        f.flush()
        with pytest.raises(ValueError, match="Failed to parse YAML"):
            TrainingJobConfig.load_yaml(f.name)

tests/config/test_model_config.py (1)

58-69: ⚠️ Potential issue

Fix incomplete test implementation.

The test creates a structured configuration but doesn't actually test anything. The config variable is unused.

def test_update_config(default_config):
    """Test updating configuration attributes."""
-    config = OmegaConf.structured(
-        ModelConfig(
-            backbone_type="unet",
-            init_weight="default",
-            pre_trained_weights=None,
-            backbone_config=BackboneConfig(),
-            head_configs=HeadConfig(),
-        )
-    )
+    # Test updating backbone_type
+    default_config.backbone_type = "convnext"
+    assert default_config.backbone_type == "convnext"
+
+    # Test updating pre_trained_weights
+    default_config.pre_trained_weights = "ConvNeXt_Tiny_Weights"
+    assert default_config.pre_trained_weights == "ConvNeXt_Tiny_Weights"
+
+    # Test invalid update
+    with pytest.raises(ValueError, match="Invalid backbone_type"):
+        default_config.backbone_type = "invalid"

🧰 Tools

🪛 Ruff (0.8.2)

60-60: Local variable config is assigned to but never used

Remove assignment to unused variable config

(F841)

sleap_nn/config/training_job_config.py (2)

61-73: 🛠️ Refactor suggestion

Add error handling for YAML parsing and schema validation.

The method should handle potential errors during YAML parsing and schema validation.

 @classmethod
 def from_yaml(cls, yaml_data: Text) -> "TrainingJobConfig":
-    schema = OmegaConf.structured(cls)
-    config = OmegaConf.create(yaml_data)
-    OmegaConf.merge(schema, config)
-    return cls(**OmegaConf.to_container(config, resolve=True))
+    try:
+        schema = OmegaConf.structured(cls)
+        config = OmegaConf.create(yaml_data)
+        OmegaConf.merge(schema, config)
+        return cls(**OmegaConf.to_container(config, resolve=True))
+    except Exception as e:
+        raise ValueError(f"Failed to parse YAML data: {e}")

---

123-125: ⚠️ Potential issue

Fix parameter mismatch in load_config.

The load_yaml method doesn't accept the load_training_config parameter, which will cause a TypeError.

-    return TrainingJobConfig.load_yaml(
-        filename, load_training_config=load_training_config
-    )
+    if os.path.isdir(filename):
+        config_file = os.path.join(
+            filename,
+            "training_job.yaml" if load_training_config else "initial_config.yaml"
+        )
+        if not os.path.exists(config_file):
+            raise FileNotFoundError(
+                f"Configuration file not found in directory: {config_file}"
+            )
+        filename = config_file
+    return TrainingJobConfig.load_yaml(filename)

tests/config/test_model_config.py (2)

52-55: 🛠️ Refactor suggestion

Fix docstring and add specific error message assertion.

1. The docstring is incorrectly copied from test_invalid_pre_trained_weights.
2. Add specific error message assertion.

def test_invalid_backbonetype():
-    """Test validation failure with an invalid pre_trained_weights."""
+    """Test validation failure with an invalid backbone_type."""
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match='backbone_type must be one of "unet", "convnext", "swint"'):
        ModelConfig(backbone_type="net")

---

40-43: 🛠️ Refactor suggestion

Enhance test coverage and fix None comparison.

1. Use is operator for None comparison.
2. Test all attributes of the default configuration.

def test_default_initialization(default_config):
    """Test default initialization of ModelConfig."""
    assert default_config.init_weight == "default"
-    assert default_config.pre_trained_weights == None
+    assert default_config.pre_trained_weights is None
+    assert default_config.backbone_type == "unet"
+    assert isinstance(default_config.backbone_config, BackboneConfig)
+    assert isinstance(default_config.head_configs, HeadConfig)

🧰 Tools

🪛 Ruff (0.8.2)

43-43: Comparison to None should be cond is None

Replace with cond is None

(E711)

sleap_nn/config/training_job_config.py (1)

27-35: Clean up unused imports.

Several imports are unused and can be removed to improve code clarity:

-import os
 from attrs import define, field, asdict
 import sleap_nn
 from sleap_nn.config.data_config import DataConfig
 from sleap_nn.config.model_config import ModelConfig
 from sleap_nn.config.trainer_config import TrainerConfig
-import yaml
-from typing import Text, Dict, Any, Optional
+from typing import Text, Optional
 from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

27-27: os imported but unused

Remove unused import: os

(F401)

---

28-28: attrs.field imported but unused

Remove unused import: attrs.field

(F401)

---

33-33: yaml imported but unused

Remove unused import: yaml

(F401)

---

34-34: typing.Dict imported but unused

Remove unused import

(F401)

---

34-34: typing.Any imported but unused

Remove unused import

(F401)

tests/config/test_model_config.py (3)

9-25: Clean up unused imports.

Several imported configuration classes are not used in the test file. Remove the following:

• UNetConfig, ConvNextConfig, SwinTConfig
• SingleInstanceConfig, CentroidConfig, CenteredInstanceConfig, BottomUpConfig
• SingleInstanceConfMapsConfig, CentroidConfMapsConfig, CenteredInstanceConfMapsConfig, BottomUpConfMapsConfig, PAFConfig

from sleap_nn.config.model_config import (
    ModelConfig,
    BackboneConfig,
-    UNetConfig,
-    ConvNextConfig,
-    SwinTConfig,
    HeadConfig,
-    SingleInstanceConfig,
-    CentroidConfig,
-    CenteredInstanceConfig,
-    BottomUpConfig,
-    SingleInstanceConfMapsConfig,
-    CentroidConfMapsConfig,
-    CenteredInstanceConfMapsConfig,
-    BottomUpConfMapsConfig,
-    PAFConfig,
)

🧰 Tools

🪛 Ruff (0.8.2)

12-12: sleap_nn.config.model_config.UNetConfig imported but unused

Remove unused import

(F401)

---

13-13: sleap_nn.config.model_config.ConvNextConfig imported but unused

Remove unused import

(F401)

---

14-14: sleap_nn.config.model_config.SwinTConfig imported but unused

Remove unused import

(F401)

---

16-16: sleap_nn.config.model_config.SingleInstanceConfig imported but unused

Remove unused import

(F401)

---

17-17: sleap_nn.config.model_config.CentroidConfig imported but unused

Remove unused import

(F401)

---

18-18: sleap_nn.config.model_config.CenteredInstanceConfig imported but unused

Remove unused import

(F401)

---

19-19: sleap_nn.config.model_config.BottomUpConfig imported but unused

Remove unused import

(F401)

---

20-20: sleap_nn.config.model_config.SingleInstanceConfMapsConfig imported but unused

Remove unused import

(F401)

---

21-21: sleap_nn.config.model_config.CentroidConfMapsConfig imported but unused

Remove unused import

(F401)

---

22-22: sleap_nn.config.model_config.CenteredInstanceConfMapsConfig imported but unused

Remove unused import

(F401)

---

23-23: sleap_nn.config.model_config.BottomUpConfMapsConfig imported but unused

Remove unused import

(F401)

---

24-24: sleap_nn.config.model_config.PAFConfig imported but unused

Remove unused import

(F401)

---

28-37: Enhance fixture documentation.

The fixture's docstring could be more descriptive about the default values.

@pytest.fixture
def default_config():
-    """Fixture for a default ModelConfig instance."""
+    """Fixture for a default ModelConfig instance.
+    
+    Returns:
+        ModelConfig: A configuration instance with:
+            - init_weight: "default"
+            - pre_trained_weights: None
+            - backbone_type: "unet"
+            - backbone_config: Default BackboneConfig
+            - head_configs: Default HeadConfig
+    """

---

46-49: Add specific error message assertion.

The test should verify the specific error message for better debugging.

def test_invalid_pre_trained_weights():
    """Test validation failure with an invalid pre_trained_weights."""
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="Invalid pre-trained weights for ConvNext"):
        ModelConfig(pre_trained_weights="here", backbone_type="unet")

sleap_nn/config/model_config.py (2)

269-272: Use Enum for backbone_type validation.

Consider using an Enum for backbone_type to improve type safety and maintainability.

+class BackboneType(Enum):
+    UNET = "unet"
+    CONVNEXT = "convnext"
+    SWINT = "swint"
+
-    backbone_type: str = field(
+    backbone_type: BackboneType = field(
-        default="unet",
+        default=BackboneType.UNET,
         validator=lambda instance, attr, value: instance.validate_backbone_type(value),
     )

---

283-334: Simplify validation methods.

1. Use the proposed BackboneType enum.
2. Move weight lists to class constants.
3. Simplify validation using a mapping.

+    CONVNEXT_WEIGHTS = [
+        "ConvNeXt_Base_Weights",
+        "ConvNeXt_Tiny_Weights",
+        "ConvNeXt_Small_Weights",
+        "ConvNeXt_Large_Weights",
+    ]
+    SWINT_WEIGHTS = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]
+    VALID_WEIGHTS = {
+        BackboneType.CONVNEXT: CONVNEXT_WEIGHTS,
+        BackboneType.SWINT: SWINT_WEIGHTS,
+        BackboneType.UNET: None,
+    }
+
     def validate_pre_trained_weights(self, value):
         if value is None:
             return
 
-        convnext_weights = [...]  # Remove duplicated lists
-        swint_weights = [...]
-
-        if self.backbone_type == "convnext":
-            if value not in convnext_weights:
-                raise ValueError(...)
-        elif self.backbone_type == "swint":
-            if value not in swint_weights:
-                raise ValueError(...)
-        elif self.backbone_type == "unet":
+        allowed_weights = self.VALID_WEIGHTS.get(self.backbone_type)
+        if allowed_weights is None and value is not None:
             raise ValueError("UNet does not support pre-trained weights.")
+        elif allowed_weights and value not in allowed_weights:
+            raise ValueError(
+                f"Invalid pre-trained weights for {self.backbone_type.value}. "
+                f"Must be one of {allowed_weights}"
+            )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

sleap_nn/config/model_config.py (2)

326-331: 🛠️ Refactor suggestion

Fix inconsistent docstring.

The docstring mentions configurations that don't match the actual attributes:

• multi_instance, multi_class_bottomup, and multi_class_topdown are mentioned but not defined
• bottomup is defined but not mentioned

     Attributes:
         single_instance: An instance of `SingleInstanceConfmapsHeadConfig`.
         centroid: An instance of `CentroidsHeadConfig`.
         centered_instance: An instance of `CenteredInstanceConfmapsHeadConfig`.
-        multi_instance: An instance of `MultiInstanceConfig`.
-        multi_class_bottomup: An instance of `MultiClassBottomUpConfig`.
-        multi_class_topdown: An instance of `MultiClassTopDownConfig`.
+        bottomup: An instance of `BottomUpConfig`.

---

43-43: 🛠️ Refactor suggestion

Set default value for max_stride.

According to past review comments, max_stride shouldn't be optional and should have a default value of 16.

-    max_stride: Optional[int] = None
+    max_stride: int = 16

sleap_nn/config/training_job_config.py (1)

114-127: ⚠️ Potential issue

Fix load_config function implementation.

The function has incorrect parameter passing and missing file validation.

 def load_config(filename: Text, load_training_config: bool = True) -> TrainingJobConfig:
+    if not os.path.exists(filename):
+        raise FileNotFoundError(f"Configuration file not found: {filename}")
+
+    if os.path.isdir(filename):
+        config_file = os.path.join(
+            filename,
+            "training_job.yaml" if load_training_config else "initial_config.yaml"
+        )
+        if not os.path.exists(config_file):
+            raise FileNotFoundError(
+                f"Configuration file not found in directory: {config_file}"
+            )
+        filename = config_file
+
     return TrainingJobConfig.load_yaml(
-        filename, load_training_config=load_training_config
+        filename
     )

tests/config/test_training_job_config.py (1)

84-90: ⚠️ Potential issue

Fix assertions in test_from_yaml.

The assertions incorrectly treat config attributes as dictionaries instead of proper config objects.

-    assert config.data["train_labels_path"] == sample_config["data"].train_labels_path
-    assert config.data["val_labels_path"] == sample_config["data"].val_labels_path
-    assert config.model["backbone_type"] == sample_config["model"].backbone_type
-    assert (
-        config.trainer["early_stopping"]["patience"]
-        == sample_config["trainer"].early_stopping.patience
-    )
+    assert isinstance(config.data, DataConfig)
+    assert config.data.train_labels_path == sample_config["data"].train_labels_path
+    assert config.data.val_labels_path == sample_config["data"].val_labels_path
+    assert isinstance(config.model, ModelConfig)
+    assert config.model.backbone_type == sample_config["model"].backbone_type
+    assert isinstance(config.trainer, TrainerConfig)
+    assert config.trainer.early_stopping.patience == sample_config["trainer"].early_stopping.patience

sleap_nn/config/model_config.py (6)

7-10: Remove unused import.

The Enum import is not used directly in the code.

-from enum import Enum

🧰 Tools

🪛 Ruff (0.8.2)

8-8: enum.Enum imported but unused

Remove unused import: enum.Enum

(F401)

---

77-77: Add validation for model_type.

Add validation to ensure model_type is one of the allowed values.

-    model_type: str = "tiny"  # Options: tiny, small, base, large
+    model_type: str = field(
+        default="tiny",
+        validator=attrs.validators.in_(["tiny", "small", "base", "large"])
+    )

---

114-114: Add validation for model_type.

Add validation to ensure model_type is one of the allowed values.

-    model_type: str = "tiny"  # Options: tiny, small, base
+    model_type: str = field(
+        default="tiny",
+        validator=attrs.validators.in_(["tiny", "small", "base"])
+    )

---

122-124: Improve type hints for list attributes.

Add type hints for list attributes to specify their element types.

-    patch_size: list = field(factory=lambda: [4, 4])
-    window_size: list = field(factory=lambda: [7, 7])
+    patch_size: List[int] = field(factory=lambda: [4, 4])
+    window_size: List[int] = field(factory=lambda: [7, 7])

---

376-379: Use enum for backbone_type validation.

Consider using an enum for backbone_type to make the code more maintainable and type-safe.

+class BackboneType(str, Enum):
+    UNET = "unet"
+    CONVNEXT = "convnext"
+    SWINT = "swint"

-    backbone_type: str = field(
+    backbone_type: BackboneType = field(
-        default="unet",
+        default=BackboneType.UNET,
         validator=lambda instance, attr, value: instance.validate_backbone_type(value),
     )

     def validate_backbone_type(self, value):
-        valid_types = ["unet", "convnext", "swint"]
-        if value not in valid_types:
-            raise ValueError(f"Invalid backbone_type. Must be one of {valid_types}")
+        if not isinstance(value, BackboneType):
+            raise ValueError(f"Invalid backbone_type. Must be a BackboneType enum.")

Also applies to: 390-398

---

399-441: Simplify pre_trained_weights validation using a dictionary.

The validation logic can be simplified and made more maintainable using a dictionary mapping.

     def validate_pre_trained_weights(self, value):
         if value is None:
             return

-        convnext_weights = [
-            "ConvNeXt_Base_Weights",
-            "ConvNeXt_Tiny_Weights",
-            "ConvNeXt_Small_Weights",
-            "ConvNeXt_Large_Weights",
-        ]
-        swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]
-
-        if self.backbone_type == "convnext":
-            if value not in convnext_weights:
-                raise ValueError(
-                    f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
-                )
-        elif self.backbone_type == "swint":
-            if value not in swint_weights:
-                raise ValueError(
-                    f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
-                )
-        elif self.backbone_type == "unet":
-            raise ValueError("UNet does not support pre-trained weights.")
+        VALID_WEIGHTS = {
+            BackboneType.CONVNEXT: [
+                "ConvNeXt_Base_Weights",
+                "ConvNeXt_Tiny_Weights",
+                "ConvNeXt_Small_Weights",
+                "ConvNeXt_Large_Weights",
+            ],
+            BackboneType.SWINT: [
+                "Swin_T_Weights",
+                "Swin_S_Weights",
+                "Swin_B_Weights"
+            ],
+            BackboneType.UNET: None
+        }
+        
+        allowed_weights = VALID_WEIGHTS.get(self.backbone_type)
+        if allowed_weights is None:
+            raise ValueError(f"{self.backbone_type.value} does not support pre-trained weights.")
+        elif value not in allowed_weights:
+            raise ValueError(
+                f"Invalid pre-trained weights for {self.backbone_type.value}. "
+                f"Must be one of {allowed_weights}"
+            )

sleap_nn/config/training_job_config.py (1)

27-35: Clean up imports.

Several imports are unused and can be removed:

• os (used in load_config but not implemented correctly)
• field from attrs (not used)
• yaml (replaced by OmegaConf)
• Dict, Any from typing (not used)

-import os
-from attrs import define, field, asdict
+from attrs import define, asdict
 import sleap_nn
 from sleap_nn.config.data_config import DataConfig
 from sleap_nn.config.model_config import ModelConfig
 from sleap_nn.config.trainer_config import TrainerConfig
-import yaml
-from typing import Text, Dict, Any, Optional
+from typing import Text, Optional
 from omegaconf import OmegaConf

🧰 Tools

🪛 Ruff (0.8.2)

27-27: os imported but unused

Remove unused import: os

(F401)

---

28-28: attrs.field imported but unused

Remove unused import: attrs.field

(F401)

---

33-33: yaml imported but unused

Remove unused import: yaml

(F401)

---

34-34: typing.Dict imported but unused

Remove unused import

(F401)

---

34-34: typing.Any imported but unused

Remove unused import

(F401)

tests/config/test_training_job_config.py (2)

35-35: Remove unused import.

The asdict import from dataclasses is not used in the test file.

-from dataclasses import asdict

🧰 Tools

🪛 Ruff (0.8.2)

35-35: dataclasses.asdict imported but unused

Remove unused import: dataclasses.asdict

(F401)

---

194-194: Remove unused variable assignment.

The config variable is assigned but never used.

-        config = TrainingJobConfig.from_yaml(yaml_data)
+        TrainingJobConfig.from_yaml(yaml_data)

🧰 Tools

🪛 Ruff (0.8.2)

194-194: Local variable config is assigned to but never used

Remove assignment to unused variable config

(F841)