Spaces:
Sleeping
Sleeping
File size: 2,745 Bytes
a47c171 | 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 | 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
|