Skip to content

smpmodule

SMPModule

Bases: SegmentationModule

Source code in src/tcd_pipeline/models/smpmodule.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class SMPModule(SegmentationModule):
    def configure_models(self, init_pretrained=False) -> None:
        """Configures the task baself, inited on kwargs parameters passed to the constructor."""

        if self.hparams["model"] == "unet":
            self.model = smp.Unet(
                encoder_name=self.hparams["backbone"],
                encoder_weights=self.hparams["weights"] if init_pretrained else None,
                in_channels=self.hparams["in_channels"],
                classes=self.hparams["num_classes"],
            )
        elif self.hparams["model"] == "deeplabv3+":
            self.model = smp.DeepLabV3Plus(
                encoder_name=self.hparams["backbone"],
                encoder_weights=self.hparams["weights"] if init_pretrained else None,
                in_channels=self.hparams["in_channels"],
                classes=self.hparams["num_classes"],
            )
        elif self.hparams["model"] == "unet++":
            self.model = smp.UnetPlusPlus(
                encoder_name=self.hparams["backbone"],
                encoder_weights=self.hparams["weights"] if init_pretrained else None,
                in_channels=self.hparams["in_channels"],
                classes=self.hparams["num_classes"],
            )
        else:
            raise ValueError(
                f"Model type '{self.hparams['model']}' is not valid. "
                f"Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
            )

    def configure_losses(self) -> None:
        """Initialize the loss criterion.

        Raises:
            ValueError: If *loss* is invalid.
        """
        loss: str = self.hparams["loss"]

        ignore_index = self.hparams.get("ignore_index")
        weight = self.hparams.get("class_weights")

        if loss == "ce":
            ignore_value = -1000 if ignore_index is None else ignore_index
            self.criterion = nn.CrossEntropyLoss(
                ignore_index=ignore_value, weight=weight
            )
        elif loss == "jaccard":
            self.criterion = smp.losses.JaccardLoss(
                mode="multiclass", classes=self.hparams["num_classes"]
            )
        elif loss == "focal":
            self.criterion = smp.losses.FocalLoss(
                "multiclass", ignore_index=ignore_index, normalized=True
            )
        elif loss == "tversky":
            self.criterion = smp.losses.TverskyLoss(
                "multiclass", ignore_index=ignore_index, normalized=True
            )
        else:
            raise ValueError(
                f"Loss type '{loss}' is not valid. "
                "Currently, supports 'ce', 'jaccard' or 'focal' loss."
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of the model.

        Args:
            x: Image tensor

        Returns:
            Tensor with semantic segmentation predictionss
        """
        return self.predict_step(x)

    def predict_step(
        self,
        batch: Union[List, torch.Tensor],
        batch_idx: int = 0,
        dataloader_idx: int = 0,
    ) -> torch.Tensor:
        if isinstance(batch, list):
            batch = torch.stack(batch)

        logits = self.model(batch)
        return logits

configure_losses()

Initialize the loss criterion.

Raises:

Type Description
ValueError

If loss is invalid.
Source code in src/tcd_pipeline/models/smpmodule.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def configure_losses(self) -> None:
    """Initialize the loss criterion.

    Raises:
        ValueError: If *loss* is invalid.
    """
    loss: str = self.hparams["loss"]

    ignore_index = self.hparams.get("ignore_index")
    weight = self.hparams.get("class_weights")

    if loss == "ce":
        ignore_value = -1000 if ignore_index is None else ignore_index
        self.criterion = nn.CrossEntropyLoss(
            ignore_index=ignore_value, weight=weight
        )
    elif loss == "jaccard":
        self.criterion = smp.losses.JaccardLoss(
            mode="multiclass", classes=self.hparams["num_classes"]
        )
    elif loss == "focal":
        self.criterion = smp.losses.FocalLoss(
            "multiclass", ignore_index=ignore_index, normalized=True
        )
    elif loss == "tversky":
        self.criterion = smp.losses.TverskyLoss(
            "multiclass", ignore_index=ignore_index, normalized=True
        )
    else:
        raise ValueError(
            f"Loss type '{loss}' is not valid. "
            "Currently, supports 'ce', 'jaccard' or 'focal' loss."
        )

configure_models(init_pretrained=False)

Configures the task baself, inited on kwargs parameters passed to the constructor.

Source code in src/tcd_pipeline/models/smpmodule.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def configure_models(self, init_pretrained=False) -> None:
    """Configures the task baself, inited on kwargs parameters passed to the constructor."""

    if self.hparams["model"] == "unet":
        self.model = smp.Unet(
            encoder_name=self.hparams["backbone"],
            encoder_weights=self.hparams["weights"] if init_pretrained else None,
            in_channels=self.hparams["in_channels"],
            classes=self.hparams["num_classes"],
        )
    elif self.hparams["model"] == "deeplabv3+":
        self.model = smp.DeepLabV3Plus(
            encoder_name=self.hparams["backbone"],
            encoder_weights=self.hparams["weights"] if init_pretrained else None,
            in_channels=self.hparams["in_channels"],
            classes=self.hparams["num_classes"],
        )
    elif self.hparams["model"] == "unet++":
        self.model = smp.UnetPlusPlus(
            encoder_name=self.hparams["backbone"],
            encoder_weights=self.hparams["weights"] if init_pretrained else None,
            in_channels=self.hparams["in_channels"],
            classes=self.hparams["num_classes"],
        )
    else:
        raise ValueError(
            f"Model type '{self.hparams['model']}' is not valid. "
            f"Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
        )

forward(x)

Forward pass of the model.

Parameters:

Name Type Description Default
x Tensor

Image tensor

 required

Returns:

Type Description
Tensor

Tensor with semantic segmentation predictionss
Source code in src/tcd_pipeline/models/smpmodule.py 
80
81
82
83
84
85
86
87
88
89
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass of the model.

    Args:
        x: Image tensor

    Returns:
        Tensor with semantic segmentation predictionss
    """
    return self.predict_step(x)