Improvements/alternatives to U-net for medical images segmentation?

Submitted by viertys t3_125ximj in deeplearning

Hello,

I am working on a project in which I'm detecting cavities in X-rays.

The dataset I have is pretty limited (~100 images). Each X-ray has a black and white mask that shows where in the image are the cavities.

I'm trying to improve my results.

What I've tried so far:

  1. different loss functions: BCE, dice loss, bce+dice, tversky loss, focal tversky loss
  2. modifying the images' gamma to make the cavities more visible
  3. trying out different U-Nets: U-net, V-net, U-net++, UNET 3+, Attention U-net, R2U-net, ResUnet-a, U^2-Net, TransUNET, and Swin-UNET

None of the new U-nets that I've tried improved the results. Probably because they are more suited for a larger dataset.

I'm now looking for other things to try to improve my results. Currently my network is detecting cavities, but it has trouble with the smaller ones.

2

Comments

You must log in or register to comment.

Seahorsejockey t1_je6m3oa wrote

How Big are your images (resolution HxW)?

1

viertys OP t1_je6oxpk wrote

512x512, but I can modify their dimensions

1

trajo123 t1_je79rme wrote

Have you tried using the segmentation models from the SMP package (Iakubovskii, P. (2019)? I built a segmentation model for dermoscopy images and pre-trained models consistently outperformed anything else, architecture didn't matter that much. I got best results with "U-Net with SegFormer pre-trained encoder".

It depends how much training data you have, but unless you have millions of samples, pre-training usually trumps architecture.

1

Environmental_Ice422 t1_je8fe4r wrote

I would suggest you to apply more dramatic image augmentation methods

1

viertys OP t1_je9lv4x wrote

I am currently using the albumentations module. I rotate, shift, rotate, blur, horizontal flip, downscale and use gauss noise. I get around 400 images after doing this. Is there anything you would suggest?

1

Environmental_Ice422 t1_jed1cqb wrote

You should apply those transforms on each batch while training, rather than transforming the data before training. This approach is called doing the augmentation on the fly.

1

Yeinstein20 t1_je8trp8 wrote

Your dataset is rather small and it seems you are not really doing Augmentations? I would try different Augmentations, that should improve your results regardless of the used model. Have you looked at some frameworks for medical image segmentation? nnUNet comes to mind which would give you a solid baseline. How good are your results currently?

1

viertys OP t1_je9m4ao wrote

I didn't mention it in the post, but I'm using the albumentations module. I rotate, shift, rotate, blur, horizontal flip, downscale and use gauss noise. I get around 400 images after doing this. Is there anything you would suggest?

I have an accuracy of 98.50 and I have dice of around 0.30-0.65 in each image

1