Skip to content

semantic_segmentation

Semantic segmentation model framework, using smp models

SemanticSegmentationModel

Bases: Model

Tiled model subclass for smp semantic segmentation models.

Source code in src/tcd_pipeline/models/semantic_segmentation.py 
 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
class SemanticSegmentationModel(Model):
    """Tiled model subclass for smp semantic segmentation models."""

    def __init__(self, config):
        """Initialize the model. Does not load weights - this is
        performed automatically at prediction or train time.

        Args:
            config (dict): Configuration dictionary
        """
        super().__init__(config)
        self.post_processor = SemanticSegmentationPostProcessor(config)
        self._cfg = config

    @abstractmethod
    def forward(self):
        """Model forward pass (i.e. predict), sub-classed by specific architectures"""

    def predict_batch(self, image_tensor: List[torch.Tensor]) -> List[torch.Tensor]:
        """Run inference on a batch of images

        Args:
            image_tensor (List[torch.Tensor]): List of images to predict

        Returns:
            predictions (List[torch.Tensor]): List of output prediction tensors
        """

        self.model.eval()
        self.should_reload = False
        predictions = None

        t_start_s = time.time()

        with torch.no_grad():
            # removing alpha channel
            inputs = [im[:3, :, :].to(self.device) for im in image_tensor]

            try:
                predictions = self.forward(inputs)
            except RuntimeError as e:
                logger.error("Runtime error: %s", e)
                self.should_reload = True
            except Exception as e:  # pylint: disable=broad-except
                logger.error(
                    "Failed to run inference: %s. Attempting to reload model.", e
                )
                self.should_reload = True

        t_elapsed_s = time.time() - t_start_s
        logger.debug("Predicted tile in %1.2fs", t_elapsed_s)

        assert len(predictions.shape) == 4

        return [p for p in predictions]

    def evaluate(self, augment=False):
        """
        Evaluate the model on the dataset provided in the config.

        Does not log to wandb.
        """

        import os

        import lightning.pytorch as pl
        from lightning.pytorch.loggers import CSVLogger

        from tcd_pipeline.data.datamodule import COCODataModule

        self.load_model()

        # Rely on Lightning for directory setup within this folder
        log_dir = self.config.data.output

        os.makedirs(log_dir, exist_ok=True)

        # Evaluate without augmentation
        data_module = COCODataModule(
            self.config.data.root,
            train_path=self.config.data.train,
            val_path=self.config.data.validation,
            test_path=self.config.data.validation,
            augment=augment,
            batch_size=int(self.config.model.batch_size),
            num_workers=int(self.config.model.num_workers),
            tile_size=int(self.config.data.tile_size),
        )

        csv_logger = CSVLogger(save_dir=log_dir, name="logs")

        # Eval "trainer"
        evaluator = pl.Trainer(
            logger=[csv_logger],
            default_root_dir=log_dir,
            accelerator=self.device,
            devices=1,
        )

        from tcd_pipeline.models.segformermodule import SegformerModule
        from tcd_pipeline.models.smpmodule import SMPModule

        if self.config.model.name == "segformer":
            module = SegformerModule(
                model=self.config.model.name,
                backbone=self.config.model.backbone,
                ignore_index=None,
                id2label=os.path.join(
                    os.path.dirname(__file__), "index_to_name_binary.json"
                ),
            )
        elif self.config.model.name == "unet":
            module = SMPModule(
                model=self.config.model.name,
                backbone=self.config.model.backbone,
                weights=self.config.model.pretrained,
                in_channels=int(self.config.model.in_channels),
                num_classes=int(self.config.model.num_classes),
                loss=self.config.model.loss,
            )
        else:
            raise NotImplementedError

        # Drop in model we've just loaded
        logger.info("Initialising empty Lightning module")
        module.configure_models(init_pretrained=False)
        module.configure_losses()
        module.configure_metrics()
        logger.info(f"Mapping model weights from {self.config.model.weights}")
        module.model = self.model

        try:
            logger.info("Starting evaluation on test data")
            evaluator.test(model=module, datamodule=data_module)
        # pylint: disable=broad-except
        except Exception as e:
            logger.error("Evaluation failed")
            logger.error(e)
            logger.error(traceback.print_exc())

__init__(config)

Initialize the model. Does not load weights - this is performed automatically at prediction or train time.

Parameters:

Name Type Description Default
config dict

Configuration dictionary

 required
