im2sim.models.ReverseHalfUNet

im2sim.models.ReverseHalfUNet#

class ReverseHalfUNet(in_channels, out_channels, rank, cfg, supervision_levels=0)[source]#

Bases: Module

A flexible implementation of a ‘Reverse HalfUNet’, which is inspired by a HalfUNet[1] architecture but where the HalfUNet does away with the decoders of a conventional UNet[2], this model instead removes the encoders. Can be used for image segmentation and reconstruction tasks where blurring caused by additive fusion as in a HalfUNet is undesirable and/or deep supervision is required.

Parameters:

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • rank (int) – Spatial rank (1D, 2D, 3D).

  • cfg (ReverseHalfUNetConfig) – Configuration object for the Reverse ReverseHalfUNet.

  • supervision_levels (int | list[int]) – Levels at which to apply deep supervision. 0 corresponds to the highest resolution output, 1 to the next lower resolution, and so on. Default is 0 (no deep supervision).

Examples

To create a ReverseHalfUNet model with a specific configuration, you can first create a ReverseHalfUNetConfig object and then pass it to the ReverseHalfUNet constructor. For example, to create a ReverseHalfUNet with 3 levels of depth, ReLU activation, and softmax output activation:

cfg = ReverseHalfUNetConfig(
           hidden_channels=32,
           n_levels=3,
           decoder_block_cfg=ImageConvBlockConfig(depth=3, activation="ReLU"),
           out_block_cfg=ImageConvBlockConfig(depth=1, activation="ReLU", out_activation="softmax")
      )

Since the configs are rankless, you could use the same config for a 1D, 2D, or 3D convolutional block by changing the rank parameter when creating the ReverseHalfUNet instance.

model1D = ReverseHalfUNet(
       rank=1,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )
model2D = ReverseHalfUNet(
       rank=2,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )
model3D = ReverseHalfUNet(
       rank=3,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
   )

The ReverseHalfUNet model can be used for both segmentation and reconstruction tasks. For segmentation, you can use the single_class_segmentation_mode() or multiclass_segmentation_mode() methods of the ReverseHalfUNetConfig to set the appropriate output activation function (sigmoid for single-class, softmax for multi-class). For reconstruction tasks, you can use the reconstruction_mode() method to set the output activation to None.

cfg_segmentation = ReverseHalfUNetConfig().single_class_segmentation_mode()
model_segmentation = ReverseHalfUNet(
       rank=2,
       in_channels=32,
       out_channels=1,
       cfg=cfg_segmentation,
   )

cfg_reconstruction = ReverseHalfUNetConfig().reconstruction_mode()
model_reconstruction = ReverseHalfUNet(
       rank=2,
       in_channels=32,
       out_channels=1,
       cfg=cfg_reconstruction,
   )

If deep supervision is desired, you can specify the levels at which to apply it using the supervision_levels argument.

model_deep_supervision = ReverseHalfUNet(
       rank=2,
       in_channels=32,
       out_channels=32,
       cfg=cfg,
       supervision_levels=[0, 1],  # Apply deep supervision at top 2 levels
   )

Models can be saved and loaded using the standard PyTorch methods:

torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth"))

Configs can also be saved and loaded using the methods provided in the im2sim.configs.UNetConfig class:

cfg.save("my_config.yaml")
loaded_cfg = ReverseHalfUNetConfig.load("my_config.yaml")
model = ReverseHalfUNet(
       rank=2,
       in_channels=32,
       out_channels=32,
       cfg=loaded_cfg,
   )

References

forward(x)[source]#

Forward pass through the Reverse ReverseHalfUNet.

Contents