Class segmentation_task_predict

src/flair/task_module.py:174–213

class segmentation_task_predict(pl.LightningModule):
    def __init__(
        self,
        model,
        num_classes : int,
        use_metadata : bool = False,
    ):

        super().__init__()
        self.model = model
        self.num_classes = num_classes
        self.use_metadata=use_metadata

    def forward(self, input_im, input_met):
        logits = self.model(input_im, input_met)
        return logits

    def step(self, batch):
        if self.use_metadata == True:
            images, metadata, targets = batch["img"], batch["mtd"], batch["msk"]
        else:
            images, metadata, targets = batch["img"], '', batch["msk"]            
        logits = self.forward(images, metadata)

        with torch.no_grad():
            proba = torch.softmax(logits, dim=1)
            preds = torch.argmax(proba, dim=1)
            targets = torch.argmax(targets, dim=1)
            preds = preds.flatten(start_dim=1)  # Change shapes and cast target to integer for metrics computation
            targets = targets.flatten(start_dim=1).type(torch.int32)
        return preds, targets

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        if self.use_metadata == True:
            logits = self.forward(batch["img"], batch["mtd"])
        else:
            logits = self.forward(batch["img"], '')
        proba = torch.softmax(logits, dim=1)
        batch["preds"] =  torch.argmax(proba, dim=1)
        return batch