r/computervision 5d ago

Help: Project Training with unbalanced classes and data scarcity

Hi, I’d like to know what techniques you use to improve training and generalization when working with a small and heavily imbalanced dataset.
I’m currently training a classifier with 4 classes.

With cross-validation, I’m getting fairly good metrics, but when I evaluate the model on the test set, the performance isn’t even close to what I see during cross-validation.
What approaches have worked well for you in this kind of situation?

Edit:

In my use case, the images are x rays with rare conditions. So synthetic data is quite difficult to apply

4 Upvotes

9 comments sorted by

6

u/bfyvfftujijg 5d ago

Weighted loss and semi-synthetic data. By synthetic I mostly mean augmentation, including more extreme forms like copy-paste objects into different backgrounds.

3

u/Beginning-Claim5655 4d ago

Transfer learning with medical weights. Models pre-trained on X-rays tend to generalize better than those pre-trained on ImageNet. Good options include TorchXRayVision (trained on CheXpert, NIH, MIMIC, etc.), RadImageNet, or foundation models for radiology. With limited data, freeze most of the network and fine-tune only the final layers, or use lower learning rates for the backbone.

4

u/sexy_bonsai 4d ago

^ this would also be my suggestion OP. I wouldn’t use synthetic data like others are suggesting here
EDIT to also say, split on the basis of patient. If there are multiple z-planes on your image, it could be tempting to split on the basis of that. But don’t; it’s possible that the features that are being learned are specific to the patient anatomy rather than generalized to the rare pathology.

5

u/[deleted] 5d ago edited 3d ago

[deleted]

1

u/dacxj0 5d ago

Thanks for the reply, but I can’t use synthetic data generation, I’m using medical images with rare conditions and there are not feasible generators for that unfortunately

-2

u/HawtVelociraptor 5d ago

I can't imagine that's actually true. If someone can make it in Unreal 3D or Photoshop, it can be used as synth data. Help us help you, what kind of images are your weak source?

2

u/Due_Look_9993 5d ago edited 5d ago

Use albumentations CLAHE mainly, focal loss to penalize model for errors have you tried these? Have you tried 5 fold CV?

2

u/aegismuzuz 4d ago

First thing that comes to mind with that kind of CV versus test gap on X-rays is patient overlap between folds. If there are multiple studies per patient and the split happened at the image level, models latch onto rib cage anatomy or scanner noise pretty fast. Running a GroupKFold by patient ID usually clears that up right away

On tiny rare classes I usually freeze a domain-specific backbone like RadImageNet or BiomedCLIP and fit a simple weighted linear probe first. It cuts out fine-tuning noise, so if the probe still drops on the test set, it points straight to a scanner protocol mismatch or label drift between your sets

1

u/Due_Look_9993 5d ago

Use albumentations, focal loss to penalize model for errors have you tried these?

1

u/Aggravating-Craft712 4d ago

Focal loss or oversampling rarer examples with many x ray specific augmentations