I am working on a image classification problem where the images of grains are captured in a controlled environment where we can set the light intensity and background. It is a 3 way classification problem. The distinguishing feature between Class-A and Class-B is color while between Class-B and Class-C is size. Shape of all the three class objects is similar.
I used a fixed color background and lighting to obtain the training set. I have 1200 images of each category.
I trained using Resnet34 with transfer learning with accuracy of 96%. However I find that the heatmap is inconsistent i.e for some images it points to background. I suspect that NN is not able to learn the correct features and is instead overfitting to some features in the background.
Is there any combination of augmentations or other techniques which can avoid this overfitting and force the NN to learn features from the grain rather than the background??