Upload baselines/models/__init__.py with huggingface_hub

baselines/models/__init__.py

@@ -0,0 +1,93 @@
+"""Model factory for deterministic segmentation baselines."""
+
+
+def get_model(model_name, in_channels=1, num_classes=1):
+    """
+    Get a segmentation model by name using official libraries.
+    
+    Args:
+        model_name: One of 'unet', 'attention_unet', 'unetpp', 'transunet'
+        in_channels: Number of input channels (1 for grayscale)
+        num_classes: Number of output classes (1 for binary)
+    
+    Returns:
+        PyTorch model
+    """
+    if model_name == "unet":
+        # From segmentation_models_pytorch
+        import segmentation_models_pytorch as smp
+        model = smp.Unet(
+            encoder_name="resnet34",
+            encoder_weights="imagenet",
+            in_channels=in_channels,
+            classes=num_classes,
+            activation=None,  # We apply sigmoid in loss/prediction
+        )
+        return model
+    
+    elif model_name == "unetpp":
+        # From segmentation_models_pytorch
+        import segmentation_models_pytorch as smp
+        model = smp.UnetPlusPlus(
+            encoder_name="resnet34",
+            encoder_weights="imagenet",
+            in_channels=in_channels,
+            classes=num_classes,
+            activation=None,
+        )
+        return model
+    
+    elif model_name == "attention_unet":
+        # From MONAI
+        from monai.networks.nets import AttentionUnet
+        model = AttentionUnet(
+            spatial_dims=2,
+            in_channels=in_channels,
+            out_channels=num_classes,
+            channels=(64, 128, 256, 512),
+            strides=(2, 2, 2),
+            dropout=0.1,
+        )
+        return model
+    
+    elif model_name == "transunet":
+        # Custom TransUNet using smp with a ViT-style encoder
+        # We use smp's MAnet with a MiT (Mix Transformer) encoder 
+        # which is the closest official arch to TransUNet in smp
+        import segmentation_models_pytorch as smp
+        model = smp.MAnet(
+            encoder_name="mit_b2",
+            encoder_weights="imagenet",
+            in_channels=in_channels,
+            classes=num_classes,
+            activation=None,
+        )
+        return model
+    
+    elif model_name == "nnunet":
+        # nnU-Net architecture (PlainConvUNet) from dynamic_network_architectures
+        # Using the same config that nnU-Net auto-planned for our 128x128 dataset
+        from dynamic_network_architectures.architectures.unet import PlainConvUNet
+        import torch.nn as nn
+        model = PlainConvUNet(
+            input_channels=in_channels,
+            n_stages=6,
+            features_per_stage=(32, 64, 128, 256, 512, 512),
+            conv_op=nn.Conv2d,
+            kernel_sizes=[[3, 3]] * 6,
+            strides=[[1, 1]] + [[2, 2]] * 5,
+            n_conv_per_stage=[2, 2, 2, 2, 2, 2],
+            num_classes=num_classes,
+            n_conv_per_stage_decoder=[2, 2, 2, 2, 2],
+            conv_bias=True,
+            norm_op=nn.InstanceNorm2d,
+            norm_op_kwargs={"eps": 1e-5, "affine": True},
+            dropout_op=None,
+            nonlin=nn.LeakyReLU,
+            nonlin_kwargs={"inplace": True},
+            deep_supervision=False,
+        )
+        return model
+
+    else:
+        raise ValueError(f"Unknown model: {model_name}. Choose from: unet, attention_unet, unetpp, transunet, nnunet")