segmentation

SegmentationDataModule(data_path, name='segmentation_datamodule', test_size=0.3, val_size=0.3, seed=42, dataset=SegmentationDataset, batch_size=32, num_workers=6, train_transform=None, test_transform=None, val_transform=None, train_split_file=None, test_split_file=None, val_split_file=None, num_data_class=None, exclude_good=False, **kwargs)

Bases: BaseDataModule

Base class for segmentation datasets.

Parameters:

• data_path (str) –

  Path to the data main folder.

• name (str, default: 'segmentation_datamodule' ) –

  The name for the data module. Defaults to "segmentation_datamodule".

• val_size (float, default: 0.3 ) –

  The validation split. Defaults to 0.2.

• test_size (float, default: 0.3 ) –

  The test split. Defaults to 0.2.

• seed (int, default: 42 ) –

  Random generator seed. Defaults to 42.

• dataset (type[SegmentationDataset], default: SegmentationDataset ) –

  Dataset class.

• batch_size (int, default: 32 ) –

  Batch size. Defaults to 32.

• num_workers (int, default: 6 ) –

  Number of workers for dataloaders. Defaults to 16.

• train_transform (Compose | None, default: None ) –

  Transformations for train dataset. Defaults to None.

• val_transform (Compose | None, default: None ) –

  Transformations for validation dataset. Defaults to None.

• test_transform (Compose | None, default: None ) –

  Transformations for test dataset. Defaults to None.

• num_data_class (int | None, default: None ) –

  The number of samples per class. Defaults to None.

• exclude_good (bool, default: False ) –

  If True, exclude good samples from the dataset. Defaults to False.

Source code in quadra/datamodules/segmentation.py

def __init__(
    self,
    data_path: str,
    name: str = "segmentation_datamodule",
    test_size: float = 0.3,
    val_size: float = 0.3,
    seed: int = 42,
    dataset: type[SegmentationDataset] = SegmentationDataset,
    batch_size: int = 32,
    num_workers: int = 6,
    train_transform: albumentations.Compose | None = None,
    test_transform: albumentations.Compose | None = None,
    val_transform: albumentations.Compose | None = None,
    train_split_file: str | None = None,
    test_split_file: str | None = None,
    val_split_file: str | None = None,
    num_data_class: int | None = None,
    exclude_good: bool = False,
    **kwargs: Any,
):
    super().__init__(
        data_path=data_path,
        name=name,
        seed=seed,
        batch_size=batch_size,
        num_workers=num_workers,
        train_transform=train_transform,
        val_transform=val_transform,
        test_transform=test_transform,
        **kwargs,
    )
    self.test_size = test_size
    self.val_size = val_size
    self.num_data_class = num_data_class
    self.exclude_good = exclude_good
    self.train_split_file = train_split_file
    self.test_split_file = test_split_file
    self.val_split_file = val_split_file
    self.dataset = dataset
    self.train_dataset: SegmentationDataset
    self.val_dataset: SegmentationDataset
    self.test_dataset: SegmentationDataset

predict_dataloader()

Returns a dataloader used for predictions.

Source code in quadra/datamodules/segmentation.py

def predict_dataloader(self) -> DataLoader:
    """Returns a dataloader used for predictions."""
    return self.test_dataloader()

setup(stage=None)

Setup data module based on stages of training.

Source code in quadra/datamodules/segmentation.py

def setup(self, stage=None):
    """Setup data module based on stages of training."""
    if stage in ["fit", "train"]:
        self.train_dataset = self.dataset(
            image_paths=self.data[self.data["split"] == "train"]["samples"].tolist(),
            mask_paths=self.data[self.data["split"] == "train"]["masks"].tolist(),
            mask_preprocess=self._preprocess_mask,
            labels=self.data[self.data["split"] == "train"]["targets"].tolist(),
            object_masks=None,
            transform=self.train_transform,
            batch_size=None,
            defect_transform=None,
            resize=None,
        )
        self.val_dataset = self.dataset(
            image_paths=self.data[self.data["split"] == "val"]["samples"].tolist(),
            mask_paths=self.data[self.data["split"] == "val"]["masks"].tolist(),
            defect_transform=None,
            labels=self.data[self.data["split"] == "val"]["targets"].tolist(),
            object_masks=None,
            batch_size=None,
            mask_preprocess=self._preprocess_mask,
            transform=self.test_transform,
            resize=None,
        )
    elif stage == "test":
        self.test_dataset = self.dataset(
            image_paths=self.data[self.data["split"] == "test"]["samples"].tolist(),
            mask_paths=self.data[self.data["split"] == "test"]["masks"].tolist(),
            labels=self.data[self.data["split"] == "test"]["targets"].tolist(),
            object_masks=None,
            batch_size=None,
            mask_preprocess=self._preprocess_mask,
            transform=self.test_transform,
            resize=None,
        )
    elif stage == "predict":
        pass
    else:
        raise ValueError(f"Unknown stage {stage}")

