Update ResNet_for_CC.py

ResNet_for_CC.py

@@ -2,92 +2,58 @@ import torch
 import torch.nn as nn
 import torchvision.models as models
 
+
 class ResClassifier(nn.Module):
-    """
-    A classifier with two fully connected layers followed by a final linear layer.
-    Uses BatchNorm, ReLU activations, and Dropout for better generalization.
-    """
-    def __init__(self, num_classes=14):
+    def __init__(self, class_num=14):
         super(ResClassifier, self).__init__()
-        
-        # First fully connected layer: reduces 128D features to 64D
         self.fc1 = nn.Sequential(
             nn.Linear(128, 64),
             nn.BatchNorm1d(64, affine=True),
             nn.ReLU(inplace=True),
             nn.Dropout()
         )
-        
-        # Second fully connected layer: retains 64D features
         self.fc2 = nn.Sequential(
             nn.Linear(64, 64),
             nn.BatchNorm1d(64, affine=True),
             nn.ReLU(inplace=True),
             nn.Dropout()
         )
-        
-        # Final classification layer mapping 64D features to class logits
-        self.fc3 = nn.Linear(64, num_classes)
+        self.fc3 = nn.Linear(64, class_num)
 
     def forward(self, x):
-        """
-        Forward pass through the classifier.
-        Returns class logits after two hidden layers.
-        """
-        x = self.fc1(x)  # First FC layer
-        x = self.fc2(x)  # Second FC layer
-        output = self.fc3(x)  # Final classification layer
-        return output
-
+        fc1_emb = self.fc1(x)
+        fc2_emb = self.fc2(fc1_emb)
+        logit = self.fc3(fc2_emb)
+        return logit
 
 class CC_model(nn.Module):
-    """
-    Clothing Classification Model based on ResNet50.
-    Extracts deep features and uses two independent classifiers for predictions.
-    """
     def __init__(self, num_classes1=14, num_classes2=None):
+
+        if num_classes2 is None:
+            num_classes2 = num_classes1
+
         super(CC_model, self).__init__()
-        
-        # If num_classes2 is not specified, default to num_classes1
-        num_classes2 = num_classes2 if num_classes2 else num_classes1
-        assert num_classes1 == num_classes2  # Ensure both classifiers predict the same categories
-        
+        assert num_classes1 == num_classes2
         self.num_classes = num_classes1
-        
-        # Load a pretrained ResNet-50 model as the feature extractor
         self.model_resnet = models.resnet50(weights='ResNet50_Weights.DEFAULT')
-        
-        # Remove ResNet's original classification layer to use as a feature extractor
         num_ftrs = self.model_resnet.fc.in_features
-        self.model_resnet.fc = nn.Identity()  # Identity layer keeps feature dimensions
-
-        # Additional transformation layer reducing feature size to 128D
+        self.model_resnet.fc = nn.Identity()
+        self.classification_fc = nn.Linear(num_ftrs, num_classes1)
         self.dr = nn.Linear(num_ftrs, 128)
-
-        # Two independent classifiers
         self.fc1 = ResClassifier(num_classes1)
         self.fc2 = ResClassifier(num_classes1)
 
     def forward(self, x, detach_feature=False):
-        """
-        Forward pass through the model.
-        Extracts deep features from ResNet and processes them through classifiers.
-        """
+
         with torch.no_grad():
-            # Extract deep features using ResNet-50 (without its original classification head)
             feature = self.model_resnet(x)
-
-        # Generate transformed features (128D) using the custom linear layer
-        dr_feature = self.dr(feature)
-
-        if detach_feature:
-            dr_feature = dr_feature.detach()  # Detach feature for non-trainable forward pass
-
-        # Pass features through two independent classifiers
-        out1 = self.fc1(dr_feature)
-        out2 = self.fc2(dr_feature)
-
-        # Compute the mean prediction from both classifiers
-        output_mean = (out1 + out2) / 2
-
-        return dr_feature, output_mean  # Returning feature embeddings and final prediction
+            res_out = self.classification_fc(feature)
+            if detach_feature:
+                feature = feature.detach()
+            dr_feature = self.dr(feature)
+            out1 = self.fc1(dr_feature)
+            out2 = self.fc2(dr_feature)
+            output_mean = (out1 + out2) / 2
+        return dr_feature, output_mean
+
+        #return dr_feature