Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.models as models | |
| class CustomModel(nn.Module): | |
| def __init__(self, num_classes=4): | |
| super(CustomModel, self).__init__() | |
| self.efficientnet = models.efficientnet_v2_s(weights=models.EfficientNet_V2_S_Weights.IMAGENET1K_V1) | |
| # Get the number of features from the last layer of EfficientNetV2 | |
| num_features = self.efficientnet.classifier[1].in_features | |
| # Remove the classifier | |
| self.efficientnet = nn.Sequential(*list(self.efficientnet.children())[:-1]) | |
| self.gap = nn.AdaptiveAvgPool2d(1) | |
| self.fc1 = nn.Linear(num_features, 512) | |
| self.dropout1 = nn.Dropout(0.5) | |
| self.fc2 = nn.Linear(512, 256) | |
| self.dropout2 = nn.Dropout(0.3) | |
| self.fc3 = nn.Linear(256, num_classes) | |
| def forward(self, x): | |
| x = self.efficientnet(x) | |
| x = self.gap(x) | |
| x = torch.flatten(x, 1) | |
| x = F.relu(self.fc1(x)) | |
| x = self.dropout1(x) | |
| x = F.relu(self.fc2(x)) | |
| x = self.dropout2(x) | |
| x = self.fc3(x) | |
| return x | |
| def load_model(model_path): | |
| model = CustomModel() | |
| model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) | |
| model.eval() | |
| return model | |