src/models/classification.py

from __future__ import annotations

from typing import Dict

import torch
import timm
from torchvision.models import (
    ConvNeXt_Tiny_Weights,
    EfficientNet_B0_Weights,
    MobileNet_V3_Small_Weights,
    ResNet18_Weights,
    Swin_T_Weights,
    ViT_B_16_Weights,
    convnext_tiny,
    efficientnet_b0,
    mobilenet_v3_small,
    resnet18,
    swin_t,
    vit_b_16,
)

from src.models.lora import (
    apply_lora_to_timm_attention,
    enable_lora_suffix_no_grad,
    freeze_all_parameters_except_lora_and_heads,
)

_TIMM_CLASSIFIER_MODELS = {
    "vit_s16_cls": "vit_small_patch16_224",
    "vit_l16_cls": "vit_large_patch16_224",
}


def _resolve_timm_classifier_name(model_cfg: Dict) -> str | None:
    model_name = str(model_cfg.get("name", "resnet18_cls")).lower()
    default_timm_name = _TIMM_CLASSIFIER_MODELS.get(model_name)
    if default_timm_name is None:
        return None
    return str(model_cfg.get("timm_model_name", default_timm_name))


def _pretrained_weights(pretrained: bool, weights_enum):
    return weights_enum.DEFAULT if pretrained else None


def _make_head(in_features: int, num_classes: int, dropout_p: float) -> torch.nn.Module:
    if dropout_p > 0.0:
        return torch.nn.Sequential(
            torch.nn.Dropout(p=dropout_p),
            torch.nn.Linear(in_features, num_classes),
        )
    return torch.nn.Linear(in_features, num_classes)


class MultiHeadResNet18Classifier(torch.nn.Module):
    def __init__(
        self,
        pretrained: bool,
        num_classes: int,
        aux_dims: dict[str, int],
        dropout_p: float,
    ) -> None:
        super().__init__()
        self.backbone = resnet18(weights=_pretrained_weights(pretrained, ResNet18_Weights))
        in_features = self.backbone.fc.in_features
        self.backbone.fc = torch.nn.Identity()
        self.class_head = _make_head(in_features, num_classes, dropout_p)
        self.brand_head = _make_head(in_features, aux_dims["brand"], dropout_p)
        self.year_head = _make_head(in_features, aux_dims["year"], dropout_p)

    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        features = self.backbone(x)
        return {
            "class": self.class_head(features),
            "brand": self.brand_head(features),
            "year": self.year_head(features),
        }


def get_classifier_data_config(model_cfg: Dict) -> dict:
    timm_model_name = _resolve_timm_classifier_name(model_cfg)
    if timm_model_name is None:
        return {}

    pretrained_cfg = timm.get_pretrained_cfg(timm_model_name)
    if pretrained_cfg is None:
        return {}

    return {
        "source": "timm_pretrained_cfg",
        "input_size": list(pretrained_cfg.input_size),
        "interpolation": pretrained_cfg.interpolation,
        "crop_pct": float(pretrained_cfg.crop_pct),
        "normalization": {
            "mean": [float(v) for v in pretrained_cfg.mean],
            "std": [float(v) for v in pretrained_cfg.std],
        },
    }


def build_classifier(model_cfg: Dict, num_classes: int, aux_dims: dict[str, int] | None = None):
    model_name = str(model_cfg.get("name", "resnet18_cls")).lower()
    pretrained = bool(model_cfg.get("pretrained", False))
    dropout_p = float(model_cfg.get("dropout", 0.0))
    use_aux_heads = bool(model_cfg.get("aux_heads", {}).get("enabled", False))
    freeze_backbone = bool(model_cfg.get("freeze_backbone", False))
    lora_cfg = model_cfg.get("lora", {})
    lora_enabled = bool(lora_cfg.get("enabled", False))

    if lora_enabled and model_name not in ("vit_s16_cls", "vit_l16_cls"):
        raise ValueError("LoRA attention injection is currently supported only for vit_s16_cls and vit_l16_cls (timm VisionTransformer)")

    if model_name == "resnet18_cls":
        if use_aux_heads:
            if aux_dims is None:
                raise ValueError("aux_dims is required when aux_heads.enabled=true")
            model = MultiHeadResNet18Classifier(
                pretrained=pretrained,
                num_classes=num_classes,
                aux_dims=aux_dims,
                dropout_p=dropout_p,
            )
        else:
            model = resnet18(weights=_pretrained_weights(pretrained, ResNet18_Weights))
            in_features = model.fc.in_features
            model.fc = _make_head(in_features, num_classes, dropout_p)
    elif model_name == "mobilenet_v3_small_cls":
        model = mobilenet_v3_small(weights=_pretrained_weights(pretrained, MobileNet_V3_Small_Weights))
        in_features = model.classifier[-1].in_features
        model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
    elif model_name == "efficientnet_b0_cls":
        model = efficientnet_b0(weights=_pretrained_weights(pretrained, EfficientNet_B0_Weights))
        in_features = model.classifier[-1].in_features
        model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
    elif model_name == "convnext_tiny_cls":
        model = convnext_tiny(weights=_pretrained_weights(pretrained, ConvNeXt_Tiny_Weights))
        in_features = model.classifier[-1].in_features
        model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
    elif model_name == "vit_b16_cls":
        model = vit_b_16(weights=_pretrained_weights(pretrained, ViT_B_16_Weights))
        in_features = model.heads.head.in_features
        model.heads.head = _make_head(in_features, num_classes, dropout_p)
    elif model_name in ("vit_s16_cls", "vit_l16_cls"):
        timm_model_name = _resolve_timm_classifier_name(model_cfg)
        if timm_model_name is None:
            raise ValueError(f"timm_model_name could not be resolved for {model_name}")
        model = timm.create_model(
            timm_model_name,
            pretrained=pretrained,
            num_classes=num_classes,
            drop_rate=dropout_p,
        )
        if lora_enabled:
            lora_targets = tuple(lora_cfg.get("target_modules", ("qkv", "proj")))
            injected_layers = apply_lora_to_timm_attention(
                model,
                rank=int(lora_cfg.get("rank", 8)),
                alpha=float(lora_cfg.get("alpha", lora_cfg.get("rank", 8))),
                dropout_p=float(lora_cfg.get("dropout", 0.0)),
                target_modules=lora_targets,
                last_n_blocks=int(lora_cfg["last_n_blocks"]) if lora_cfg.get("last_n_blocks") is not None else None,
                block_indices=lora_cfg.get("block_indices"),
            )
            model._lora_injected_layers = injected_layers
            target_block_indices = tuple(getattr(model, "_lora_target_block_indices", ()))
            if freeze_backbone and target_block_indices:
                try:
                    frozen_prefix_blocks = enable_lora_suffix_no_grad(model, target_block_indices)
                    model._lora_no_grad_prefix_count = frozen_prefix_blocks
                except ValueError:
                    model._lora_no_grad_prefix_count = 0
    elif model_name == "swin_t_cls":
        model = swin_t(weights=_pretrained_weights(pretrained, Swin_T_Weights))
        in_features = model.head.in_features
        model.head = _make_head(in_features, num_classes, dropout_p)
    else:
        raise ValueError(
            f"Unsupported classification model '{model_name}'. "
            "Supported: resnet18_cls, mobilenet_v3_small_cls, efficientnet_b0_cls, "
            "convnext_tiny_cls, vit_b16_cls, vit_s16_cls, swin_t_cls"
        )

    if freeze_backbone:
        freeze_all_parameters_except_lora_and_heads(model)

    return model