Add application file

core.py

@@ -16,27 +16,15 @@ def load_model(pretrained_dict, new):
     new.load_state_dict(model_dict)
 
 
-# if torch.cuda.is_available():
-#     device = torch.device("cuda")
-# else:
 device = torch.device("cpu")
 print("use cpu")
 
-# model_ckpt_path = "./models/resnet18.pth"
-# model_ckpt_path = "https://huggingface.co/M4869/beauty_prediction_fpb5k/blob/main/resnet18.pth"
-
-# model_ckpt_path = "M4869/beauty_prediction_fpb5k"
-# model = torch.hub.load("huggingface/transformers", model_ckpt_path)
-
-
-
-
 # ans = hf_hub_download(repo_id="google/pegasus-xsum", filename="config.json")
 # /home/my/.cache/huggingface/hub/models--google--pegasus-xsum/snapshots/8d8ffc158a3bee9fbb03afacdfc347c823c5ec8b/config.json
 
 model_ckpt_path = hf_hub_download(repo_id="M4869/beauty_prediction_fpb5k", filename="resnet18.pth")
 net = Nets.ResNet(block=Nets.BasicBlock, layers=[2, 2, 2, 2], num_classes=1).to(device)
-load_model(torch.load(model_ckpt_path, encoding='latin1'), net)
+load_model(torch.load(model_ckpt_path, encoding='latin1', map_location=torch.device('cpu')), net)
 net.eval()
 
 transform = transforms.Compose([
@@ -53,17 +41,3 @@ def fun(img_path):
         img = img.unsqueeze(0).to(device)
         output = net(img).squeeze(1).cpu().numpy()[0]
         return output
-
-# def main():
-#     for i in range(6, 7):
-#         img = Image.open("./data2/%d.jpg" % i).convert('RGB')
-#         img = transform(img)
-#
-#         with torch.no_grad():
-#             img = img.unsqueeze(0).to(device)
-#             output = net(img).squeeze(1).cpu().numpy()[0]
-#             print(i, output * 20)
-
-
-# if __name__ == '__main__':
-# main()