| import torch |
| import torchvision.models as models |
| from torchvision import transforms |
| from PIL import Image |
|
|
| model = models.resnet50() |
| model.fc = torch.nn.Linear(model.fc.in_features, 6) |
| model_path = "resnet50_c3_lr3e-04_bs32_aug_heavy_opt_adam_drop0.5_ls0.1_6class.pth" |
| model.load_state_dict(torch.load(model_path, map_location="cpu")) |
| model.eval() |
|
|
| labels = ["angry","fear","happy","sad","surprise","neutral"] |
|
|
| transform = transforms.Compose([ |
| transforms.Resize((224,224)), |
| transforms.ToTensor() |
| ]) |
|
|
| def predict(image): |
| image = transform(image).unsqueeze(0) |
|
|
| with torch.no_grad(): |
| output = model(image) |
| pred = torch.argmax(output,1).item() |
|
|
| return {"emotion": labels[pred]} |
