DataBlock for segmentation

I am working on a segmentation task, where my original images and masked images are in separate folder namely images and segmentation. I tried to build the DataBlock with the code,

tree = DataBlock(blocks = (ImageBlock() , ImageBlock()), 
                 get_items = get_image_files , 
                 splitter = RandomSplitter(valid_pct= 0.2), 
                 item_tfms = Resize(224), 
                 get_x = get_img,
                 get_y = get_seg)

tree.summary('data/train')

This is how I built my learner and dls:

# Dataloaders 
dls = tree.dataloaders('data/train/' , bs=8)
dls.c = 1

# Learner 
learn = unet_learner(dls, resnet18 , n_out= dls.c , loss_func= nn.MSELoss() , metrics = dice)

The loss function I am using is Dice but when I fit the model, I am getting an assertion error.

I am not sure what I am doing wrong, can anyone help me with this?

If you’re doing segmentation, should you not use the MaskBlock? Why are you using the ImageBlock in this instance?

Some segmentation resources:

1 Like

Because I have the segmented images in a separate folder itself. I can only use MaskBlock() if I have a mask function right?
I am not sure sorry!

This is a bit of a guess, so my explanation may not be 100%

But I think @muellerzr is right that is the fact you are using ImageBlock and not MaskBlock

The check is making sure the predicted mask is the same size as your labels
(Assuming that your images are colour) As you have passed in your labels as images with the ImageBlock, they still have the 3 colour channels.

So when fastai flattens them to check that they are the same size as your 1 channel predicted mask it is 3 times larger than it should be

1204224 / 401408 = 3

This might be wrong, so I would suggest you double check your self - using the ipython debug to confirm the types/ shapes of the variables should give you the answer.

Oh and just a little extra thing, your loss function isn’t Dice. It’s nn.MSELoss()
Your metric is Dice, which won’t effect how your model trains, just what get printed out

I think Dice loss has been added in the latest update, but don’t quote me!

2 Likes

Ok I think (not totally sure) that if your segmented images are in a separate folder, you need a label_func which can map from the input image path to the output image path.

def label_func(o):
     """ takes in a file as found from `get_image_files`
         and maps to the path name of the segmented image 
     """
     return Path("segmented_images") / o.name # or whatever is right for your directory setup. 

DataBlock(blocks=(ImageBlock, ImageBlock),
        splitter=RandomSplitter(0.2, seed=42),
        get_items=get_image_files,
        get_y=label_func,)

Hey @idraja yup I did that everything works fine, but I get an error when I use Dice metric.

Not sure how to solve this.

Hey @lukemshepherd yup it Dice is a metric sorry about that.

If I didnt use this metric my model runs smoothly, but using Dice as a metric bringing some assertion errors.

This is how I built my learner,

learn = unet_learner(dls= tree , arch = resnet18 , loss_func= loss_gen , norm_type= NormType.Weight , 
                     self_attention = True , y_range = y_range , metrics = [Dice)])

Its all good - I only mention that because I’ve made that mistake and waisted a couple a hours training with the wrong loss function!

Where is your dice metic from?

Tbh you were right, the assertion error was because I was using ImageBlock for my target block but the dice function expects a TensorMask (1D).

I got the metric from here Metrics | fastai

Trying out different things, but the pain point is I have my segmented images in a separate folder. Thanks for your answer, it really helps and now I was able to figure out where I was wrong.

You should be able to have masks and images in different folders. Have a read of walkwithfastai lesson 4

I use something like (taken from the lesson 4):
get_y_fn = lambda x: path/'masks'/f'{x.stem}_mask{x.suffix}'

then my data block

and finally:
dls = block.dataloaders(path/'images', bs = bs path=path)

1 Like

There are several steps that need to be followed in order to create data blocks.

The steps are defined by the data block API. They can be asked in the form of questions while looking at the data:

  1. What is the types of your inputs/targets?
  2. Where is your data?
  3. Does something need to be applied to inputs?
  4. Does something need to be applied to the target?
  5. How to split the data?
  6. Do we need to apply something on formed items?
  7. Do we need to apply something on formed batches?
    This is it!!

dgcustomerfirst