0
$\begingroup$

I am building a MANet model using pytorch lightning. For getting the model I use the library segmentation models. As my objective is to do binary semantic segmentation, during the test phase I calculate the IoU score, however all the images have a reduced mask which means the mask's positive pixels percentage is about 3% so the vanilla IoU score doesnt take that in account during the calculation. In order to overcome this limitation I am trying to calculate the weighted IoU score but I do not understand how to pass the weights class list as input to the iou_score function:

smp.metrics.iou_score(tp, fp, fn, tn, reduction='weighted', class_weights=[0.52, 16.67])

For completeness that is the code called by the training, validation and test step inside the model:

def shared_step(self, batch, stage):
    # batch[0] -> image
    # batch[1] -> mask

    image = batch[0]

    # (batch_size, num_channels, height, width)
    assert image.ndim == 4

    # Check that image dimensions are divisible by 32
    h, w = image.shape[2:]
    assert h % 32 == 0 and w % 32 == 0

    mask = batch[1]

    assert mask.ndim == 4 # [batch_size, num_classes, height, width] = dim 4

    # Check that mask values in between 0 and 1, NOT 0 and 255 for binary segmentation
    assert mask.max() <= 1.0 and mask.min() >= 0

    logits_mask = self.forward(image) # -> output grezzo del modello
    loss = self.loss_fn(logits_mask, mask) # calcolo funzione di loss = errore
    self.log_dict({f"{stage}/loss": loss.detach().item()}, batch_size=config["batch_size"]) # log metrica loss
    prob_mask = logits_mask.sigmoid() # funzione di attivazione -> probabilità pixel per pixel
    pred_mask = (prob_mask > 0.5).float()
    pred_mask = pred_mask.permute(0, 3, 1, 2) #  "NCHW" (dove N è la dimensione del batch, C è il numero di canali, H è l'altezza e W è la larghezza)
    mask = mask.permute(0, 3, 1, 2)


    tp, fp, fn, tn = smp.metrics.get_stats(pred_mask.long(), mask.long(), mode="binary")

    if stage == "train":
        self.log_dict(
            {
              "train/batch-IOU-img" : smp.metrics.iou_score(tp, fp, fn, tn, reduction="macro-imagewise"),
              "train/batch-IOU" : smp.metrics.iou_score(tp, fp, fn, tn, reduction="macro")
            },
            prog_bar=True,
            batch_size=config["batch_size"]
        )
    elif stage == "test":
        self.log_dict(
            {
              "test/IoU-weighted" : smp.metrics.iou_score(tp, fp, fn, tn, reduction='weighted', class_weights=[0.52, 16.67]),
              "test/mcc": self.mcc(pred_mask, mask)
            },
            on_step=False, on_epoch=True, prog_bar=True, batch_size=config["batch_size"]
        )

    return {
        "loss": loss,
        "tp": tp,
        "fp": fp,
        "fn": fn,
        "tn": tn,
    }

And this is the error I get Gcolab error

$\endgroup$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.