test_dataloader()

Returns the test dataloader.

Raises:

• ValueError –

  If test dataset is not initialized.

Returns:

• DataLoader –

  test dataloader.

Source code in quadra/datamodules/segmentation.py

def test_dataloader(self) -> DataLoader:
    """Returns the test dataloader.

    Raises:
        ValueError: If test dataset is not initialized.


    Returns:
        test dataloader.
    """
    if not self.test_dataset_available:
        raise ValueError("Test dataset is not initialized")

    loader = DataLoader(
        self.test_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )
    return loader

train_dataloader()

Returns the train dataloader.

Raises:

• ValueError –

  If train dataset is not initialized.

Returns:

• DataLoader –

  Train dataloader.

Source code in quadra/datamodules/segmentation.py

def train_dataloader(self) -> DataLoader:
    """Returns the train dataloader.

    Raises:
        ValueError: If train dataset is not initialized.

    Returns:
        Train dataloader.
    """
    if not self.train_dataset_available:
        raise ValueError("Train dataset is not initialized")

    return DataLoader(
        self.train_dataset,
        batch_size=self.batch_size,
        shuffle=True,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )

val_dataloader()

Returns the validation dataloader.

Raises:

• ValueError –

  If validation dataset is not initialized.

Returns:

• DataLoader –

  val dataloader.

Source code in quadra/datamodules/segmentation.py

def val_dataloader(self) -> DataLoader:
    """Returns the validation dataloader.

    Raises:
        ValueError: If validation dataset is not initialized.

    Returns:
        val dataloader.
    """
    if not self.val_dataset_available:
        raise ValueError("Validation dataset is not initialized")

    return DataLoader(
        self.val_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )

SegmentationMulticlassDataModule(data_path, idx_to_class, name='multiclass_segmentation_datamodule', dataset=SegmentationDatasetMulticlass, batch_size=32, test_size=0.3, val_size=0.3, seed=42, num_workers=6, train_transform=None, test_transform=None, val_transform=None, train_split_file=None, test_split_file=None, val_split_file=None, exclude_good=False, num_data_train=None, one_hot_encoding=False, **kwargs)

Bases: BaseDataModule

Base class for segmentation datasets with multiple classes.

Parameters:

• data_path –

  Path to the data main folder.

• idx_to_class (dict) –

  dict with corrispondence btw mask index and classes: {1: class_1, 2: class_2, ..., N: class_N} except background class which is 0.

• name –

  The name for the data module. Defaults to "multiclass_segmentation_datamodule".

• dataset (type[SegmentationDatasetMulticlass], default: SegmentationDatasetMulticlass ) –

  Dataset class.

• batch_size –

  Batch size. Defaults to 32.

• val_size –

  The validation split. Defaults to 0.3.

• test_size –

  The test split. Defaults to 0.3.

• seed –

  Random generator seed. Defaults to 42.

• num_workers (int, default: 6 ) –

  Number of workers for dataloaders. Defaults to 6.

• train_transform (Compose | None, default: None ) –

  Transformations for train dataset. Defaults to None.

• val_transform –

  Transformations for validation dataset. Defaults to None.

• test_transform –

  Transformations for test dataset. Defaults to None.

• train_split_file (str | None, default: None ) –

  path to txt file with training samples list

• val_split_file (str | None, default: None ) –

  path to txt file with validation samples list

• test_split_file (str | None, default: None ) –

  path to txt file with test samples list

• exclude_good –

  If True, exclude good samples from the dataset. Defaults to False.

• num_data_train (int | None, default: None ) –

  number of samples to use in the train split (shuffle the samples and pick the first num_data_train)

• one_hot_encoding (bool, default: False ) –

  if True, the labels are one-hot encoded to N channels, where N is the number of classes. If False, masks are single channel that contains values as class indexes. Defaults to True.

Source code in quadra/datamodules/segmentation.py

