How to avoid confusion between two classes

vision
1

I am using a pretrained HRNet model from timm for multi -class classification problem.My current dataset’s distribution is as follows

{'C1': 480, 'C2: 444, 'C3: 348, 'C4: 479, 'C5': 456, 'C6: 372, 'C7: 468, 'C8 396
And I am using weighted sample in the dataloader with weights as follows
[0.0021, 0.0023, 0.0029, 0.0021, 0.0022, 0.0027, 0.0021, 0.0025

My model classifier is as follows

(classifier): Sequential(
    (0): Linear(in_features=2048, out_features=625, bias=True)
    (1): ReLU()
    (2): Dropout(p=0.3, inplace=False)
    (3): Linear(in_features=625, out_features=256, bias=True)
    (4): ReLU()
    (5): Linear(in_features=256, out_features=128, bias=True)
    (6): ReLU()
    (7): Linear(in_features=128, out_features=8, bias=True)

I have noticed that when it comes to class2,class 4, class 6 and 8,the model ends up getting confused between these classes.
Images for class C2 and C4

C2
![model_44_3|500x500](upload://gHbVdyPDZFgpogwNQbjw40sQTLu.jpeg)

C4
![model_20_18|500x500](upload://uhS4xVOYPSXSmu9UMyR3ruv08l2.jpeg)

For example for the two images shown below

C6
![model_24_7|500x500](upload://15qkWoRfFDYsLkVCOlI0E62UJqO.jpeg)

C8
![model_14_8|500x500](upload://uDXEu7Q9FTZJAck2rBGJkJopa4n.jpeg)


When I use the model in inference stage the prediction for class 6 come out to be as follows
C6 score .33 C8 score .67
Prediction result for C8
C8 score .67 C6 score .33

I have trained the model using cross entropy loss,and I have used weighted sampling to handle class imbalance.

Any ideas on what I can do to ensure that the model is able to discrimante between C6 and C8 sufficently. Note this happens only for these four classes for the rest of the classes the confusion isn’t that significant