Create databunch with multiple segmentation mask as label

I got past this stage with the following code:

get_y_fn_fish = lambda x: train['EncodedPixels'].loc[train['Image_Label'].str.contains(x.as_posix().split('/')[-1]+"*",regex=True)].values[0] if train['EncodedPixels'].loc[train['Image_Label'].str.contains(x.as_posix().split('/')[-1]+"*",regex=True)].values[0] == train['EncodedPixels'].loc[train['Image_Label'].str.contains(x.as_posix().split('/')[-1]+"*",regex=True)].values[0] else ''

so that function returns the rle string or ‘’ corresponding to an image. After that, if you want to open the mask you can do (img_f is a PosixPath object):

mask_fish = open_mask_rle(get_y_fn_fish(img_f), shape=(1400, 2100)).resize((1,128,128))
mask_fish

I then proceeded to convert the rle mask into an image and saved them:

for image in path_img.iterdir():
    mask_fish = open_mask_rle(get_y_fn_fish(image), shape=(2100, 1400))
    mask_fish.save(path_img/'..'/'fish_masks'/f'{image.stem}_fish_mask.png')

After that, we can create a function that maps an actual image to its image mask:

get_y_fn_f = lambda x: path_img/'..'/'fish_masks'/f'{x.stem}_fish_mask.png'

we need another modification to make this work with masks with values 0 and 1:

and voila:

   src = (SegItemListCustom.from_folder(path_img)
       .split_by_rand_pct(0.2)
       .label_from_func(get_y_fn_f,classes=[0,1]))

data = (src.transform(tfms, size=128)
        .databunch().normalize(imagenet_stats))
1 Like