Droid210/FleetVision

1"""Swin model architecture for damage detection."""2from torch import nn3from transformers import AutoModelForImageClassification4 5from .config import MODEL_ID, NUM_CLASSES6 7 8def build_model() -> nn.Module:9    """Build Swin Transformer model for binary damage classification.10 11    Returns:12        ViT model with custom classification head.13    """14    # Load pre-trained Swin15    model = AutoModelForImageClassification.from_pretrained(16        MODEL_ID,17        num_labels=NUM_CLASSES,18        ignore_mismatched_sizes=True,19    )20 21    # Freeze transformer backbone (works for Swin and other HF image models)22    if hasattr(model, "swin"):23        for param in model.swin.parameters():24            param.requires_grad = False25    elif hasattr(model, "base_model"):26        for param in model.base_model.parameters():27            param.requires_grad = False28 29    # Only train the classification head30    for param in model.classifier.parameters():31        param.requires_grad = True32 33    return model34