diff --git a/visdet/engine/config/schema_generator.py b/visdet/engine/config/schema_generator.py index 896edeb1..a0a3e87b 100644 --- a/visdet/engine/config/schema_generator.py +++ b/visdet/engine/config/schema_generator.py @@ -17,20 +17,32 @@ def __init__(self) -> None: """Initialize the schema registry.""" self._schemas: Dict[Type, Type[BaseModel]] = {} self._manual_overrides: Dict[Type, Type[BaseModel]] = {} - - def register_schema(self, component_cls: Type, schema_cls: Type[BaseModel], is_manual: bool = False) -> None: + self._type_name_to_schema: Dict[str, Type[BaseModel]] = {} + + def register_schema( + self, + component_cls: Type, + schema_cls: Type[BaseModel], + is_manual: bool = False, + type_name: Optional[str] = None, + ) -> None: """Register a Pydantic schema for a component class. Args: component_cls: The component class to register a schema for schema_cls: The Pydantic model schema is_manual: Whether this is a manually created schema (overrides auto-generated) + type_name: Optional type name for reverse lookup (defaults to class name) """ if is_manual: self._manual_overrides[component_cls] = schema_cls else: self._schemas[component_cls] = schema_cls + # Register reverse lookup by type name + name = type_name or component_cls.__name__ + self._type_name_to_schema[name] = schema_cls + def get_schema(self, component_cls: Type) -> Optional[Type[BaseModel]]: """Get the Pydantic schema for a component class. @@ -45,6 +57,23 @@ def get_schema(self, component_cls: Type) -> Optional[Type[BaseModel]]: return self._manual_overrides[component_cls] return self._schemas.get(component_cls) + def get_schema_by_type_name(self, type_name: str) -> Optional[Type[BaseModel]]: + """Get the Pydantic schema by type name (for config 'type' field). + + This enables validation of config dicts without knowing the component class. + + Args: + type_name: The 'type' field value from a config dict + + Returns: + The Pydantic model schema, or None if not found + + Example: + >>> schema = registry.get_schema_by_type_name('SwinTransformer') + >>> validated = schema(**config_dict) + """ + return self._type_name_to_schema.get(type_name) + def has_schema(self, component_cls: Type) -> bool: """Check if a schema exists for a component class. @@ -56,6 +85,14 @@ def has_schema(self, component_cls: Type) -> bool: """ return component_cls in self._manual_overrides or component_cls in self._schemas + def list_type_names(self) -> list[str]: + """List all registered type names. + + Returns: + List of type names with registered schemas. + """ + return list(self._type_name_to_schema.keys()) + # Global schema registry _schema_registry = SchemaRegistry() @@ -202,7 +239,8 @@ def validate_config_with_schema( Args: config: Configuration dictionary to validate component_cls: The component class (if known) - type_name: The 'type' value from config (for registry lookup) + type_name: The 'type' value from config (for registry lookup). + If not provided, will try to extract from config['type']. Returns: Validated Pydantic model instance @@ -214,6 +252,8 @@ def validate_config_with_schema( Example: >>> config = {'type': 'SwinTransformer', 'embed_dims': 96, 'depths': [2,2,6,2]} >>> validated = validate_config_with_schema(config, type_name='SwinTransformer') + >>> # Or auto-detect from config: + >>> validated = validate_config_with_schema(config) """ # Try to get schema by component class if component_cls: @@ -222,15 +262,18 @@ def validate_config_with_schema( return schema(**config) # Try to get schema by type name + if type_name is None: + # Try to extract from config + type_name = config.get("type") + if type_name: - # This would require a reverse lookup in registries - # For now, raise an error - this can be implemented later - raise NotImplementedError( - "Schema lookup by type name not yet implemented. Please provide component_cls instead." - ) + schema = _schema_registry.get_schema_by_type_name(type_name) + if schema: + return schema(**config) raise ValueError( - "Cannot validate config: no component_cls or type_name provided, or no schema registered for the component" + f"Cannot validate config: no schema registered for type '{type_name}'. " + f"Available types: {_schema_registry.list_type_names()}" ) diff --git a/visdet/py_configs/__init__.py b/visdet/py_configs/__init__.py new file mode 100644 index 00000000..6003855a --- /dev/null +++ b/visdet/py_configs/__init__.py @@ -0,0 +1,93 @@ +"""Python-first configuration API for visdet. + +This package provides a programmatic way to create configurations +with full IDE autocomplete and type safety. + +Key Features: +1. Builder functions for creating component configs +2. Pre-defined experiment presets +3. Sweep generators for hyperparameter search +4. Full IDE autocomplete via Pydantic models + +Example: + >>> from visdet.py_configs import mask_rcnn_swin_tiny_coco + >>> # Get a complete experiment config with one line + >>> cfg = mask_rcnn_swin_tiny_coco(data_root='/data/coco') + >>> + >>> # Customize with full autocomplete + >>> cfg.train_cfg.max_epochs = 24 + >>> cfg.optim_wrapper.optimizer.lr = 2e-4 + >>> + >>> # Use with SimpleRunner + >>> from visdet import SimpleRunner + >>> runner = SimpleRunner(config=cfg) + >>> runner.train() +""" + +# Builders - factory functions for components +from visdet.py_configs.builders import ( + # Backbones + swin_tiny, + swin_small, + swin_base, + resnet50, + resnet101, + # Necks + fpn_for_swin, + fpn_for_resnet, + # Heads + standard_rpn_head, + standard_roi_head, + # Models + mask_rcnn, + # Optimizers + adamw_default, + one_cycle_scheduler, + # Data + coco_train_pipeline, + coco_test_pipeline, + coco_dataset, + train_dataloader, + val_dataloader, +) + +# Presets - ready-to-use experiment configs +from visdet.py_configs.presets import ( + mask_rcnn_swin_tiny_coco, + mask_rcnn_swin_small_coco, + # Sweeps + lr_sweep, + batch_size_sweep, +) + +__all__ = [ + # Backbone builders + "swin_tiny", + "swin_small", + "swin_base", + "resnet50", + "resnet101", + # Neck builders + "fpn_for_swin", + "fpn_for_resnet", + # Head builders + "standard_rpn_head", + "standard_roi_head", + # Model builders + "mask_rcnn", + # Optimizer builders + "adamw_default", + "one_cycle_scheduler", + # Data builders + "coco_train_pipeline", + "coco_test_pipeline", + "coco_dataset", + "train_dataloader", + "val_dataloader", + # Presets + "mask_rcnn_swin_tiny_coco", + "mask_rcnn_swin_small_coco", + # Sweeps + "lr_sweep", + "batch_size_sweep", +] diff --git a/visdet/py_configs/builders.py b/visdet/py_configs/builders.py new file mode 100644 index 00000000..7a846af2 --- /dev/null +++ b/visdet/py_configs/builders.py @@ -0,0 +1,489 @@ +"""Factory functions for building configuration objects. + +This module provides builder functions that create properly configured +detection models with sensible defaults and full IDE autocomplete support. + +Example: + >>> from visdet.py_configs.builders import mask_rcnn, swin_tiny + >>> # Create Mask R-CNN with Swin-Tiny backbone + >>> model = mask_rcnn(backbone=swin_tiny(), num_classes=80) + >>> model.to_dict() # Convert to dict for MMEngine +""" + +from visdet.schemas import ( + AdamWConfig, + CocoDatasetConfig, + DataLoaderConfig, + DefaultSamplerConfig, + DetDataPreprocessorConfig, + EpochBasedTrainLoopConfig, + ExperimentConfig, + FCNMaskHeadConfig, + FPNConfig, + LoadAnnotationsConfig, + LoadImageFromFileConfig, + MaskRCNNConfig, + OneCycleLRConfig, + OptimWrapperConfig, + PackDetInputsConfig, + RandomFlipConfig, + RCNNTestConfig, + RCNNTrainConfig, + ResizeConfig, + ResNetConfig, + RPNHeadConfig, + RPNProposalConfig, + RPNTestConfig, + RPNTrainConfig, + Shared2FCBBoxHeadConfig, + SingleRoIExtractorConfig, + StandardRoIHeadConfig, + SwinTransformerConfig, + TwoStageTestConfig, + TwoStageTrainConfig, + ValLoopConfig, +) + + +# ============================================================================= +# Backbone Builders +# ============================================================================= + + +def swin_tiny( + drop_path_rate: float = 0.2, + frozen_stages: int = -1, + **kwargs, +) -> SwinTransformerConfig: + """Create Swin Transformer Tiny backbone config. + + Args: + drop_path_rate: Stochastic depth rate. + frozen_stages: Freeze stages up to this index (-1 = none). + **kwargs: Additional overrides. + + Returns: + Configured SwinTransformerConfig. + + Example: + >>> backbone = swin_tiny(drop_path_rate=0.3) + """ + return SwinTransformerConfig( + embed_dims=96, + depths=(2, 2, 6, 2), + num_heads=(3, 6, 12, 24), + window_size=7, + drop_path_rate=drop_path_rate, + frozen_stages=frozen_stages, + convert_weights=True, + **kwargs, + ) + + +def swin_small( + drop_path_rate: float = 0.3, + frozen_stages: int = -1, + **kwargs, +) -> SwinTransformerConfig: + """Create Swin Transformer Small backbone config. + + Deeper than Tiny (18 blocks vs 6 in stage 3). + + Args: + drop_path_rate: Stochastic depth rate. + frozen_stages: Freeze stages up to this index. + **kwargs: Additional overrides. + + Returns: + Configured SwinTransformerConfig. + """ + return SwinTransformerConfig( + embed_dims=96, + depths=(2, 2, 18, 2), + num_heads=(3, 6, 12, 24), + window_size=7, + drop_path_rate=drop_path_rate, + frozen_stages=frozen_stages, + convert_weights=True, + **kwargs, + ) + + +def swin_base( + drop_path_rate: float = 0.5, + frozen_stages: int = -1, + **kwargs, +) -> SwinTransformerConfig: + """Create Swin Transformer Base backbone config. + + Wider channels (128 vs 96) than Small. + + Args: + drop_path_rate: Stochastic depth rate. + frozen_stages: Freeze stages up to this index. + **kwargs: Additional overrides. + + Returns: + Configured SwinTransformerConfig. + """ + return SwinTransformerConfig( + embed_dims=128, + depths=(2, 2, 18, 2), + num_heads=(4, 8, 16, 32), + window_size=7, + drop_path_rate=drop_path_rate, + frozen_stages=frozen_stages, + convert_weights=True, + **kwargs, + ) + + +def resnet50( + frozen_stages: int = 1, + norm_eval: bool = True, + **kwargs, +) -> ResNetConfig: + """Create ResNet-50 backbone config. + + Standard ResNet-50 with frozen stem (stage 0). + + Args: + frozen_stages: Freeze stages up to this index. + norm_eval: Freeze BatchNorm stats during training. + **kwargs: Additional overrides. + + Returns: + Configured ResNetConfig. + """ + return ResNetConfig( + depth=50, + frozen_stages=frozen_stages, + norm_eval=norm_eval, + **kwargs, + ) + + +def resnet101( + frozen_stages: int = 1, + norm_eval: bool = True, + **kwargs, +) -> ResNetConfig: + """Create ResNet-101 backbone config.""" + return ResNetConfig( + depth=101, + frozen_stages=frozen_stages, + norm_eval=norm_eval, + **kwargs, + ) + + +# ============================================================================= +# Neck Builders +# ============================================================================= + + +def fpn_for_swin(embed_dims: int = 96, out_channels: int = 256, **kwargs) -> FPNConfig: + """Create FPN neck configured for Swin Transformer backbone. + + Auto-computes in_channels based on Swin embed_dims. + + Args: + embed_dims: Swin backbone embed_dims (96=Tiny/Small, 128=Base). + out_channels: FPN output channels. + **kwargs: Additional overrides. + + Returns: + Configured FPNConfig. + """ + return FPNConfig( + in_channels=[embed_dims, embed_dims * 2, embed_dims * 4, embed_dims * 8], + out_channels=out_channels, + num_outs=5, + **kwargs, + ) + + +def fpn_for_resnet(depth: int = 50, out_channels: int = 256, **kwargs) -> FPNConfig: + """Create FPN neck configured for ResNet backbone. + + Auto-computes in_channels based on ResNet depth. + + Args: + depth: ResNet depth (18, 34, 50, 101, 152). + out_channels: FPN output channels. + **kwargs: Additional overrides. + + Returns: + Configured FPNConfig. + """ + if depth in [18, 34]: + in_channels = [64, 128, 256, 512] + else: # 50, 101, 152 + in_channels = [256, 512, 1024, 2048] + + return FPNConfig( + in_channels=in_channels, + out_channels=out_channels, + num_outs=5, + **kwargs, + ) + + +# ============================================================================= +# Head Builders +# ============================================================================= + + +def standard_rpn_head(in_channels: int = 256, **kwargs) -> RPNHeadConfig: + """Create standard RPN head config. + + Args: + in_channels: Input feature channels from neck. + **kwargs: Additional overrides. + + Returns: + Configured RPNHeadConfig. + """ + return RPNHeadConfig( + in_channels=in_channels, + feat_channels=256, + **kwargs, + ) + + +def standard_roi_head(num_classes: int, with_mask: bool = True, **kwargs) -> StandardRoIHeadConfig: + """Create standard RoI head config for Mask R-CNN. + + Args: + num_classes: Number of object classes. + with_mask: Include mask head for instance segmentation. + **kwargs: Additional overrides. + + Returns: + Configured StandardRoIHeadConfig. + """ + mask_roi_extractor = None + mask_head = None + + if with_mask: + mask_roi_extractor = SingleRoIExtractorConfig( + roi_layer={"type": "RoIAlign", "output_size": 14, "sampling_ratio": 0}, + out_channels=256, + featmap_strides=[4, 8, 16, 32], + ) + mask_head = FCNMaskHeadConfig( + num_convs=4, + in_channels=256, + conv_out_channels=256, + num_classes=num_classes, + ) + + return StandardRoIHeadConfig( + bbox_roi_extractor=SingleRoIExtractorConfig( + roi_layer={"type": "RoIAlign", "output_size": 7, "sampling_ratio": 0}, + out_channels=256, + featmap_strides=[4, 8, 16, 32], + ), + bbox_head=Shared2FCBBoxHeadConfig( + in_channels=256, + fc_out_channels=1024, + roi_feat_size=7, + num_classes=num_classes, + ), + mask_roi_extractor=mask_roi_extractor, + mask_head=mask_head, + **kwargs, + ) + + +# ============================================================================= +# Model Builders +# ============================================================================= + + +def mask_rcnn( + backbone: SwinTransformerConfig | ResNetConfig, + num_classes: int = 80, + neck_out_channels: int = 256, + **kwargs, +) -> MaskRCNNConfig: + """Create Mask R-CNN config with auto-configured components. + + This builder automatically configures the FPN neck based on the + backbone type, and sets up RPN and RoI heads with standard settings. + + Args: + backbone: Backbone configuration (Swin or ResNet). + num_classes: Number of object classes (default: 80 for COCO). + neck_out_channels: FPN output channels. + **kwargs: Additional overrides. + + Returns: + Complete MaskRCNNConfig. + + Example: + >>> model = mask_rcnn(backbone=swin_tiny(), num_classes=20) + """ + # Auto-configure neck based on backbone + if isinstance(backbone, SwinTransformerConfig): + neck = fpn_for_swin(backbone.embed_dims, neck_out_channels) + else: + neck = fpn_for_resnet(backbone.depth, neck_out_channels) + + return MaskRCNNConfig( + backbone=backbone, + neck=neck, + rpn_head=standard_rpn_head(neck_out_channels), + roi_head=standard_roi_head(num_classes, with_mask=True), + **kwargs, + ) + + +# ============================================================================= +# Optimizer and Scheduler Builders +# ============================================================================= + + +def adamw_default(lr: float = 1e-4, weight_decay: float = 0.05, **kwargs) -> OptimWrapperConfig: + """Create AdamW optimizer wrapper with default settings. + + Args: + lr: Learning rate. + weight_decay: Weight decay coefficient. + **kwargs: Additional optimizer overrides. + + Returns: + Configured OptimWrapperConfig. + """ + return OptimWrapperConfig( + optimizer=AdamWConfig(lr=lr, weight_decay=weight_decay, **kwargs), + ) + + +def one_cycle_scheduler(max_lr: float = 1e-3, **kwargs) -> OneCycleLRConfig: + """Create 1cycle LR scheduler config. + + Args: + max_lr: Maximum learning rate at peak. + **kwargs: Additional scheduler overrides. + + Returns: + Configured OneCycleLRConfig. + """ + return OneCycleLRConfig(max_lr=max_lr, **kwargs) + + +# ============================================================================= +# Data Builders +# ============================================================================= + + +def coco_train_pipeline() -> list[dict]: + """Create standard COCO training data pipeline.""" + return [ + LoadImageFromFileConfig().to_dict(), + LoadAnnotationsConfig(with_bbox=True, with_mask=True).to_dict(), + ResizeConfig(scale=(1333, 800), keep_ratio=True).to_dict(), + RandomFlipConfig(prob=0.5).to_dict(), + PackDetInputsConfig().to_dict(), + ] + + +def coco_test_pipeline() -> list[dict]: + """Create standard COCO test/val data pipeline.""" + return [ + LoadImageFromFileConfig().to_dict(), + ResizeConfig(scale=(1333, 800), keep_ratio=True).to_dict(), + LoadAnnotationsConfig(with_bbox=True, with_mask=True).to_dict(), + PackDetInputsConfig( + meta_keys=("img_id", "img_path", "ori_shape", "img_shape", "scale_factor") + ).to_dict(), + ] + + +def coco_dataset( + data_root: str, + ann_file: str, + img_prefix: str = "", + pipeline: list[dict] | None = None, + test_mode: bool = False, + **kwargs, +) -> CocoDatasetConfig: + """Create COCO dataset config. + + Args: + data_root: Root data directory. + ann_file: Annotation file path (relative to data_root). + img_prefix: Image directory prefix. + pipeline: Data pipeline (auto-set if None). + test_mode: Test mode (no GT loading). + **kwargs: Additional dataset overrides. + + Returns: + Configured CocoDatasetConfig. + """ + if pipeline is None: + pipeline = coco_test_pipeline() if test_mode else coco_train_pipeline() + + return CocoDatasetConfig( + data_root=data_root, + ann_file=ann_file, + data_prefix={"img": img_prefix}, + pipeline=pipeline, + test_mode=test_mode, + **kwargs, + ) + + +def train_dataloader( + dataset: CocoDatasetConfig, + batch_size: int = 2, + num_workers: int = 2, + **kwargs, +) -> DataLoaderConfig: + """Create training dataloader config. + + Args: + dataset: Dataset configuration. + batch_size: Samples per batch per GPU. + num_workers: Data loading workers. + **kwargs: Additional dataloader overrides. + + Returns: + Configured DataLoaderConfig. + """ + return DataLoaderConfig( + batch_size=batch_size, + num_workers=num_workers, + persistent_workers=True, + sampler=DefaultSamplerConfig(shuffle=True), + dataset=dataset, + **kwargs, + ) + + +def val_dataloader( + dataset: CocoDatasetConfig, + batch_size: int = 1, + num_workers: int = 2, + **kwargs, +) -> DataLoaderConfig: + """Create validation dataloader config. + + Args: + dataset: Dataset configuration. + batch_size: Samples per batch per GPU. + num_workers: Data loading workers. + **kwargs: Additional dataloader overrides. + + Returns: + Configured DataLoaderConfig. + """ + return DataLoaderConfig( + batch_size=batch_size, + num_workers=num_workers, + persistent_workers=True, + sampler=DefaultSamplerConfig(shuffle=False), + dataset=dataset, + **kwargs, + ) diff --git a/visdet/py_configs/presets.py b/visdet/py_configs/presets.py new file mode 100644 index 00000000..b913f5a0 --- /dev/null +++ b/visdet/py_configs/presets.py @@ -0,0 +1,294 @@ +"""Pre-defined experiment configurations. + +This module provides ready-to-use experiment configs for common setups. +All presets return ExperimentConfig objects with full IDE autocomplete. + +Example: + >>> from visdet.py_configs import mask_rcnn_swin_tiny_coco + >>> cfg = mask_rcnn_swin_tiny_coco(data_root='/data/coco') + >>> # Customize with autocomplete + >>> cfg.train_cfg.max_epochs = 24 + >>> cfg.optim_wrapper.optimizer.lr = 2e-4 +""" + +from typing import Optional + +from visdet.py_configs.builders import ( + adamw_default, + coco_dataset, + coco_test_pipeline, + coco_train_pipeline, + mask_rcnn, + one_cycle_scheduler, + swin_base, + swin_small, + swin_tiny, + train_dataloader, + val_dataloader, +) +from visdet.schemas import ( + EpochBasedTrainLoopConfig, + ExperimentConfig, + ValLoopConfig, +) + + +def mask_rcnn_swin_tiny_coco( + data_root: str = "/data/coco", + train_ann_file: str = "annotations/instances_train2017.json", + val_ann_file: Optional[str] = "annotations/instances_val2017.json", + train_img_prefix: str = "train2017", + val_img_prefix: str = "val2017", + num_classes: int = 80, + batch_size: int = 2, + num_workers: int = 2, + max_epochs: int = 12, + lr: float = 1e-4, + work_dir: str = "./work_dirs/mask_rcnn_swin_tiny_coco", +) -> ExperimentConfig: + """Mask R-CNN with Swin-Tiny backbone on COCO. + + Standard configuration for instance segmentation training. + + Args: + data_root: COCO data root directory. + train_ann_file: Training annotation file path. + val_ann_file: Validation annotation file path (None to skip validation). + train_img_prefix: Training image directory prefix. + val_img_prefix: Validation image directory prefix. + num_classes: Number of object classes. + batch_size: Batch size per GPU. + num_workers: Data loading workers. + max_epochs: Maximum training epochs. + lr: Learning rate. + work_dir: Output directory. + + Returns: + Complete ExperimentConfig ready for training. + + Example: + >>> cfg = mask_rcnn_swin_tiny_coco( + ... data_root='/my/coco', + ... max_epochs=24, + ... lr=2e-4 + ... ) + >>> from visdet import SimpleRunner + >>> runner = SimpleRunner(config=cfg) + >>> runner.train() + """ + # Model + model = mask_rcnn(backbone=swin_tiny(), num_classes=num_classes) + + # Training data + train_dataset = coco_dataset( + data_root=data_root, + ann_file=train_ann_file, + img_prefix=train_img_prefix, + pipeline=coco_train_pipeline(), + ) + train_dl = train_dataloader( + dataset=train_dataset, + batch_size=batch_size, + num_workers=num_workers, + ) + + # Validation data + val_dl = None + val_cfg = None + val_evaluator = None + if val_ann_file: + val_dataset = coco_dataset( + data_root=data_root, + ann_file=val_ann_file, + img_prefix=val_img_prefix, + pipeline=coco_test_pipeline(), + test_mode=True, + ) + val_dl = val_dataloader( + dataset=val_dataset, + batch_size=1, + num_workers=num_workers, + ) + val_cfg = ValLoopConfig() + val_evaluator = { + "type": "CocoMetric", + "ann_file": f"{data_root}/{val_ann_file}", + "metric": ["bbox", "segm"], + } + + return ExperimentConfig( + model=model, + train_dataloader=train_dl, + val_dataloader=val_dl, + optim_wrapper=adamw_default(lr=lr), + param_scheduler=one_cycle_scheduler(max_lr=lr * 10), + train_cfg=EpochBasedTrainLoopConfig(max_epochs=max_epochs, val_interval=1), + val_cfg=val_cfg, + val_evaluator=val_evaluator, + work_dir=work_dir, + default_hooks={ + "timer": {"type": "IterTimerHook"}, + "logger": {"type": "LoggerHook", "interval": 50}, + "param_scheduler": {"type": "ParamSchedulerHook"}, + "checkpoint": {"type": "CheckpointHook", "interval": 1}, + "sampler_seed": {"type": "DistSamplerSeedHook"}, + "visualization": {"type": "DetVisualizationHook"}, + }, + ) + + +def mask_rcnn_swin_small_coco( + data_root: str = "/data/coco", + train_ann_file: str = "annotations/instances_train2017.json", + val_ann_file: Optional[str] = "annotations/instances_val2017.json", + train_img_prefix: str = "train2017", + val_img_prefix: str = "val2017", + num_classes: int = 80, + batch_size: int = 2, + num_workers: int = 2, + max_epochs: int = 12, + lr: float = 1e-4, + work_dir: str = "./work_dirs/mask_rcnn_swin_small_coco", +) -> ExperimentConfig: + """Mask R-CNN with Swin-Small backbone on COCO. + + Deeper than Swin-Tiny (18 blocks vs 6 in stage 3). + + Args: + data_root: COCO data root directory. + train_ann_file: Training annotation file path. + val_ann_file: Validation annotation file path. + train_img_prefix: Training image directory prefix. + val_img_prefix: Validation image directory prefix. + num_classes: Number of object classes. + batch_size: Batch size per GPU. + num_workers: Data loading workers. + max_epochs: Maximum training epochs. + lr: Learning rate. + work_dir: Output directory. + + Returns: + Complete ExperimentConfig. + """ + # Model with Swin-Small backbone + model = mask_rcnn(backbone=swin_small(), num_classes=num_classes) + + # Training data + train_dataset = coco_dataset( + data_root=data_root, + ann_file=train_ann_file, + img_prefix=train_img_prefix, + pipeline=coco_train_pipeline(), + ) + train_dl = train_dataloader( + dataset=train_dataset, + batch_size=batch_size, + num_workers=num_workers, + ) + + # Validation data + val_dl = None + val_cfg = None + val_evaluator = None + if val_ann_file: + val_dataset = coco_dataset( + data_root=data_root, + ann_file=val_ann_file, + img_prefix=val_img_prefix, + pipeline=coco_test_pipeline(), + test_mode=True, + ) + val_dl = val_dataloader( + dataset=val_dataset, + batch_size=1, + num_workers=num_workers, + ) + val_cfg = ValLoopConfig() + val_evaluator = { + "type": "CocoMetric", + "ann_file": f"{data_root}/{val_ann_file}", + "metric": ["bbox", "segm"], + } + + return ExperimentConfig( + model=model, + train_dataloader=train_dl, + val_dataloader=val_dl, + optim_wrapper=adamw_default(lr=lr), + param_scheduler=one_cycle_scheduler(max_lr=lr * 10), + train_cfg=EpochBasedTrainLoopConfig(max_epochs=max_epochs, val_interval=1), + val_cfg=val_cfg, + val_evaluator=val_evaluator, + work_dir=work_dir, + default_hooks={ + "timer": {"type": "IterTimerHook"}, + "logger": {"type": "LoggerHook", "interval": 50}, + "param_scheduler": {"type": "ParamSchedulerHook"}, + "checkpoint": {"type": "CheckpointHook", "interval": 1}, + "sampler_seed": {"type": "DistSamplerSeedHook"}, + "visualization": {"type": "DetVisualizationHook"}, + }, + ) + + +# ============================================================================= +# Sweep Generators +# ============================================================================= + + +def lr_sweep( + base_preset_fn, + learning_rates: list[float] | None = None, + **preset_kwargs, +): + """Generate configs for learning rate sweep. + + Args: + base_preset_fn: Preset function to use as base (e.g., mask_rcnn_swin_tiny_coco). + learning_rates: List of learning rates to try. + **preset_kwargs: Arguments passed to preset function. + + Yields: + ExperimentConfig for each learning rate. + + Example: + >>> from visdet.py_configs import mask_rcnn_swin_tiny_coco, lr_sweep + >>> for cfg in lr_sweep( + ... mask_rcnn_swin_tiny_coco, + ... learning_rates=[1e-5, 1e-4, 1e-3], + ... data_root='/data/coco' + ... ): + ... # Train each config + ... SimpleRunner(config=cfg).train() + """ + if learning_rates is None: + learning_rates = [1e-5, 3e-5, 1e-4, 3e-4, 1e-3] + + for lr in learning_rates: + cfg = base_preset_fn(lr=lr, **preset_kwargs) + cfg.work_dir = f"{cfg.work_dir}_lr{lr}" + yield cfg + + +def batch_size_sweep( + base_preset_fn, + batch_sizes: list[int] | None = None, + **preset_kwargs, +): + """Generate configs for batch size sweep. + + Args: + base_preset_fn: Preset function to use as base. + batch_sizes: List of batch sizes to try. + **preset_kwargs: Arguments passed to preset function. + + Yields: + ExperimentConfig for each batch size. + """ + if batch_sizes is None: + batch_sizes = [1, 2, 4, 8] + + for bs in batch_sizes: + cfg = base_preset_fn(batch_size=bs, **preset_kwargs) + cfg.work_dir = f"{cfg.work_dir}_bs{bs}" + yield cfg diff --git a/visdet/runner.py b/visdet/runner.py index c4bda6ef..9a9f9baf 100644 --- a/visdet/runner.py +++ b/visdet/runner.py @@ -551,6 +551,61 @@ def _sync_num_classes(self) -> None: f"across {len(roi_head['bbox_head'])} stages (from {source})" ) + @classmethod + def from_config(cls, config: "ExperimentConfig") -> "SimpleRunner": + """Create a SimpleRunner from a Pydantic ExperimentConfig. + + This method provides full IDE autocomplete and type safety for configuration. + Use with configs from visdet.py_configs or visdet.schemas. + + Args: + config: A Pydantic ExperimentConfig from visdet.schemas or visdet.py_configs. + + Returns: + Configured SimpleRunner ready for training. + + Example: + >>> from visdet.py_configs import mask_rcnn_swin_tiny_coco + >>> from visdet import SimpleRunner + >>> + >>> # Create config with full IDE autocomplete + >>> cfg = mask_rcnn_swin_tiny_coco(data_root='/data/coco') + >>> cfg.train_cfg.max_epochs = 24 # Modify with autocomplete + >>> + >>> # Create runner from Pydantic config + >>> runner = SimpleRunner.from_config(cfg) + >>> runner.train() + """ + from visdet.engine import Config + from visdet.schemas.base import VisdetBaseConfig + + # Convert Pydantic model to dict + if isinstance(config, VisdetBaseConfig): + config_dict = config.to_dict() + elif hasattr(config, "model_dump"): + # Pydantic v2 + config_dict = config.model_dump(exclude_none=True) + elif hasattr(config, "dict"): + # Pydantic v1 fallback + config_dict = config.dict(exclude_none=True) + else: + raise TypeError( + f"config must be a Pydantic model (VisdetBaseConfig or similar), " + f"got {type(config).__name__}" + ) + + # Create a new instance that bypasses the normal __init__ + instance = object.__new__(cls) + + # Set attributes directly from the config + instance.epochs = config_dict.get("train_cfg", {}).get("max_epochs", 12) + instance.work_dir = config_dict.get("work_dir", "./work_dirs") + + # Create Config object from dict + instance.cfg = Config(config_dict) + + return instance + def train(self) -> None: """Start training using the assembled configuration. @@ -568,7 +623,8 @@ def train(self) -> None: print("Building runner from config...") runner = MMEngineRunner.from_cfg(self.cfg) - print(f"Starting training for {self.epochs} epochs...") + epochs = getattr(self, "epochs", 12) + print(f"Starting training for {epochs} epochs...") runner.train() # Discoverability class methods diff --git a/visdet/schemas/__init__.py b/visdet/schemas/__init__.py new file mode 100644 index 00000000..19edb208 --- /dev/null +++ b/visdet/schemas/__init__.py @@ -0,0 +1,205 @@ +"""Pydantic configuration schemas for visdet. + +This package provides type-safe configuration with full IDE support. + +Key Benefits: +1. Validation/Type Safety - Catch config errors before runtime +2. Reduce Repetition - Factory functions, inheritance, composition +3. IDE Experience - Autocomplete, jump-to-definition, inline docs +4. Performance - Pydantic v2 uses Rust core, fast validation + +Example: + >>> from visdet.schemas import ( + ... SwinTransformerConfig, + ... FPNConfig, + ... MaskRCNNConfig, + ... ExperimentConfig, + ... ) + >>> # Create a model config with full IDE autocomplete + >>> backbone = SwinTransformerConfig(embed_dims=96, depths=(2, 2, 6, 2)) + >>> neck = FPNConfig(in_channels=[96, 192, 384, 768], out_channels=256, num_outs=5) +""" + +# Base classes +from visdet.schemas.base import ( + BackboneConfig, + ComponentConfig, + DatasetConfig, + HeadConfig, + LossConfig, + NeckConfig, + OptionalConfig, + TransformConfig, + VisdetBaseConfig, +) + +# Backbones +from visdet.schemas.backbones import ( + BackboneType, + ResNetConfig, + ResNeXtConfig, + SwinTransformerConfig, +) + +# Necks +from visdet.schemas.necks import FPNConfig, NeckType + +# Heads and components +from visdet.schemas.heads import ( + AnchorGeneratorConfig, + CrossEntropyLossConfig, + DeltaXYWHBBoxCoderConfig, + FCNMaskHeadConfig, + L1LossConfig, + MaxIoUAssignerConfig, + NMSConfig, + RandomSamplerConfig, + RCNNTestConfig, + RCNNTrainConfig, + RoIAlignConfig, + RoIHeadType, + RPNHeadConfig, + RPNHeadType, + RPNProposalConfig, + RPNTestConfig, + RPNTrainConfig, + Shared2FCBBoxHeadConfig, + SingleRoIExtractorConfig, + SmoothL1LossConfig, + StandardRoIHeadConfig, +) + +# Data and transforms +from visdet.schemas.data import ( + CocoDatasetConfig, + DataLoaderConfig, + DatasetType, + DefaultSamplerConfig, + DetDataPreprocessorConfig, + LoadAnnotationsConfig, + LoadImageFromFileConfig, + PackDetInputsConfig, + RandomFlipConfig, + ResizeConfig, + TransformType, +) + +# Training +from visdet.schemas.training import ( + AdamConfig, + AdamW8bitConfig, + AdamWConfig, + CheckpointHookConfig, + CosineAnnealingLRConfig, + EpochBasedTrainLoopConfig, + IterBasedTrainLoopConfig, + LinearLRConfig, + LoggerHookConfig, + MultiStepLRConfig, + OneCycleLRConfig, + OptimizerType, + OptimWrapperConfig, + SchedulerType, + SGDConfig, + TestLoopConfig, + TrainLoopType, + ValLoopConfig, +) + +# Complete models +from visdet.schemas.models import ( + ExperimentConfig, + FasterRCNNConfig, + MaskRCNNConfig, + ModelType, + TwoStageTestConfig, + TwoStageTrainConfig, +) + +__all__ = [ + # Base + "VisdetBaseConfig", + "ComponentConfig", + "BackboneConfig", + "NeckConfig", + "HeadConfig", + "LossConfig", + "DatasetConfig", + "TransformConfig", + "OptionalConfig", + # Backbones + "SwinTransformerConfig", + "ResNetConfig", + "ResNeXtConfig", + "BackboneType", + # Necks + "FPNConfig", + "NeckType", + # Heads + "RPNHeadConfig", + "StandardRoIHeadConfig", + "Shared2FCBBoxHeadConfig", + "FCNMaskHeadConfig", + "RPNHeadType", + "RoIHeadType", + # Head components + "AnchorGeneratorConfig", + "DeltaXYWHBBoxCoderConfig", + "MaxIoUAssignerConfig", + "RandomSamplerConfig", + "RoIAlignConfig", + "SingleRoIExtractorConfig", + "NMSConfig", + # Losses + "CrossEntropyLossConfig", + "L1LossConfig", + "SmoothL1LossConfig", + # Train/test configs + "RPNTrainConfig", + "RPNProposalConfig", + "RCNNTrainConfig", + "RPNTestConfig", + "RCNNTestConfig", + "TwoStageTrainConfig", + "TwoStageTestConfig", + # Data + "CocoDatasetConfig", + "DataLoaderConfig", + "DetDataPreprocessorConfig", + "DefaultSamplerConfig", + "DatasetType", + # Transforms + "LoadImageFromFileConfig", + "LoadAnnotationsConfig", + "ResizeConfig", + "RandomFlipConfig", + "PackDetInputsConfig", + "TransformType", + # Training + "AdamWConfig", + "SGDConfig", + "AdamConfig", + "AdamW8bitConfig", + "OptimWrapperConfig", + "OptimizerType", + # Schedulers + "OneCycleLRConfig", + "MultiStepLRConfig", + "CosineAnnealingLRConfig", + "LinearLRConfig", + "SchedulerType", + # Loops + "EpochBasedTrainLoopConfig", + "IterBasedTrainLoopConfig", + "ValLoopConfig", + "TestLoopConfig", + "TrainLoopType", + # Hooks + "CheckpointHookConfig", + "LoggerHookConfig", + # Complete models + "MaskRCNNConfig", + "FasterRCNNConfig", + "ExperimentConfig", + "ModelType", +] diff --git a/visdet/schemas/backbones.py b/visdet/schemas/backbones.py new file mode 100644 index 00000000..9ae32286 --- /dev/null +++ b/visdet/schemas/backbones.py @@ -0,0 +1,239 @@ +"""Backbone network configuration schemas. + +This module provides Pydantic schemas for backbone networks with full +IDE autocomplete and validation support. + +Example: + >>> from visdet.schemas.backbones import SwinTransformerConfig + >>> backbone = SwinTransformerConfig(embed_dims=128, depths=(2, 2, 18, 2)) + >>> backbone.to_dict() +""" + +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field + +from visdet.schemas.base import BackboneConfig, OptionalConfig + + +class SwinTransformerConfig(BackboneConfig): + """Swin Transformer backbone configuration. + + A PyTorch implementation of "Swin Transformer: Hierarchical Vision + Transformer using Shifted Windows" (https://arxiv.org/abs/2103.14030). + + Common presets: + - Swin-Tiny: embed_dims=96, depths=(2,2,6,2), num_heads=(3,6,12,24) + - Swin-Small: embed_dims=96, depths=(2,2,18,2), num_heads=(3,6,12,24) + - Swin-Base: embed_dims=128, depths=(2,2,18,2), num_heads=(4,8,16,32) + - Swin-Large: embed_dims=192, depths=(2,2,18,2), num_heads=(6,12,24,48) + + Attributes: + embed_dims: Feature embedding dimension. + depths: Number of blocks at each stage. + num_heads: Number of attention heads at each stage. + window_size: Window size for local attention. + drop_path_rate: Stochastic depth rate. + + Example: + >>> # Swin-Tiny configuration + >>> cfg = SwinTransformerConfig( + ... embed_dims=96, + ... depths=(2, 2, 6, 2), + ... num_heads=(3, 6, 12, 24), + ... ) + """ + + type: Literal["SwinTransformer"] = "SwinTransformer" + + # Image and patch settings + pretrain_img_size: int | tuple[int, int] = Field( + default=224, description="Input image size for pretraining" + ) + in_channels: int = Field(default=3, ge=1, description="Number of input channels") + patch_size: int | tuple[int, int] = Field(default=4, ge=1, description="Patch size") + + # Architecture settings + embed_dims: int = Field( + default=96, gt=0, description="Feature embedding dimension (96=Tiny, 128=Base, 192=Large)" + ) + depths: tuple[int, ...] = Field( + default=(2, 2, 6, 2), + min_length=1, + description="Number of transformer blocks at each stage", + ) + num_heads: tuple[int, ...] = Field( + default=(3, 6, 12, 24), + min_length=1, + description="Number of attention heads at each stage", + ) + window_size: int = Field(default=7, gt=0, description="Window size for local attention") + mlp_ratio: int = Field(default=4, gt=0, description="Ratio of MLP hidden dim to embedding dim") + + # Stride and output settings + strides: tuple[int, ...] = Field( + default=(4, 2, 2, 2), description="Patch merging stride at each stage" + ) + out_indices: tuple[int, ...] = Field( + default=(0, 1, 2, 3), description="Output feature map indices" + ) + + # Attention settings + qkv_bias: bool = Field(default=True, description="Add learnable bias to Q, K, V") + qk_scale: Optional[float] = Field(default=None, description="Override default qk scale") + patch_norm: bool = Field(default=True, description="Apply normalization to patch embedding") + + # Regularization + drop_rate: float = Field(default=0.0, ge=0.0, le=1.0, description="Dropout rate") + attn_drop_rate: float = Field(default=0.0, ge=0.0, le=1.0, description="Attention dropout rate") + drop_path_rate: float = Field( + default=0.1, ge=0.0, le=1.0, description="Stochastic depth rate" + ) + + # Position embedding + use_abs_pos_embed: bool = Field( + default=False, description="Use absolute position embedding" + ) + + # Layer configurations + act_cfg: OptionalConfig = Field( + default_factory=lambda: {"type": "GELU"}, description="Activation layer config" + ) + norm_cfg: OptionalConfig = Field( + default_factory=lambda: {"type": "LN"}, description="Normalization layer config" + ) + + # Training settings + with_cp: bool = Field( + default=False, description="Use gradient checkpointing to save memory" + ) + attn_backend: Literal["torch", "flash"] = Field( + default="torch", description="Attention backend" + ) + + # Pretrained weights + pretrained: Optional[str] = Field(default=None, description="Pretrained checkpoint path") + convert_weights: bool = Field( + default=False, description="Convert weights from official repo format" + ) + + # Freezing + frozen_stages: int = Field( + default=-1, ge=-1, description="Freeze stages up to this index (-1 = none)" + ) + + # Initialization + init_cfg: OptionalConfig = Field(default=None, description="Weight initialization config") + + +class ResNetConfig(BackboneConfig): + """ResNet backbone configuration. + + Standard ResNet architecture with support for ResNet-18/34/50/101/152. + + Common presets: + - ResNet-50: depth=50 (most common for detection) + - ResNet-101: depth=101 (better accuracy, more compute) + + Attributes: + depth: ResNet depth (18, 34, 50, 101, or 152). + out_indices: Which stages to output features from. + frozen_stages: Freeze parameters up to this stage. + norm_eval: Freeze BatchNorm running stats. + + Example: + >>> cfg = ResNetConfig(depth=50, frozen_stages=1) + """ + + type: Literal["ResNet"] = "ResNet" + + # Architecture + depth: Literal[18, 34, 50, 101, 152] = Field( + ..., description="ResNet depth (18, 34, 50, 101, or 152)" + ) + in_channels: int = Field(default=3, ge=1, description="Number of input channels") + stem_channels: Optional[int] = Field( + default=None, description="Stem channels (defaults to base_channels)" + ) + base_channels: int = Field(default=64, gt=0, description="Base channel count") + num_stages: int = Field(default=4, ge=1, le=4, description="Number of ResNet stages") + + # Stride and dilation + strides: tuple[int, ...] = Field( + default=(1, 2, 2, 2), description="Stride of first block in each stage" + ) + dilations: tuple[int, ...] = Field( + default=(1, 1, 1, 1), description="Dilation rate for each stage" + ) + out_indices: tuple[int, ...] = Field( + default=(0, 1, 2, 3), description="Output feature map indices" + ) + + # Style + style: Literal["pytorch", "caffe"] = Field( + default="pytorch", description="ResNet style (affects stride placement)" + ) + deep_stem: bool = Field( + default=False, description="Replace 7x7 conv with three 3x3 convs" + ) + avg_down: bool = Field( + default=False, description="Use AvgPool for downsampling in bottleneck" + ) + + # Freezing and normalization + frozen_stages: int = Field( + default=-1, ge=-1, description="Freeze stages up to this index (-1 = none)" + ) + norm_eval: bool = Field( + default=True, description="Freeze BatchNorm running stats during training" + ) + + # Layer configurations + conv_cfg: OptionalConfig = Field(default=None, description="Convolution layer config") + norm_cfg: OptionalConfig = Field( + default_factory=lambda: {"type": "BN", "requires_grad": True}, + description="Normalization layer config", + ) + + # DCN (Deformable Convolution) - rarely used + dcn: OptionalConfig = Field(default=None, description="DCN config (if using deformable convs)") + stage_with_dcn: tuple[bool, ...] = Field( + default=(False, False, False, False), description="Which stages use DCN" + ) + + # Training settings + with_cp: bool = Field( + default=False, description="Use gradient checkpointing to save memory" + ) + zero_init_residual: bool = Field( + default=True, description="Zero-initialize last norm layer in residual blocks" + ) + + # Pretrained weights + pretrained: Optional[str] = Field(default=None, description="Pretrained checkpoint path") + init_cfg: OptionalConfig = Field(default=None, description="Weight initialization config") + + +class ResNeXtConfig(ResNetConfig): + """ResNeXt backbone configuration. + + ResNeXt extends ResNet with grouped convolutions. + Inherits all ResNet parameters plus group settings. + + Attributes: + groups: Number of groups for grouped convolutions. + base_width: Base width for each group. + """ + + type: Literal["ResNeXt"] = "ResNeXt" + + groups: int = Field(default=32, gt=0, description="Number of groups in grouped convolutions") + base_width: int = Field(default=4, gt=0, description="Base width for each group") + + +# Discriminated union for any backbone type +BackboneType = Annotated[ + Union[SwinTransformerConfig, ResNetConfig, ResNeXtConfig], + Field(discriminator="type"), +] +"""Type alias for any backbone configuration (with discriminator on 'type' field).""" diff --git a/visdet/schemas/base.py b/visdet/schemas/base.py new file mode 100644 index 00000000..157f9eb7 --- /dev/null +++ b/visdet/schemas/base.py @@ -0,0 +1,148 @@ +"""Base Pydantic schema classes for visdet configuration. + +This module provides the foundation for type-safe configuration management. +All component configs inherit from these base classes. + +Key Benefits: +1. Validation/Type Safety - Catch config errors before runtime, clear error messages +2. Reduce Repetition - Factory functions, inheritance, composition +3. IDE Experience - Autocomplete, jump-to-definition, inline docs +4. Performance - Pydantic v2 uses Rust core, fast validation +""" + +from typing import Any, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class VisdetBaseConfig(BaseModel): + """Base class for all visdet configuration schemas. + + Provides common configuration and serialization behavior. + + Attributes: + model_config: Pydantic configuration that: + - extra='forbid': Catches typos by rejecting unknown fields + - validate_default=True: Validates default values + - use_enum_values=True: Serializes enums to their values + + Example: + >>> class MyConfig(VisdetBaseConfig): + ... learning_rate: float = Field(gt=0, description="Learning rate") + >>> cfg = MyConfig(learning_rate=0.001) + >>> cfg.to_dict() + {'learning_rate': 0.001} + """ + + model_config = ConfigDict( + extra="forbid", + validate_default=True, + use_enum_values=True, + populate_by_name=True, + ) + + def to_dict(self) -> dict[str, Any]: + """Convert config to dictionary for MMEngine compatibility. + + Returns: + Dictionary representation of the config. + """ + return self.model_dump(exclude_none=True) + + def to_yaml(self, path: str) -> None: + """Export config to YAML file. + + Args: + path: Output file path. + """ + import yaml + + with open(path, "w") as f: + yaml.dump(self.to_dict(), f, default_flow_style=False, sort_keys=False) + + @classmethod + def from_yaml(cls, path: str) -> "VisdetBaseConfig": + """Load config from YAML file. + + Args: + path: Path to YAML file. + + Returns: + Validated config instance. + """ + import yaml + + with open(path) as f: + data = yaml.safe_load(f) + return cls(**data) + + +class ComponentConfig(VisdetBaseConfig): + """Base class for registry-buildable component configurations. + + Components are classes registered in the visdet registry system. + The 'type' field specifies which class to instantiate. + + Attributes: + type: Registry key for the component class. + + Example: + >>> class SwinConfig(ComponentConfig): + ... type: Literal['SwinTransformer'] = 'SwinTransformer' + ... embed_dims: int = 96 + """ + + type: str = Field(..., description="Registry key for component class") + + +class BackboneConfig(ComponentConfig): + """Base class for backbone network configurations. + + Backbones are feature extraction networks (e.g., ResNet, Swin Transformer). + """ + + pass + + +class NeckConfig(ComponentConfig): + """Base class for neck configurations. + + Necks connect backbones to heads (e.g., FPN). + """ + + pass + + +class HeadConfig(ComponentConfig): + """Base class for detection head configurations. + + Heads perform the final predictions (e.g., RPNHead, RoIHead). + """ + + pass + + +class LossConfig(ComponentConfig): + """Base class for loss function configurations.""" + + pass + + +class DatasetConfig(ComponentConfig): + """Base class for dataset configurations.""" + + pass + + +class TransformConfig(ComponentConfig): + """Base class for data transform configurations.""" + + pass + + +# Type aliases for optional configs (common pattern in MMEngine) +OptionalConfig = Optional[dict[str, Any]] +"""Type alias for optional component config dicts.""" + +ConfigList = list[dict[str, Any]] +"""Type alias for lists of component configs (e.g., data pipelines).""" diff --git a/visdet/schemas/data.py b/visdet/schemas/data.py new file mode 100644 index 00000000..9d7cd955 --- /dev/null +++ b/visdet/schemas/data.py @@ -0,0 +1,220 @@ +"""Dataset and data loading configuration schemas. + +This module provides schemas for datasets, data pipelines, and dataloaders. + +Example: + >>> from visdet.schemas.data import CocoDatasetConfig, DataLoaderConfig + >>> dataset = CocoDatasetConfig( + ... data_root='/data/coco', + ... ann_file='annotations/instances_train2017.json' + ... ) +""" + +from typing import Annotated, Any, Literal, Optional, Union + +from pydantic import Field + +from visdet.schemas.base import ( + ComponentConfig, + ConfigList, + DatasetConfig, + OptionalConfig, + TransformConfig, + VisdetBaseConfig, +) + + +# ============================================================================= +# Transform Configurations +# ============================================================================= + + +class LoadImageFromFileConfig(TransformConfig): + """Load image from file transform.""" + + type: Literal["LoadImageFromFile"] = "LoadImageFromFile" + to_float32: bool = Field(default=False, description="Convert to float32") + color_type: str = Field(default="color", description="Color type") + backend_args: OptionalConfig = Field(default=None, description="Backend args") + + +class LoadAnnotationsConfig(TransformConfig): + """Load annotations transform.""" + + type: Literal["LoadAnnotations"] = "LoadAnnotations" + with_bbox: bool = Field(default=True, description="Load bounding boxes") + with_mask: bool = Field(default=False, description="Load instance masks") + with_seg: bool = Field(default=False, description="Load semantic segmentation") + poly2mask: bool = Field(default=True, description="Convert polygons to masks") + + +class ResizeConfig(TransformConfig): + """Resize image transform.""" + + type: Literal["Resize"] = "Resize" + scale: tuple[int, int] = Field( + default=(1333, 800), description="Target scale (w, h)" + ) + keep_ratio: bool = Field(default=True, description="Keep aspect ratio") + backend: str = Field(default="cv2", description="Resize backend") + + +class RandomFlipConfig(TransformConfig): + """Random horizontal flip transform.""" + + type: Literal["RandomFlip"] = "RandomFlip" + prob: float = Field(default=0.5, ge=0.0, le=1.0, description="Flip probability") + direction: str = Field(default="horizontal", description="Flip direction") + + +class PackDetInputsConfig(TransformConfig): + """Pack detection inputs transform.""" + + type: Literal["PackDetInputs"] = "PackDetInputs" + meta_keys: tuple[str, ...] = Field( + default=("img_id", "img_path", "ori_shape", "img_shape", "scale_factor"), + description="Keys to include in meta info", + ) + + +# ============================================================================= +# Dataset Configurations +# ============================================================================= + + +class CocoDatasetConfig(DatasetConfig): + """COCO format dataset configuration. + + Attributes: + data_root: Root directory containing images and annotations. + ann_file: Path to annotation file (relative to data_root). + data_prefix: Prefix dict for data paths. + pipeline: Data loading and augmentation pipeline. + metainfo: Dataset metadata (classes, palette). + + Example: + >>> cfg = CocoDatasetConfig( + ... data_root='/data/coco', + ... ann_file='annotations/instances_train2017.json', + ... data_prefix={'img': 'train2017'} + ... ) + """ + + type: Literal["CocoDataset"] = "CocoDataset" + data_root: str = Field(..., description="Root directory for data") + ann_file: str = Field(..., description="Annotation file path (relative to data_root)") + data_prefix: dict[str, str] = Field( + default_factory=lambda: {"img": ""}, + description="Prefix paths for data (e.g., {'img': 'train2017'})", + ) + pipeline: ConfigList = Field( + default_factory=list, description="Data loading/augmentation pipeline" + ) + metainfo: OptionalConfig = Field( + default=None, description="Dataset metadata (classes, palette)" + ) + filter_cfg: OptionalConfig = Field( + default_factory=lambda: {"filter_empty_gt": True, "min_size": 32}, + description="Filter config for invalid samples", + ) + backend_args: OptionalConfig = Field(default=None, description="Backend arguments") + test_mode: bool = Field(default=False, description="Test mode (no GT loading)") + + +# ============================================================================= +# Data Preprocessor Configuration +# ============================================================================= + + +class DetDataPreprocessorConfig(ComponentConfig): + """Detection data preprocessor configuration. + + Handles normalization, padding, and format conversion. + """ + + type: Literal["DetDataPreprocessor"] = "DetDataPreprocessor" + mean: tuple[float, float, float] = Field( + default=(123.675, 116.28, 103.53), description="Normalization mean (RGB)" + ) + std: tuple[float, float, float] = Field( + default=(58.395, 57.12, 57.375), description="Normalization std (RGB)" + ) + bgr_to_rgb: bool = Field(default=True, description="Convert BGR to RGB") + pad_mask: bool = Field(default=True, description="Pad instance masks") + pad_size_divisor: int = Field( + default=32, gt=0, description="Pad to multiple of this value" + ) + + +# ============================================================================= +# Sampler Configurations +# ============================================================================= + + +class DefaultSamplerConfig(ComponentConfig): + """Default data sampler configuration.""" + + type: Literal["DefaultSampler"] = "DefaultSampler" + shuffle: bool = Field(default=True, description="Shuffle data") + + +class InfiniteSamplerConfig(ComponentConfig): + """Infinite data sampler configuration.""" + + type: Literal["InfiniteSampler"] = "InfiniteSampler" + shuffle: bool = Field(default=True, description="Shuffle data") + + +# ============================================================================= +# DataLoader Configuration +# ============================================================================= + + +class DataLoaderConfig(VisdetBaseConfig): + """DataLoader configuration. + + Configures batch loading, workers, and sampling. + + Attributes: + batch_size: Samples per batch per GPU. + num_workers: Number of data loading workers. + persistent_workers: Keep workers alive between epochs. + sampler: Sampling strategy configuration. + dataset: Dataset configuration. + + Example: + >>> cfg = DataLoaderConfig( + ... batch_size=2, + ... num_workers=4, + ... dataset=CocoDatasetConfig( + ... data_root='/data/coco', + ... ann_file='annotations/instances_train2017.json' + ... ) + ... ) + """ + + batch_size: int = Field(default=2, ge=1, description="Batch size per GPU") + num_workers: int = Field(default=2, ge=0, description="Number of data loading workers") + persistent_workers: bool = Field( + default=True, description="Keep workers alive between epochs" + ) + sampler: DefaultSamplerConfig | InfiniteSamplerConfig = Field( + default_factory=DefaultSamplerConfig, description="Sampler config" + ) + dataset: CocoDatasetConfig = Field(..., description="Dataset config") + pin_memory: bool = Field(default=True, description="Pin memory for faster transfer") + drop_last: bool = Field(default=False, description="Drop incomplete final batch") + + +# Type aliases +DatasetType = Annotated[Union[CocoDatasetConfig], Field(discriminator="type")] +TransformType = Annotated[ + Union[ + LoadImageFromFileConfig, + LoadAnnotationsConfig, + ResizeConfig, + RandomFlipConfig, + PackDetInputsConfig, + ], + Field(discriminator="type"), +] diff --git a/visdet/schemas/heads.py b/visdet/schemas/heads.py new file mode 100644 index 00000000..36dd3c1a --- /dev/null +++ b/visdet/schemas/heads.py @@ -0,0 +1,410 @@ +"""Detection head configuration schemas. + +This module provides schemas for RPN heads, RoI heads, and their components. +These are the most complex configs in detection as they involve nested +components for anchors, coders, losses, extractors, and samplers. + +Example: + >>> from visdet.schemas.heads import RPNHeadConfig, StandardRoIHeadConfig + >>> rpn = RPNHeadConfig(in_channels=256, feat_channels=256) + >>> roi = StandardRoIHeadConfig(num_classes=80) +""" + +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field + +from visdet.schemas.base import ComponentConfig, HeadConfig, LossConfig, OptionalConfig + + +# ============================================================================= +# Loss Configurations +# ============================================================================= + + +class CrossEntropyLossConfig(LossConfig): + """Cross entropy loss configuration.""" + + type: Literal["CrossEntropyLoss"] = "CrossEntropyLoss" + use_sigmoid: bool = Field(default=False, description="Use sigmoid instead of softmax") + use_mask: bool = Field(default=False, description="Use mask-based loss") + loss_weight: float = Field(default=1.0, gt=0, description="Loss weight") + + +class L1LossConfig(LossConfig): + """L1 (smooth) loss configuration.""" + + type: Literal["L1Loss"] = "L1Loss" + loss_weight: float = Field(default=1.0, gt=0, description="Loss weight") + + +class SmoothL1LossConfig(LossConfig): + """Smooth L1 loss configuration.""" + + type: Literal["SmoothL1Loss"] = "SmoothL1Loss" + beta: float = Field(default=1.0, gt=0, description="Smooth L1 beta parameter") + loss_weight: float = Field(default=1.0, gt=0, description="Loss weight") + + +# ============================================================================= +# Anchor and BBox Coder Configurations +# ============================================================================= + + +class AnchorGeneratorConfig(ComponentConfig): + """Anchor generator configuration. + + Generates anchor boxes at each feature map location. + + Attributes: + scales: Anchor scales (multiplied by stride). + ratios: Anchor aspect ratios. + strides: Feature map strides (one per FPN level). + """ + + type: Literal["AnchorGenerator"] = "AnchorGenerator" + scales: list[float] = Field( + default=[8], description="Anchor scales (multiplied by stride)" + ) + ratios: list[float] = Field( + default=[0.5, 1.0, 2.0], description="Anchor aspect ratios" + ) + strides: list[int] = Field( + default=[4, 8, 16, 32, 64], + description="Feature map strides (typically one per FPN level)", + ) + + +class DeltaXYWHBBoxCoderConfig(ComponentConfig): + """Delta XYWH bounding box coder configuration. + + Encodes/decodes bounding boxes as deltas from anchors. + """ + + type: Literal["DeltaXYWHBBoxCoder"] = "DeltaXYWHBBoxCoder" + target_means: tuple[float, float, float, float] = Field( + default=(0.0, 0.0, 0.0, 0.0), description="Target mean for normalization" + ) + target_stds: tuple[float, float, float, float] = Field( + default=(1.0, 1.0, 1.0, 1.0), description="Target std for normalization" + ) + + +# ============================================================================= +# Assigner and Sampler Configurations +# ============================================================================= + + +class MaxIoUAssignerConfig(ComponentConfig): + """MaxIoU assigner configuration. + + Assigns ground truth to anchors based on IoU thresholds. + """ + + type: Literal["MaxIoUAssigner"] = "MaxIoUAssigner" + pos_iou_thr: float = Field( + default=0.7, ge=0.0, le=1.0, description="IoU threshold for positive assignment" + ) + neg_iou_thr: float = Field( + default=0.3, ge=0.0, le=1.0, description="IoU threshold for negative assignment" + ) + min_pos_iou: float = Field( + default=0.3, ge=0.0, le=1.0, description="Minimum IoU for positive assignment" + ) + match_low_quality: bool = Field( + default=True, description="Match low quality proposals to GT" + ) + ignore_iof_thr: float = Field( + default=-1, description="IoF threshold for ignoring (-1 = disabled)" + ) + + +class RandomSamplerConfig(ComponentConfig): + """Random sampler configuration. + + Randomly samples positive and negative proposals. + """ + + type: Literal["RandomSampler"] = "RandomSampler" + num: int = Field(default=256, gt=0, description="Total number of samples") + pos_fraction: float = Field( + default=0.5, ge=0.0, le=1.0, description="Fraction of positive samples" + ) + neg_pos_ub: int = Field( + default=-1, description="Upper bound on negatives per positive (-1 = unlimited)" + ) + add_gt_as_proposals: bool = Field( + default=False, description="Add GT boxes as proposals" + ) + + +# ============================================================================= +# RoI Extractor Configurations +# ============================================================================= + + +class RoIAlignConfig(ComponentConfig): + """RoI Align layer configuration.""" + + type: Literal["RoIAlign"] = "RoIAlign" + output_size: int = Field(default=7, gt=0, description="Output feature size") + sampling_ratio: int = Field( + default=0, ge=0, description="Sampling ratio (0 = adaptive)" + ) + + +class SingleRoIExtractorConfig(ComponentConfig): + """Single-level RoI feature extractor configuration.""" + + type: Literal["SingleRoIExtractor"] = "SingleRoIExtractor" + roi_layer: RoIAlignConfig = Field( + default_factory=lambda: RoIAlignConfig(), + description="RoI pooling layer config", + ) + out_channels: int = Field(default=256, gt=0, description="Output channels") + featmap_strides: list[int] = Field( + default=[4, 8, 16, 32], description="Feature map strides" + ) + + +# ============================================================================= +# BBox and Mask Head Configurations +# ============================================================================= + + +class Shared2FCBBoxHeadConfig(HeadConfig): + """Shared 2-FC bbox head configuration. + + Standard bbox head with 2 shared fully-connected layers. + """ + + type: Literal["Shared2FCBBoxHead"] = "Shared2FCBBoxHead" + in_channels: int = Field(default=256, gt=0, description="Input channels") + fc_out_channels: int = Field(default=1024, gt=0, description="FC layer output channels") + roi_feat_size: int = Field(default=7, gt=0, description="RoI feature size") + num_classes: int = Field(..., gt=0, description="Number of object classes") + bbox_coder: DeltaXYWHBBoxCoderConfig = Field( + default_factory=lambda: DeltaXYWHBBoxCoderConfig( + target_stds=(0.1, 0.1, 0.2, 0.2) + ), + description="BBox coder config", + ) + reg_class_agnostic: bool = Field( + default=False, description="Use class-agnostic regression" + ) + loss_cls: CrossEntropyLossConfig = Field( + default_factory=CrossEntropyLossConfig, description="Classification loss" + ) + loss_bbox: L1LossConfig = Field( + default_factory=L1LossConfig, description="Regression loss" + ) + + +class FCNMaskHeadConfig(HeadConfig): + """FCN mask head configuration. + + Fully convolutional mask prediction head. + """ + + type: Literal["FCNMaskHead"] = "FCNMaskHead" + num_convs: int = Field(default=4, ge=1, description="Number of conv layers") + in_channels: int = Field(default=256, gt=0, description="Input channels") + conv_out_channels: int = Field(default=256, gt=0, description="Conv output channels") + num_classes: int = Field(..., gt=0, description="Number of object classes") + loss_mask: CrossEntropyLossConfig = Field( + default_factory=lambda: CrossEntropyLossConfig(use_mask=True), + description="Mask loss config", + ) + + +# ============================================================================= +# RPN Head Configuration +# ============================================================================= + + +class RPNHeadConfig(HeadConfig): + """Region Proposal Network (RPN) head configuration. + + Generates object proposals from feature maps. + + Attributes: + in_channels: Number of input feature channels (from neck). + feat_channels: Number of feature channels in conv layers. + anchor_generator: Anchor generation config. + bbox_coder: BBox encoding/decoding config. + loss_cls: Classification loss for objectness. + loss_bbox: Regression loss for box refinement. + + Example: + >>> cfg = RPNHeadConfig( + ... in_channels=256, + ... feat_channels=256, + ... anchor_generator=AnchorGeneratorConfig( + ... strides=[4, 8, 16, 32, 64] + ... ) + ... ) + """ + + type: Literal["RPNHead"] = "RPNHead" + in_channels: int = Field(default=256, gt=0, description="Input feature channels") + feat_channels: int = Field(default=256, gt=0, description="Conv feature channels") + num_convs: int = Field(default=1, ge=1, description="Number of conv layers") + anchor_generator: AnchorGeneratorConfig = Field( + default_factory=AnchorGeneratorConfig, description="Anchor generator config" + ) + bbox_coder: DeltaXYWHBBoxCoderConfig = Field( + default_factory=DeltaXYWHBBoxCoderConfig, description="BBox coder config" + ) + loss_cls: CrossEntropyLossConfig = Field( + default_factory=lambda: CrossEntropyLossConfig(use_sigmoid=True), + description="Classification loss", + ) + loss_bbox: L1LossConfig = Field( + default_factory=L1LossConfig, description="Regression loss" + ) + init_cfg: OptionalConfig = Field( + default_factory=lambda: {"type": "Normal", "layer": "Conv2d", "std": 0.01}, + description="Weight initialization config", + ) + + +# ============================================================================= +# RoI Head Configuration +# ============================================================================= + + +class StandardRoIHeadConfig(HeadConfig): + """Standard RoI head configuration. + + Two-stage detection head with bbox and optional mask branches. + + Attributes: + bbox_roi_extractor: Extracts RoI features for bbox prediction. + bbox_head: Bbox classification and regression head. + mask_roi_extractor: Extracts RoI features for mask prediction (optional). + mask_head: Mask prediction head (optional). + + Example: + >>> cfg = StandardRoIHeadConfig(num_classes=80) + """ + + type: Literal["StandardRoIHead"] = "StandardRoIHead" + bbox_roi_extractor: SingleRoIExtractorConfig = Field( + default_factory=SingleRoIExtractorConfig, + description="BBox RoI feature extractor", + ) + bbox_head: Shared2FCBBoxHeadConfig = Field( + ..., description="BBox prediction head" + ) + mask_roi_extractor: Optional[SingleRoIExtractorConfig] = Field( + default=None, description="Mask RoI feature extractor (None = share with bbox)" + ) + mask_head: Optional[FCNMaskHeadConfig] = Field( + default=None, description="Mask prediction head (optional)" + ) + + +# ============================================================================= +# NMS Configuration +# ============================================================================= + + +class NMSConfig(ComponentConfig): + """Non-Maximum Suppression configuration.""" + + type: Literal["nms"] = "nms" + iou_threshold: float = Field( + default=0.5, ge=0.0, le=1.0, description="IoU threshold for NMS" + ) + + +# ============================================================================= +# Train and Test Configurations +# ============================================================================= + + +class RPNTrainConfig(ComponentConfig): + """RPN training configuration.""" + + type: Literal["RPNTrainConfig"] = Field(default="RPNTrainConfig", exclude=True) + assigner: MaxIoUAssignerConfig = Field( + default_factory=lambda: MaxIoUAssignerConfig( + pos_iou_thr=0.7, neg_iou_thr=0.3, min_pos_iou=0.3 + ), + description="RPN assigner config", + ) + sampler: RandomSamplerConfig = Field( + default_factory=lambda: RandomSamplerConfig( + num=256, pos_fraction=0.5, add_gt_as_proposals=False + ), + description="RPN sampler config", + ) + allowed_border: int = Field(default=-1, description="Allowed border for anchors") + pos_weight: float = Field(default=-1, description="Positive sample weight") + debug: bool = Field(default=False, description="Enable debug mode") + + +class RPNProposalConfig(ComponentConfig): + """RPN proposal generation configuration.""" + + type: Literal["RPNProposalConfig"] = Field(default="RPNProposalConfig", exclude=True) + nms_pre: int = Field(default=2000, gt=0, description="NMS candidates before filtering") + max_per_img: int = Field(default=1000, gt=0, description="Max proposals per image") + nms: NMSConfig = Field( + default_factory=lambda: NMSConfig(iou_threshold=0.7), description="NMS config" + ) + min_bbox_size: int = Field(default=0, ge=0, description="Min bbox size") + + +class RCNNTrainConfig(ComponentConfig): + """RCNN (RoI head) training configuration.""" + + type: Literal["RCNNTrainConfig"] = Field(default="RCNNTrainConfig", exclude=True) + assigner: MaxIoUAssignerConfig = Field( + default_factory=lambda: MaxIoUAssignerConfig( + pos_iou_thr=0.5, neg_iou_thr=0.5, min_pos_iou=0.5 + ), + description="RCNN assigner config", + ) + sampler: RandomSamplerConfig = Field( + default_factory=lambda: RandomSamplerConfig( + num=512, pos_fraction=0.25, add_gt_as_proposals=True + ), + description="RCNN sampler config", + ) + mask_size: int = Field(default=28, gt=0, description="Mask output size") + pos_weight: float = Field(default=-1, description="Positive sample weight") + debug: bool = Field(default=False, description="Enable debug mode") + + +class RPNTestConfig(ComponentConfig): + """RPN test configuration.""" + + type: Literal["RPNTestConfig"] = Field(default="RPNTestConfig", exclude=True) + nms_pre: int = Field(default=1000, gt=0, description="NMS candidates before filtering") + max_per_img: int = Field(default=1000, gt=0, description="Max proposals per image") + nms: NMSConfig = Field( + default_factory=lambda: NMSConfig(iou_threshold=0.7), description="NMS config" + ) + min_bbox_size: int = Field(default=0, ge=0, description="Min bbox size") + + +class RCNNTestConfig(ComponentConfig): + """RCNN test configuration.""" + + type: Literal["RCNNTestConfig"] = Field(default="RCNNTestConfig", exclude=True) + score_thr: float = Field( + default=0.05, ge=0.0, le=1.0, description="Score threshold" + ) + nms: NMSConfig = Field( + default_factory=lambda: NMSConfig(iou_threshold=0.5), description="NMS config" + ) + max_per_img: int = Field(default=100, gt=0, description="Max detections per image") + mask_thr_binary: float = Field( + default=0.5, ge=0.0, le=1.0, description="Binary mask threshold" + ) + + +# Type aliases +RPNHeadType = Annotated[Union[RPNHeadConfig], Field(discriminator="type")] +RoIHeadType = Annotated[Union[StandardRoIHeadConfig], Field(discriminator="type")] diff --git a/visdet/schemas/models.py b/visdet/schemas/models.py new file mode 100644 index 00000000..51223b1c --- /dev/null +++ b/visdet/schemas/models.py @@ -0,0 +1,263 @@ +"""Complete model configuration schemas. + +This module provides schemas for complete detector configurations, +combining backbones, necks, heads, and training settings. + +Example: + >>> from visdet.schemas.models import MaskRCNNConfig + >>> model = MaskRCNNConfig( + ... backbone=SwinTransformerConfig(embed_dims=96), + ... num_classes=80 + ... ) +""" + +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field, model_validator + +from visdet.schemas.backbones import BackboneType, ResNetConfig, SwinTransformerConfig +from visdet.schemas.base import ComponentConfig, OptionalConfig, VisdetBaseConfig +from visdet.schemas.data import DataLoaderConfig, DetDataPreprocessorConfig +from visdet.schemas.heads import ( + FCNMaskHeadConfig, + RCNNTestConfig, + RCNNTrainConfig, + RPNHeadConfig, + RPNProposalConfig, + RPNTestConfig, + RPNTrainConfig, + Shared2FCBBoxHeadConfig, + SingleRoIExtractorConfig, + StandardRoIHeadConfig, +) +from visdet.schemas.necks import FPNConfig, NeckType +from visdet.schemas.training import ( + EpochBasedTrainLoopConfig, + OptimWrapperConfig, + SchedulerType, + ValLoopConfig, +) + + +# ============================================================================= +# Training and Test Configuration +# ============================================================================= + + +class TwoStageTrainConfig(VisdetBaseConfig): + """Training configuration for two-stage detectors.""" + + rpn: RPNTrainConfig = Field( + default_factory=RPNTrainConfig, description="RPN training config" + ) + rpn_proposal: RPNProposalConfig = Field( + default_factory=RPNProposalConfig, description="RPN proposal config" + ) + rcnn: RCNNTrainConfig = Field( + default_factory=RCNNTrainConfig, description="RCNN training config" + ) + + +class TwoStageTestConfig(VisdetBaseConfig): + """Test configuration for two-stage detectors.""" + + rpn: RPNTestConfig = Field( + default_factory=RPNTestConfig, description="RPN test config" + ) + rcnn: RCNNTestConfig = Field( + default_factory=RCNNTestConfig, description="RCNN test config" + ) + + +# ============================================================================= +# Detector Configurations +# ============================================================================= + + +class MaskRCNNConfig(ComponentConfig): + """Mask R-CNN detector configuration. + + Complete configuration for the Mask R-CNN two-stage detector + with instance segmentation support. + + Attributes: + backbone: Feature extraction backbone. + neck: Feature pyramid network. + rpn_head: Region proposal network head. + roi_head: RoI prediction head with bbox and mask branches. + train_cfg: Training configuration. + test_cfg: Testing configuration. + + Example: + >>> from visdet.schemas.backbones import SwinTransformerConfig + >>> cfg = MaskRCNNConfig( + ... backbone=SwinTransformerConfig(embed_dims=96), + ... num_classes=80 + ... ) + """ + + type: Literal["MaskRCNN"] = "MaskRCNN" + + # Data preprocessor + data_preprocessor: DetDataPreprocessorConfig = Field( + default_factory=DetDataPreprocessorConfig, + description="Data preprocessor config", + ) + + # Architecture + backbone: SwinTransformerConfig | ResNetConfig = Field( + ..., description="Backbone network config" + ) + neck: FPNConfig = Field(..., description="Neck network config") + rpn_head: RPNHeadConfig = Field( + default_factory=RPNHeadConfig, description="RPN head config" + ) + roi_head: StandardRoIHeadConfig = Field(..., description="RoI head config") + + # Training and testing + train_cfg: TwoStageTrainConfig = Field( + default_factory=TwoStageTrainConfig, description="Training config" + ) + test_cfg: TwoStageTestConfig = Field( + default_factory=TwoStageTestConfig, description="Testing config" + ) + + @model_validator(mode="before") + @classmethod + def auto_configure_neck(cls, data: dict) -> dict: + """Auto-configure FPN in_channels based on backbone.""" + if "backbone" in data and "neck" in data: + backbone = data["backbone"] + neck = data["neck"] + + # If neck doesn't have in_channels, compute from backbone + if isinstance(neck, dict) and "in_channels" not in neck: + if isinstance(backbone, dict): + backbone_type = backbone.get("type", "") + if backbone_type == "SwinTransformer": + embed_dims = backbone.get("embed_dims", 96) + neck["in_channels"] = [ + embed_dims, + embed_dims * 2, + embed_dims * 4, + embed_dims * 8, + ] + elif backbone_type == "ResNet": + depth = backbone.get("depth", 50) + if depth in [18, 34]: + neck["in_channels"] = [64, 128, 256, 512] + else: # 50, 101, 152 + neck["in_channels"] = [256, 512, 1024, 2048] + + return data + + +class FasterRCNNConfig(ComponentConfig): + """Faster R-CNN detector configuration. + + Two-stage detector without instance segmentation. + """ + + type: Literal["FasterRCNN"] = "FasterRCNN" + + data_preprocessor: DetDataPreprocessorConfig = Field( + default_factory=DetDataPreprocessorConfig, + description="Data preprocessor config", + ) + backbone: SwinTransformerConfig | ResNetConfig = Field( + ..., description="Backbone network config" + ) + neck: FPNConfig = Field(..., description="Neck network config") + rpn_head: RPNHeadConfig = Field( + default_factory=RPNHeadConfig, description="RPN head config" + ) + roi_head: StandardRoIHeadConfig = Field(..., description="RoI head config") + train_cfg: TwoStageTrainConfig = Field( + default_factory=TwoStageTrainConfig, description="Training config" + ) + test_cfg: TwoStageTestConfig = Field( + default_factory=TwoStageTestConfig, description="Testing config" + ) + + +# ============================================================================= +# Complete Experiment Configuration +# ============================================================================= + + +class ExperimentConfig(VisdetBaseConfig): + """Complete experiment configuration. + + Top-level config that combines model, data, training settings. + This is the main config type passed to SimpleRunner. + + Attributes: + model: Complete detector configuration. + train_dataloader: Training data loading config. + val_dataloader: Validation data loading config (optional). + optim_wrapper: Optimizer wrapper config. + param_scheduler: Learning rate scheduler config (optional). + train_cfg: Training loop config. + work_dir: Output directory for logs and checkpoints. + + Example: + >>> from visdet.schemas.models import ExperimentConfig + >>> cfg = ExperimentConfig( + ... model=MaskRCNNConfig(...), + ... train_dataloader=DataLoaderConfig(...), + ... optim_wrapper=OptimWrapperConfig(...), + ... work_dir='./work_dirs/my_experiment' + ... ) + """ + + # Scope + default_scope: str = Field(default="visdet", description="Default registry scope") + + # Model + model: MaskRCNNConfig | FasterRCNNConfig = Field(..., description="Model config") + + # Data + train_dataloader: DataLoaderConfig = Field(..., description="Training dataloader") + val_dataloader: Optional[DataLoaderConfig] = Field( + default=None, description="Validation dataloader" + ) + test_dataloader: Optional[DataLoaderConfig] = Field( + default=None, description="Test dataloader" + ) + + # Optimization + optim_wrapper: OptimWrapperConfig = Field(..., description="Optimizer wrapper") + param_scheduler: Optional[SchedulerType] = Field( + default=None, description="LR scheduler" + ) + + # Training loop + train_cfg: EpochBasedTrainLoopConfig = Field( + default_factory=EpochBasedTrainLoopConfig, description="Training loop config" + ) + val_cfg: Optional[ValLoopConfig] = Field( + default=None, description="Validation loop config" + ) + + # Evaluation + val_evaluator: OptionalConfig = Field( + default=None, description="Validation evaluator config" + ) + + # Output + work_dir: str = Field(default="./work_dirs", description="Output directory") + + # Hooks + default_hooks: OptionalConfig = Field(default=None, description="Default hooks") + + # Logging + log_level: str = Field(default="INFO", description="Logging level") + log_processor: OptionalConfig = Field( + default_factory=lambda: {"window_size": 50}, description="Log processor config" + ) + + +# Type aliases +ModelType = Annotated[ + Union[MaskRCNNConfig, FasterRCNNConfig], Field(discriminator="type") +] diff --git a/visdet/schemas/necks.py b/visdet/schemas/necks.py new file mode 100644 index 00000000..02a10d9d --- /dev/null +++ b/visdet/schemas/necks.py @@ -0,0 +1,127 @@ +"""Neck configuration schemas. + +Necks connect backbone feature maps to detection heads. +FPN (Feature Pyramid Network) is the most commonly used neck. + +Example: + >>> from visdet.schemas.necks import FPNConfig + >>> neck = FPNConfig(in_channels=[256, 512, 1024, 2048], out_channels=256) +""" + +from typing import Annotated, Literal, Union + +from pydantic import Field + +from visdet.schemas.base import NeckConfig, OptionalConfig + + +class FPNConfig(NeckConfig): + """Feature Pyramid Network (FPN) configuration. + + FPN creates a multi-scale feature pyramid from backbone outputs, + enabling detection at multiple scales. + + Reference: "Feature Pyramid Networks for Object Detection" + (https://arxiv.org/abs/1612.03144) + + Attributes: + in_channels: Number of input channels per backbone level. + out_channels: Number of output channels (same for all levels). + num_outs: Number of output feature maps. + + Example: + >>> # FPN for ResNet-50 backbone + >>> cfg = FPNConfig( + ... in_channels=[256, 512, 1024, 2048], + ... out_channels=256, + ... num_outs=5 + ... ) + >>> # FPN for Swin-Tiny backbone + >>> cfg = FPNConfig( + ... in_channels=[96, 192, 384, 768], + ... out_channels=256, + ... num_outs=5 + ... ) + """ + + type: Literal["FPN"] = "FPN" + + # Channel configuration + in_channels: list[int] = Field( + ..., + min_length=1, + description="Number of input channels per backbone level", + ) + out_channels: int = Field( + ..., + gt=0, + description="Number of output channels (used at each scale)", + ) + num_outs: int = Field( + ..., + ge=1, + description="Number of output feature maps (typically 5 for detection)", + ) + + # Level selection + start_level: int = Field( + default=0, + ge=0, + description="Index of the start input backbone level", + ) + end_level: int = Field( + default=-1, + description="End input backbone level index (-1 = last level)", + ) + + # Extra convolutions + add_extra_convs: bool | Literal["on_input", "on_lateral", "on_output"] = Field( + default=False, + description="Add conv layers on top of original feature maps", + ) + relu_before_extra_convs: bool = Field( + default=False, + description="Apply ReLU before extra convolutions", + ) + + # Normalization + no_norm_on_lateral: bool = Field( + default=False, + description="Skip normalization on lateral connections", + ) + + # Layer configurations + conv_cfg: OptionalConfig = Field( + default=None, + description="Convolution layer config", + ) + norm_cfg: OptionalConfig = Field( + default=None, + description="Normalization layer config", + ) + act_cfg: OptionalConfig = Field( + default=None, + description="Activation layer config", + ) + upsample_cfg: OptionalConfig = Field( + default_factory=lambda: {"mode": "nearest"}, + description="Upsampling layer config", + ) + + # Initialization + init_cfg: OptionalConfig = Field( + default_factory=lambda: { + "type": "Xavier", + "layer": "Conv2d", + "distribution": "uniform", + }, + description="Weight initialization config", + ) + + +# Type alias for neck configurations +NeckType = Annotated[ + Union[FPNConfig], + Field(discriminator="type"), +] +"""Type alias for any neck configuration.""" diff --git a/visdet/schemas/training.py b/visdet/schemas/training.py new file mode 100644 index 00000000..16bbba39 --- /dev/null +++ b/visdet/schemas/training.py @@ -0,0 +1,265 @@ +"""Training configuration schemas. + +This module provides schemas for optimizers, schedulers, and training loops. + +Example: + >>> from visdet.schemas.training import AdamWConfig, OneCycleLRConfig + >>> optimizer = AdamWConfig(lr=1e-4, weight_decay=0.05) + >>> scheduler = OneCycleLRConfig(max_lr=1e-3) +""" + +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field + +from visdet.schemas.base import ComponentConfig, OptionalConfig, VisdetBaseConfig + + +# ============================================================================= +# Optimizer Configurations +# ============================================================================= + + +class AdamWConfig(ComponentConfig): + """AdamW optimizer configuration. + + AdamW with decoupled weight decay, recommended for transformers. + + Attributes: + lr: Learning rate. + betas: Adam beta parameters. + weight_decay: Weight decay coefficient. + + Example: + >>> cfg = AdamWConfig(lr=1e-4, weight_decay=0.05) + """ + + type: Literal["AdamW"] = "AdamW" + lr: float = Field(default=1e-4, gt=0, description="Learning rate") + betas: tuple[float, float] = Field( + default=(0.9, 0.999), description="Adam beta parameters" + ) + weight_decay: float = Field( + default=0.05, ge=0, description="Weight decay coefficient" + ) + eps: float = Field(default=1e-8, gt=0, description="Epsilon for numerical stability") + + +class SGDConfig(ComponentConfig): + """SGD optimizer configuration. + + Standard SGD with optional momentum. + + Attributes: + lr: Learning rate. + momentum: Momentum factor. + weight_decay: Weight decay (L2 penalty). + """ + + type: Literal["SGD"] = "SGD" + lr: float = Field(default=0.02, gt=0, description="Learning rate") + momentum: float = Field(default=0.9, ge=0, description="Momentum factor") + weight_decay: float = Field(default=0.0001, ge=0, description="Weight decay") + nesterov: bool = Field(default=False, description="Use Nesterov momentum") + + +class AdamConfig(ComponentConfig): + """Adam optimizer configuration.""" + + type: Literal["Adam"] = "Adam" + lr: float = Field(default=1e-3, gt=0, description="Learning rate") + betas: tuple[float, float] = Field( + default=(0.9, 0.999), description="Adam beta parameters" + ) + weight_decay: float = Field(default=0, ge=0, description="Weight decay") + eps: float = Field(default=1e-8, gt=0, description="Epsilon for numerical stability") + + +# 8-bit optimizers (requires bitsandbytes) +class AdamW8bitConfig(ComponentConfig): + """8-bit AdamW optimizer configuration. + + Memory-efficient 8-bit AdamW from bitsandbytes. + Requires bitsandbytes package. + """ + + type: Literal["AdamW8bit"] = "AdamW8bit" + lr: float = Field(default=1e-4, gt=0, description="Learning rate") + betas: tuple[float, float] = Field(default=(0.9, 0.999), description="Beta parameters") + weight_decay: float = Field(default=0.05, ge=0, description="Weight decay") + eps: float = Field(default=1e-8, gt=0, description="Epsilon") + + +# ============================================================================= +# Optimizer Wrapper Configuration +# ============================================================================= + + +class OptimWrapperConfig(VisdetBaseConfig): + """Optimizer wrapper configuration. + + Wraps an optimizer with additional features like gradient clipping. + """ + + type: Literal["OptimWrapper"] = Field(default="OptimWrapper", description="Wrapper type") + optimizer: AdamWConfig | SGDConfig | AdamConfig | AdamW8bitConfig = Field( + ..., description="Optimizer config" + ) + clip_grad: OptionalConfig = Field( + default=None, description="Gradient clipping config" + ) + accumulative_counts: int = Field( + default=1, ge=1, description="Gradient accumulation steps" + ) + + +# ============================================================================= +# Learning Rate Scheduler Configurations +# ============================================================================= + + +class OneCycleLRConfig(ComponentConfig): + """OneCycle learning rate scheduler configuration. + + Fast.ai's 1cycle policy - warmup, then anneal. + + Attributes: + max_lr: Maximum learning rate at peak. + total_steps: Total training steps (or use epochs). + pct_start: Percentage of cycle spent increasing LR. + """ + + type: Literal["OneCycleLR"] = "OneCycleLR" + max_lr: float = Field(default=1e-3, gt=0, description="Maximum learning rate") + total_steps: Optional[int] = Field( + default=None, description="Total steps (auto-computed if None)" + ) + pct_start: float = Field( + default=0.3, ge=0, le=1, description="Fraction of cycle for warmup" + ) + anneal_strategy: Literal["cos", "linear"] = Field( + default="cos", description="Annealing strategy" + ) + div_factor: float = Field( + default=25.0, gt=0, description="Initial LR = max_lr / div_factor" + ) + final_div_factor: float = Field( + default=1e4, gt=0, description="Final LR = initial_lr / final_div_factor" + ) + + +class MultiStepLRConfig(ComponentConfig): + """Multi-step learning rate scheduler configuration. + + Decays LR by gamma at specified milestones. + """ + + type: Literal["MultiStepLR"] = "MultiStepLR" + milestones: list[int] = Field( + ..., min_length=1, description="Epochs to decay LR" + ) + gamma: float = Field(default=0.1, gt=0, description="Decay factor") + + +class CosineAnnealingLRConfig(ComponentConfig): + """Cosine annealing learning rate scheduler configuration.""" + + type: Literal["CosineAnnealingLR"] = "CosineAnnealingLR" + T_max: int = Field(..., gt=0, description="Maximum number of iterations") + eta_min: float = Field(default=0, ge=0, description="Minimum learning rate") + + +class LinearLRConfig(ComponentConfig): + """Linear learning rate scheduler for warmup.""" + + type: Literal["LinearLR"] = "LinearLR" + start_factor: float = Field( + default=0.001, ge=0, le=1, description="Starting factor" + ) + end_factor: float = Field(default=1.0, ge=0, description="Ending factor") + by_epoch: bool = Field(default=False, description="Step by epoch or iteration") + begin: int = Field(default=0, ge=0, description="Begin epoch/iteration") + end: int = Field(default=500, gt=0, description="End epoch/iteration") + + +# ============================================================================= +# Training Loop Configurations +# ============================================================================= + + +class EpochBasedTrainLoopConfig(ComponentConfig): + """Epoch-based training loop configuration. + + Standard training loop that iterates by epochs. + + Attributes: + max_epochs: Maximum number of training epochs. + val_interval: Validate every N epochs. + """ + + type: Literal["EpochBasedTrainLoop"] = "EpochBasedTrainLoop" + max_epochs: int = Field(default=12, ge=1, description="Maximum epochs") + val_interval: int = Field(default=1, ge=1, description="Validation interval") + + +class IterBasedTrainLoopConfig(ComponentConfig): + """Iteration-based training loop configuration. + + Training loop that iterates by iterations (steps). + """ + + type: Literal["IterBasedTrainLoop"] = "IterBasedTrainLoop" + max_iters: int = Field(..., ge=1, description="Maximum iterations") + val_interval: int = Field(default=5000, ge=1, description="Validation interval") + + +class ValLoopConfig(ComponentConfig): + """Validation loop configuration.""" + + type: Literal["ValLoop"] = "ValLoop" + + +class TestLoopConfig(ComponentConfig): + """Test loop configuration.""" + + type: Literal["TestLoop"] = "TestLoop" + + +# ============================================================================= +# Hook Configurations +# ============================================================================= + + +class CheckpointHookConfig(ComponentConfig): + """Checkpoint saving hook configuration.""" + + type: Literal["CheckpointHook"] = "CheckpointHook" + interval: int = Field(default=1, ge=1, description="Save interval (epochs)") + by_epoch: bool = Field(default=True, description="Interval by epoch or iteration") + save_best: Optional[str] = Field( + default=None, description="Metric to track for best model" + ) + max_keep_ckpts: int = Field(default=3, ge=1, description="Max checkpoints to keep") + + +class LoggerHookConfig(ComponentConfig): + """Logging hook configuration.""" + + type: Literal["LoggerHook"] = "LoggerHook" + interval: int = Field(default=50, ge=1, description="Log interval (iterations)") + log_metric_by_epoch: bool = Field(default=True, description="Log by epoch") + + +# Type aliases +OptimizerType = Annotated[ + Union[AdamWConfig, SGDConfig, AdamConfig, AdamW8bitConfig], + Field(discriminator="type"), +] +SchedulerType = Annotated[ + Union[OneCycleLRConfig, MultiStepLRConfig, CosineAnnealingLRConfig, LinearLRConfig], + Field(discriminator="type"), +] +TrainLoopType = Annotated[ + Union[EpochBasedTrainLoopConfig, IterBasedTrainLoopConfig], + Field(discriminator="type"), +]