Backbone¶

def create_backbone(cfg: BackboneConfig | str) -> nn.Module:
    """Create a timm backbone from a config or model name string.

    When *cfg* is a plain string it is treated as a model name with
    ``pretrained=True`` and ``num_classes=0`` (headless).

    Supports both timm models and Hugging Face Hub models:
        - Timm models: Use the model name directly (e.g., ``"resnet50"``).
        - HF Hub models: Use the ``hf-hub:`` prefix (e.g., ``"hf-hub:timm/resnet50.a1_in1k"``).

    Examples::

        create_backbone("resnet50")  # Standard timm model
        create_backbone("hf-hub:timm/resnet50.a1_in1k")  # HF Hub model

    Raises ``ValueError`` if the model name is not found.
    """
    if isinstance(cfg, str):
        cfg = BackboneConfig(model_name=cfg)

    # Skip validation for HF Hub models as they're not in timm.list_models()
    if not _is_hf_hub_model(cfg.model_name):
        available = timm.list_models()
        if cfg.model_name not in available:
            raise ValueError(
                f"Model '{cfg.model_name}' not found in timm. "
                f"Use autotimm.list_backbones('*pattern*') to search, "
                f"or use 'hf-hub:' prefix for Hugging Face Hub models."
            )

    return timm.create_model(
        cfg.model_name,
        pretrained=cfg.pretrained,
        num_classes=cfg.num_classes,
        drop_rate=cfg.drop_rate,
        drop_path_rate=cfg.drop_path_rate,
        **cfg.extra_kwargs,
    )

def create_feature_backbone(cfg: FeatureBackboneConfig | str) -> nn.Module:
    """Create a timm backbone for multi-scale feature extraction.

    Uses timm's ``features_only=True`` mode to extract intermediate feature maps,
    typically used for object detection, segmentation, etc.

    When *cfg* is a plain string it is treated as a model name with
    ``pretrained=True`` and default ``out_indices=(1, 2, 3, 4)``.

    Supports both timm models and Hugging Face Hub models:
        - Timm models: Use the model name directly (e.g., ``"resnet50"``).
        - HF Hub models: Use the ``hf-hub:`` prefix (e.g., ``"hf-hub:timm/resnet50.a1_in1k"``).

    Returns a model that outputs a list of feature tensors, one per stage.

    Examples::

        create_feature_backbone("resnet50")  # Standard timm model
        create_feature_backbone("hf-hub:timm/convnext_base.fb_in22k_ft_in1k")  # HF Hub model

    Raises ``ValueError`` if the model name is not found.
    """
    if isinstance(cfg, str):
        cfg = FeatureBackboneConfig(model_name=cfg)

    # Skip validation for HF Hub models as they're not in timm.list_models()
    if not _is_hf_hub_model(cfg.model_name):
        available = timm.list_models()
        if cfg.model_name not in available:
            raise ValueError(
                f"Model '{cfg.model_name}' not found in timm. "
                f"Use autotimm.list_backbones('*pattern*') to search, "
                f"or use 'hf-hub:' prefix for Hugging Face Hub models."
            )

    return timm.create_model(
        cfg.model_name,
        pretrained=cfg.pretrained,
        features_only=True,
        out_indices=cfg.out_indices,
        drop_rate=cfg.drop_rate,
        drop_path_rate=cfg.drop_path_rate,
        **cfg.extra_kwargs,
    )

def get_feature_info(backbone: nn.Module) -> list[dict[str, Any]]:
    """Return feature info for a backbone created with features_only=True.

    Each dict contains:
        - ``num_chs``: Number of output channels
        - ``reduction``: Spatial reduction factor (stride) relative to input
        - ``module``: Name of the module producing this feature

    Only returns info for features that are actually output (respecting out_indices).

    Raises ``AttributeError`` if the backbone was not created with features_only=True.
    """
    if not hasattr(backbone, "feature_info"):
        raise AttributeError(
            "Backbone does not have 'feature_info'. "
            "Ensure it was created with features_only=True."
        )
    # Use get_dicts() to get info only for the output indices
    return backbone.feature_info.get_dicts()

def get_feature_channels(backbone: nn.Module) -> list[int]:
    """Return the number of channels for each feature level.

    Convenience function that extracts just the channel counts from feature_info.
    Only returns channels for features that are actually output (respecting out_indices).
    """
    if not hasattr(backbone, "feature_info"):
        raise AttributeError(
            "Backbone does not have 'feature_info'. "
            "Ensure it was created with features_only=True."
        )
    return backbone.feature_info.channels()

def get_feature_strides(backbone: nn.Module) -> list[int]:
    """Return the stride (spatial reduction) for each feature level.

    Convenience function that extracts just the strides from feature_info.
    Only returns strides for features that are actually output (respecting out_indices).
    """
    if not hasattr(backbone, "feature_info"):
        raise AttributeError(
            "Backbone does not have 'feature_info'. "
            "Ensure it was created with features_only=True."
        )
    return backbone.feature_info.reduction()

def list_backbones(
    pattern: str = "",
    pretrained_only: bool = False,
    with_source: bool = False,
) -> list[str] | list[tuple[str, ModelSource]]:
    """List available timm backbones, optionally filtered by glob *pattern*.

    This lists models available in the timm library. To search Hugging Face Hub
    models, use ``list_hf_hub_backbones()``.

    Parameters:
        pattern: Glob pattern to filter model names (e.g., ``"*resnet*"``).
        pretrained_only: If True, only list models with pretrained weights.
        with_source: If True, return tuples of (model_name, ModelSource.TIMM).

    Returns:
        List of model names, or list of (model_name, source) tuples if
        ``with_source=True``.

    Examples::

        list_backbones("*resnet*")
        list_backbones("*efficientnet*", pretrained_only=True)
        list_backbones("*vit*", with_source=True)  # [(name, ModelSource.TIMM), ...]
    """
    models = timm.list_models(pattern or "*", pretrained=pretrained_only)
    if with_source:
        return [(name, ModelSource.TIMM) for name in models]
    return models