Spaces:
Configuration error
Configuration error
File size: 3,465 Bytes
6ec35ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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):
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)
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
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):
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
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.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
|