Source code in src/tcd_pipeline/models/semantic_segmentation.py 
23
24
25
26
27
28
29
30
31
32
def __init__(self, config):
    """Initialize the model. Does not load weights - this is
    performed automatically at prediction or train time.

    Args:
        config (dict): Configuration dictionary
    """
    super().__init__(config)
    self.post_processor = SemanticSegmentationPostProcessor(config)
    self._cfg = config

evaluate(augment=False)

Evaluate the model on the dataset provided in the config.

Does not log to wandb.

Source code in src/tcd_pipeline/models/semantic_segmentation.py 
 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def evaluate(self, augment=False):
    """
    Evaluate the model on the dataset provided in the config.

    Does not log to wandb.
    """

    import os

    import lightning.pytorch as pl
    from lightning.pytorch.loggers import CSVLogger

    from tcd_pipeline.data.datamodule import COCODataModule

    self.load_model()

    # Rely on Lightning for directory setup within this folder
    log_dir = self.config.data.output

    os.makedirs(log_dir, exist_ok=True)

    # Evaluate without augmentation
    data_module = COCODataModule(
        self.config.data.root,
        train_path=self.config.data.train,
        val_path=self.config.data.validation,
        test_path=self.config.data.validation,
        augment=augment,
        batch_size=int(self.config.model.batch_size),
        num_workers=int(self.config.model.num_workers),
        tile_size=int(self.config.data.tile_size),
    )

    csv_logger = CSVLogger(save_dir=log_dir, name="logs")

    # Eval "trainer"
    evaluator = pl.Trainer(
        logger=[csv_logger],
        default_root_dir=log_dir,
        accelerator=self.device,
        devices=1,
    )

    from tcd_pipeline.models.segformermodule import SegformerModule
    from tcd_pipeline.models.smpmodule import SMPModule

    if self.config.model.name == "segformer":
        module = SegformerModule(
            model=self.config.model.name,
            backbone=self.config.model.backbone,
            ignore_index=None,
            id2label=os.path.join(
                os.path.dirname(__file__), "index_to_name_binary.json"
            ),
        )
    elif self.config.model.name == "unet":
        module = SMPModule(
            model=self.config.model.name,
            backbone=self.config.model.backbone,
            weights=self.config.model.pretrained,
            in_channels=int(self.config.model.in_channels),
            num_classes=int(self.config.model.num_classes),
            loss=self.config.model.loss,
        )
    else:
        raise NotImplementedError

    # Drop in model we've just loaded
    logger.info("Initialising empty Lightning module")
    module.configure_models(init_pretrained=False)
    module.configure_losses()
    module.configure_metrics()
    logger.info(f"Mapping model weights from {self.config.model.weights}")
    module.model = self.model

    try:
        logger.info("Starting evaluation on test data")
        evaluator.test(model=module, datamodule=data_module)
    # pylint: disable=broad-except
    except Exception as e:
        logger.error("Evaluation failed")
        logger.error(e)
        logger.error(traceback.print_exc())

forward() abstractmethod

Model forward pass (i.e. predict), sub-classed by specific architectures

Source code in src/tcd_pipeline/models/semantic_segmentation.py 
34
35
36
@abstractmethod
def forward(self):
    """Model forward pass (i.e. predict), sub-classed by specific architectures"""

predict_batch(image_tensor)

Run inference on a batch of images

Parameters:

Name Type Description Default
image_tensor List[Tensor]

List of images to predict

 required

Returns:

Name Type Description
predictions List[Tensor]

List of output prediction tensors
Source code in src/tcd_pipeline/models/semantic_segmentation.py 
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
def predict_batch(self, image_tensor: List[torch.Tensor]) -> List[torch.Tensor]:
    """Run inference on a batch of images

    Args:
        image_tensor (List[torch.Tensor]): List of images to predict

    Returns:
        predictions (List[torch.Tensor]): List of output prediction tensors
    """

    self.model.eval()
    self.should_reload = False
    predictions = None

    t_start_s = time.time()

    with torch.no_grad():
        # removing alpha channel
        inputs = [im[:3, :, :].to(self.device) for im in image_tensor]

        try:
            predictions = self.forward(inputs)
        except RuntimeError as e:
            logger.error("Runtime error: %s", e)
            self.should_reload = True
        except Exception as e:  # pylint: disable=broad-except
            logger.error(
                "Failed to run inference: %s. Attempting to reload model.", e
            )
            self.should_reload = True

    t_elapsed_s = time.time() - t_start_s
    logger.debug("Predicted tile in %1.2fs", t_elapsed_s)

    assert len(predictions.shape) == 4

    return [p for p in predictions]