Update inference.py

inference.py

@@ -1,20 +1,21 @@
 import torch
-import torchvision.transforms as transforms
+import torchvision.models as models
+from torchvision import transforms
 from PIL import Image
-import pickle
 
-# load model
-model = pickle.load(open("resnet50_c3_lr3e-04_bs32_aug_heavy_opt_adam_dr.pkl", "rb"))
+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()
 ])
 
-labels = ["angry","happy","sad","fear","surprise","neutral"]
-
-
 def predict(image):
     image = transform(image).unsqueeze(0)