Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| from PIL import Image | |
| from torchvision import transforms | |
| #################################################################################################################### | |
| # Define your model and transform and all necessary helper functions here # | |
| # They will be imported to the exp_recognition.py file # | |
| #################################################################################################################### | |
| # Must match ImageFolder's alphabetical class_to_idx from the training notebook. | |
| classes = { | |
| 0: 'ANGER', | |
| 1: 'DISGUST', | |
| 2: 'FEAR', | |
| 3: 'HAPPINESS', | |
| 4: 'NEUTRAL', | |
| 5: 'SADNESS', | |
| 6: 'SURPRISE', | |
| } | |
| IMG_SIZE = 100 | |
| IMAGENET_MEAN = [0.485, 0.456, 0.406] | |
| IMAGENET_STD = [0.229, 0.224, 0.225] | |
| class facExpRec(nn.Module): | |
| """ResNet18 expression classifier trained in the expression notebook.""" | |
| def __init__(self, num_classes=len(classes)): | |
| super().__init__() | |
| self.backbone = models.resnet18(weights=None) | |
| in_features = self.backbone.fc.in_features | |
| self.backbone.fc = nn.Linear(in_features, num_classes) | |
| def forward(self, x): | |
| return self.backbone(x) | |
| def ensure_rgb(image): | |
| if isinstance(image, Image.Image): | |
| return image.convert('RGB') | |
| return Image.fromarray(image).convert('RGB') | |
| trnscm = transforms.Compose([ | |
| transforms.Lambda(ensure_rgb), | |
| transforms.Resize((IMG_SIZE, IMG_SIZE)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), | |
| ]) | |
| def _extract_state_dict(checkpoint): | |
| if not isinstance(checkpoint, dict): | |
| return checkpoint | |
| for key in ('state_dict', 'model_state_dict', 'net_dict'): | |
| if key in checkpoint: | |
| return checkpoint[key] | |
| return checkpoint | |
| def _normalize_state_dict_keys(state_dict): | |
| normalized = {} | |
| for key, value in state_dict.items(): | |
| if key.startswith('module.'): | |
| key = key[len('module.'):] | |
| if key.startswith('model.'): | |
| key = key[len('model.'):] | |
| normalized[key] = value | |
| return normalized | |
| def load_model(checkpoint_path, device, num_classes=len(classes)): | |
| model = facExpRec(num_classes=num_classes).to(device) | |
| checkpoint = torch.load(checkpoint_path, map_location=device) | |
| state_dict = _normalize_state_dict_keys(_extract_state_dict(checkpoint)) | |
| if any(key.startswith('backbone.') for key in state_dict): | |
| model.load_state_dict(state_dict, strict=True) | |
| else: | |
| model.backbone.load_state_dict(state_dict, strict=True) | |
| model.eval() | |
| return model | |