src/model.py

import segmentation_models_pytorch as smp


def build_model(config):
    """Build U-Net with pretrained encoder and multi-channel sigmoid output.

    Each output channel predicts one mask type independently (multi-label).
    Returns raw logits — apply sigmoid in loss/inference.
    """
    model = smp.Unet(
        encoder_name=config.encoder_name,
        encoder_weights=config.encoder_weights,
        in_channels=3,
        classes=config.num_classes,
        activation=None,
        decoder_attention_type=config.decoder_attention,
    )
    return model