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