README.md

Models

🧠 ModelFactory.py

Creates model instances from config. The factory looks up config.model['name'] in the models module namespace and passes every other key of config.model as a constructor keyword argument (flat keys, no params wrapper), e.g. "model": {"name": "UNet3D", "in_channels": 4, "num_classes": 4, "size": 32, "depth": 4}.

Available Models

🧠 UNet2D.py:

• 2D U-Net architecture
• Configurable: in_channels, num_classes, size=32 (base number of features), depth=3 (number of down-sampling levels), t_prime_init=0.1, bias_init=-10.0 (learnable SigLIP temperature/bias parameters)
• Exposes bottleneck_channels (size * 2**(depth+1)); forward(x, return_bottleneck=False) returns (logits, bottleneck) when return_bottleneck=True, as required by BaseTrainerText
• Inputs whose spatial size is not divisible by 2**depth are padded and unpadded automatically (utils/pad_unpad.py)

🧠 UNet3D.py:

• 3D U-Net architecture
• Configurable: in_channels, num_classes, size=32, depth=3, t_prime_init=0.1, bias_init=-10.0
• Inputs whose spatial size is not divisible by 2**depth are padded and unpadded automatically (utils/pad_unpad.py)

🧠 UNet3DText.py:

• UNet3D with a built-in 3D ReportGuidanceHead on the bottleneck
• Configurable: in_channels, num_classes, size=32, depth=3, text_dim=768, guidance_hidden_dim=1024, guidance_out_dim=512
• forward(x, text_emb=None): returns the segmentation logits, or (logits, z_img, z_txt) when text_emb is given (used by Trainer_3DText)

🧠 TextGuidanceHead.py:

• ReportGuidanceHead(bottleneck_channels, text_dim=768, hidden_dim=1024, out_dim=512, spatial_dims=2): not a standalone model; projects the bottleneck with a 1x1 conv + global average pooling (spatial_dims 2 or 3) and the text embedding with an MLP (Linear -> GELU -> LayerNorm -> Linear) into a shared out_dim space
• forward(feat, text_emb) returns (z_img [B, out_dim], z_txt [B, out_dim])
• Used by BaseTrainerText (2D) and UNet3DText (3D)

Adding Custom Models

1. Create a new file in models/ with your model class that inherits from BaseModel (or directly from nn.Module)
2. Implement __init__() to define layers
3. Implement forward() method that must at least return the output of the model
4. Add your model class to __init__.py in the models/ folder so the ModelFactory can find it:

   from .YourModel import YourModel

5. To use it, specify in config: "model": {"name": "YourModel", "num_classes": ..., "arg1": ...} (num_classes is required, since it is also read by the trainer and the metrics)