Best way to approach binary segmentation with multilabel mask

Hi all. Is there a way in fastai v2 to ‘group’ labels using SegmentationDataLoaders so that it will fit a binary segmentation unet?
I have some masks that I wanted to test that seem like they were created using a gradiented masking tool, which has resulted in pixel values ranging between 0 to 255 rather than being either 0 or 255. Rather than recreating the masks, I was wondering if there was a way to edit the masks using DataLoaders in Fastai V2.

I am aware that usually if you have values of 0 or 255 rather than 0 or 1 you can edit the MaskBlock to include IntToFloatTensor(div_mask = 255) to divide by 255 resulting in values that are 0 or 1, like so:

def MaskBlock(codes=None):
"A TransformBlock for segmentation masks, potentially with codes"
return TransformBlock(type_tfms=PILMask.create, item_tfms=AddMaskCodes(codes=codes), batch_tfms=IntToFloatTensor(div_mask = 255))

However, when I add this line prior to my SegmentationDataLoaders, when I show_batch() the mask is missing. I was also messing around changing div_mask to 2, which on show_batch() results in the mask result I am looking for but crashes when I try to train the model as the number of codes doesn’t match with the number of pixel values available.

I also tried rewriting the IntToFloatTensor class but even running this class without changes anything results in the actual image to disappear on show_batch() - I assume it is affecting a similarly named class for the images.

I am assuming the best approach would be to convert eg. all values above 200 into ‘1’ and all values below into ‘0’ so that the masks would work in a binary segmentation model. I am aware you could do it using SegmentationItemList in V1 of Fastai, but what would be the best way to approach this in V2?

Thank you very much for your help!

I think the better approach is make your own custom block.
It is pretty straight forward, and probably you will need to make a custom Transform to format you masks.