U-Net#

torchgeo.models.unet(weights=None, classes=None, *args, **kwargs)[source]#

U-Net model.

If you use this model in your research, please cite the following paper:

Added in version 0.8.

Parameters:

  • weights (Unet_Weights | None) – Pre-trained model weights to use.

  • classes (int | None) – Number of output classes. If not specified, the number of classes will be inferred from the weights.

  • *args (Any) – Additional arguments to pass to segmentation_models_pytorch.create_model

  • **kwargs (Any) – Additional keyword arguments to pass to segmentation_models_pytorch.create_model

Returns:

A U-Net model.

Return type:

Module

class torchgeo.models.Unet_Weights(*values)[source]#

Bases: WeightsEnum

U-Net weights.

For smp Unet implementation.

Added in version 0.8.

__new__(value)#