Hello,
I am working on a project related to medical image processing using the UNet model architecture. I am here to train my UNet model to detect fractures in CT scans. My input data to UNet is (512,512,1) and the mask which is the segmentation. As the fracture/segmentation is too small to detect while training the model. I need to improve my F1 score. Any suggestion for loss function? I used focal loss, dice loss, and BCE but no improvement. The ratio of non-segmentation and segmentation is almost 1000:1 which means the class is imbalance. Can anyone please suggest some loss function or any new model?
When you say too small to detect, do you mean the fracture is too small in size, or that there isn’t enough training examples of fractures?
I don’t think you should be focusing on the loss function here. Rather, to improve the F1 score for an imbalanced data set, you should look at trying some of the following:
- If it’s an image, you can try performing image augmentation to increase the samples in the under-represented class. This is the easiest thing to try, and I’d recommend doing this first. Make sure that the image still looks reasonable after the augmentations you apply to it.
- Can you collect more data for the under-represented class?
- You can try different segmentation models to see if they perform better, but I wouldn’t do this first. It’s unlikely (but possible) you will see a massive improvement in performance just from changing the model.
- Are you fine-tuning an existing UNet or training it from scratch? If you’re doing it from scratch, I recommend fine-tuning instead. You can try to find an open-source model for image segmentation, especially one that’s already trained on similar medical images, and then fine-tuning on that model with your data instead.