otx.algo.segmentation.segmentors#

Module for base NN segmentation models.

Classes

BaseSegmentationModel(backbone, decode_head)

Base Segmentation Model.

MeanTeacher(model[, unsup_weight, ...])

MeanTeacher for Semi-supervised learning.
class otx.algo.segmentation.segmentors.BaseSegmentationModel(backbone: Module, decode_head: Module, criterion: Module | None = None)[source]#

Bases: Module

Base Segmentation Model.

Parameters:

  • backbone (nn.Module) – The backbone of the segmentation model.

  • decode_head (nn.Module) – The decode head of the segmentation model.

  • criterion (nn.Module, optional) – The criterion of the model. Defaults to None. If None, use CrossEntropyLoss with ignore_index=255.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

calculate_loss(model_features: Tensor, img_metas: list[ImageInfo], masks: Tensor, interpolate: bool) → Tensor[source]#

Calculates the loss of the model.

Parameters:

  • model_features (Tensor) – model outputs of the model.

  • img_metas (list[ImageInfo]) – Image meta information. Defaults to None.

  • masks (Tensor) – Ground truth masks for training. Defaults to None.

Returns:

The loss of the model.

Return type:

Tensor

extract_features(inputs: Tensor) → tuple[Tensor, Tensor][source]#

Extract features from the backbone and head.

forward(inputs: Tensor, img_metas: list[ImageInfo] | None = None, masks: Tensor | None = None, mode: str = 'tensor') → Tensor[source]#

Performs the forward pass of the model.

Parameters:

  • inputs (Tensor) – Input images to the model.

  • img_metas (list[ImageInfo]) – Image meta information. Defaults to None.

  • masks (Tensor) – Ground truth masks for training. Defaults to None.

  • mode (str) – The mode of operation. Defaults to “tensor”.

Returns:

  • If mode is “tensor”, returns the model outputs.

  • If mode is “loss”, returns a dictionary of output losses.

  • If mode is “predict”, returns the predicted outputs.

  • Otherwise, returns the model outputs after interpolation.

Return type:

Depending on the mode

get_valid_label_mask(img_metas: list[ImageInfo]) → list[Tensor][source]#

Get valid label mask removing ignored classes to zero mask in a batch.

Parameters:

img_metas (List[dict]) – List of image metadata.

Returns:

List of valid label masks.

Return type:

List[torch.Tensor]

class otx.algo.segmentation.segmentors.MeanTeacher(model: Module, unsup_weight: float = 1.0, drop_unrel_pixels_percent: int = 20, semisl_start_epoch: int = 0, filter_pixels_epochs: int = 100)[source]#

Bases: Module

MeanTeacher for Semi-supervised learning.

Parameters:

  • model (nn.Module) – model

  • unsup_weight (float, optional) – unsupervised weight. Defaults to 1.0.

  • drop_unrel_pixels_percent (int, optional) – drop unrel pixels percent. Defaults to 20.

  • semisl_start_epoch (int, optional) – semisl start epoch. Defaults to 0.

  • filter_pixels_epochs (int, optional) – filter pixels epochs. Defaults to 100.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor, unlabeled_weak_images: Tensor | None = None, unlabeled_strong_images: Tensor | None = None, global_step: int | None = None, steps_per_epoch: int | None = None, img_metas: list[ImageInfo] | None = None, unlabeled_img_metas: list[ImageInfo] | None = None, masks: Tensor | None = None, mode: str = 'tensor') → Tensor[source]#

Step for model training.

Parameters:

  • inputs (Tensor) – input labeled images

  • unlabeled_weak_images (Tensor, optional) – unlabeled images with weak augmentations. Defaults to None.

  • unlabeled_strong_images (Tensor, optional) – unlabeled images with strong augmentations. Defaults to None.

  • global_step (int, optional) – global step. Defaults to None.

  • steps_per_epoch (int, optional) – steps per epoch. Defaults to None.

  • img_metas (list[ImageInfo], optional) – image meta information. Defaults to None.

  • unlabeled_img_metas (list[ImageInfo], optional) – unlabeled image meta information. Defaults to None.

  • masks (Tensor, optional) – ground truth masks for training. Defaults to None.

  • mode (str, optional) – mode of forward. Defaults to “tensor”.