Segmentation model where result must be one contiguous shape

I used the lesson3-camvid code to create a model that identifies the left ventricle in a cardiac ultrasound. One problem is that the model sometimes segments multiple parts of the heart and labels them as the left ventricle (see top row in the image below). How can I discourage this behavior and have the model constrained to predict one contiguous shape?

Here’s the code
codes = [‘Void’, ‘Ventricle’]

src = (SegmentationItemList.from_folder(path_img)
   .split_by_fname_file(path/'valids.txt')
   .label_from_func(get_y_fn, classes=codes))

train_tfms = [
    crop_pad(row_pct=(0,1), size=100, col_pct=(0,1), padding_mode='zeros'),
    symmetric_warp(magnitude=(-0.15,0.15), p=0.5),
    brightness(change=(0.3, 0.7), p=0.5),
    contrast(scale=(0.8,1.25), p=0.75)
]
valid_tfms = [
    crop_pad(row_pct=(0,1), size=100, col_pct=(0,1), padding_mode='zeros')
]

tfms = (train_tfms, valid_tfms)
data = (src.transform(tfms= tfms, tfm_y=True, size=112, padding_mode='zeros')
    .databunch(bs=bs)
    .normalize())

name2id = {v:k for k,v in enumerate(codes)}
void_code = name2id['Void']

def acc_camvid(input, target):
    target = target.squeeze(1)
    mask = target != void_code
    return (input.argmax(dim=1)[mask]==target[mask]).float().mean()

metrics=acc_camvid

learn = unet_learner(data, models.resnet34, metrics=metrics, wd=wd)

I think you may have to take a step back and determine what is most important. I think this is an example of the model having a hard time with labeling this image. What is most important to you? That the model label part of the image correct? Or the model not label things incorrectly?

I think you will have to weight void loss as more important so that the model assumes void. Otherwise I think the step after that is to use some kind of heuristic to guarantee that the visualization is correct.

In this one the model should have already had a high loss for the image you show, so it is already “discouraged”

1 Like