def __init__(
    self,
    data_path: str,
    idx_to_class: dict,
    name: str = "multiclass_segmentation_datamodule",
    dataset: type[SegmentationDatasetMulticlass] = SegmentationDatasetMulticlass,
    batch_size: int = 32,
    test_size: float = 0.3,
    val_size: float = 0.3,
    seed: int = 42,
    num_workers: int = 6,
    train_transform: albumentations.Compose | None = None,
    test_transform: albumentations.Compose | None = None,
    val_transform: albumentations.Compose | None = None,
    train_split_file: str | None = None,
    test_split_file: str | None = None,
    val_split_file: str | None = None,
    exclude_good: bool = False,
    num_data_train: int | None = None,
    one_hot_encoding: bool = False,
    **kwargs: Any,
):
    super().__init__(
        data_path=data_path,
        name=name,
        seed=seed,
        batch_size=batch_size,
        num_workers=num_workers,
        train_transform=train_transform,
        val_transform=val_transform,
        test_transform=test_transform,
        **kwargs,
    )
    self.test_size = test_size
    self.val_size = val_size
    self.exclude_good = exclude_good
    self.train_split_file = train_split_file
    self.test_split_file = test_split_file
    self.val_split_file = val_split_file
    self.dataset = dataset
    self.idx_to_class = idx_to_class
    self.num_data_train = num_data_train
    self.one_hot_encoding = one_hot_encoding
    self.train_dataset: SegmentationDataset
    self.val_dataset: SegmentationDataset
    self.test_dataset: SegmentationDataset

predict_dataloader()

Returns a dataloader used for predictions.

Source code in quadra/datamodules/segmentation.py

def predict_dataloader(self) -> DataLoader:
    """Returns a dataloader used for predictions."""
    return self.test_dataloader()

setup(stage=None)

Setup data module based on stages of training.

Source code in quadra/datamodules/segmentation.py

def setup(self, stage=None):
    """Setup data module based on stages of training."""
    if stage in ["fit", "train"]:
        train_data = self.data[self.data["split"] == "train"]
        val_data = self.data[self.data["split"] == "val"]

        self.train_dataset = self.dataset(
            image_paths=train_data["samples"].tolist(),
            mask_paths=train_data["masks"].tolist(),
            idx_to_class=self.idx_to_class,
            transform=self.train_transform,
            one_hot=self.one_hot_encoding,
        )
        self.val_dataset = self.dataset(
            image_paths=val_data["samples"].tolist(),
            mask_paths=val_data["masks"].tolist(),
            transform=self.val_transform,
            idx_to_class=self.idx_to_class,
            one_hot=self.one_hot_encoding,
        )
    elif stage == "test":
        self.test_dataset = self.dataset(
            image_paths=self.data[self.data["split"] == "test"]["samples"].tolist(),
            mask_paths=self.data[self.data["split"] == "test"]["masks"].tolist(),
            transform=self.test_transform,
            idx_to_class=self.idx_to_class,
            one_hot=self.one_hot_encoding,
        )
    elif stage == "predict":
        pass
    else:
        raise ValueError(f"Unknown stage {stage}")

test_dataloader()

Returns the test dataloader.

Raises:

• ValueError –

  If test dataset is not initialized.

Returns:

• DataLoader –

  test dataloader.

Source code in quadra/datamodules/segmentation.py

def test_dataloader(self) -> DataLoader:
    """Returns the test dataloader.

    Raises:
        ValueError: If test dataset is not initialized.


    Returns:
        test dataloader.
    """
    if not self.test_dataset_available:
        raise ValueError("Test dataset is not initialized")

    loader = DataLoader(
        self.test_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )
    return loader

train_dataloader()

Returns the train dataloader.

Raises:

• ValueError –

  If train dataset is not initialized.

Returns:

• DataLoader –

  Train dataloader.

Source code in quadra/datamodules/segmentation.py

def train_dataloader(self) -> DataLoader:
    """Returns the train dataloader.

    Raises:
        ValueError: If train dataset is not initialized.

    Returns:
        Train dataloader.
    """
    if not self.train_dataset_available:
        raise ValueError("Train dataset is not initialized")

    return DataLoader(
        self.train_dataset,
        batch_size=self.batch_size,
        shuffle=True,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )

val_dataloader()

Returns the validation dataloader.

Raises:

• ValueError –

  If validation dataset is not initialized.

Returns:

• DataLoader –

  val dataloader.

Source code in quadra/datamodules/segmentation.py

def val_dataloader(self) -> DataLoader:
    """Returns the validation dataloader.

    Raises:
        ValueError: If validation dataset is not initialized.

    Returns:
        val dataloader.
    """
    if not self.val_dataset_available:
        raise ValueError("Validation dataset is not initialized")

    return DataLoader(
        self.val_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        drop_last=False,
        pin_memory=True,
        persistent_workers=self.num_workers > 0,
    )