Segmentation masks not logging

Hi there,

I’m trying to log images and segmentation masks to Wandb, which I’ve had success with in the past. However it seems to have stopped working. The code I’m using (with pytorch-lightning framework) is:

image = Image(
      'ground_truth': {
          'mask_data': y_img,
          'class_labels': class_labels
      'predictions': {
          'mask_data': y_hat_img,
          'class_labels': class_labels
title = f'desc:{desc}:class:{j}:axis:{axis}'
self.logger.log_image(key=title, images=[image])

Any ideas what could be going wrong?

Example run: clarkbab/segmenter-miccai-convergence-part-3/runs/2ssrpujc


Hey @clarkbab, when you try to log images and segmentation masks, does it error? Or does it not log in the workspace? Any details about how your script stopped working would be super helpful!

Additionally, in the line

self.logger.log_image(key=title, images=[image])

How are you logging the images? Are you using WandbLogger from pytorch_lightning.loggers? Thank you for providing all these details!

Hi! There are no errors, it’s just failing silently. Yes that’s correct, I’m using WandbLogger from pytorch_lightning.loggers.

I am passing x_img, y_img and y_hat_img as numpy arrays and they have the same size and correct datatypes (float32 for x_img, and bool for others).

Hi @uma-wandb,

I’ve found the problem here, it seems that pytorch_lightning wraps the wandb.Image in another wandb.Image (see here) which doesn’t preserve the ‘masks’ keyword of the original image. Is this expected behaviour, or should wandb.Image preserve the ‘masks’ keyword?

I’ve added some code to reproduce:

import numpy as np
import wandb

# Create image and mask.
image = np.random.normal(size=(100, 100))
mask = np.random.choice([True, False], size=(100, 100))

# Create wandb image.
wandb_image = wandb.Image(image, masks={'mask': { 'mask_data': mask, 'class_labels': {0: 'background', 1: 'foreground'}}})

# This logs both image and mask...
wandb.init(project='test', name='test')
wandb.log({'test': [wandb_image]})

# Create wrapped image and mask.
wrapped_image = wandb.Image(wandb_image)

# This logs the image only...
wandb.log({'test': [wrapped_image]})


@clarkba Yes this behavior is expected. Thank you for sending the line this is referencing - in .log_image, we expect to receive the native image as well as the corresponding mask. When wrapping this in another wandb.Image, we override all the metadata associated with the image, masks included. For your workflow, I recommend continuing to log the native image + mask rather than a wandb.Image to avoid this problem.

Hi Brett, since we have not heard back from you we are going to close this request. If you would like to re-open the conversation, please let us know!

Thanks @uma-wandb,

That makes sense, I’m now using the pytorch-lightning ‘log_image’ method to log native image and masks instead of wandb.Image.