update test_single.py.

test_single.py

@@ -6,55 +6,66 @@ from PIL import Image
 from torchvision import transforms
 import argparse
 from models.modeling import VisionTransformer, CONFIGS
+import time
 
 
+
+#模型测试单张图片
 parser = argparse.ArgumentParser()
-parser.add_argument("--dataset", choices=["CUB_200_2011", "emptyJudge5", "emptyJudge4"], default="emptyJudge5", help="Which dataset.")
-parser.add_argument("--img_size", default=448, type=int, help="Resolution size")
+parser.add_argument("--dataset", choices=["emptyJudge5"], default="emptyJudge5", help="Which dataset.")
+parser.add_argument("--img_size", default=600, type=int, help="Resolution size")
 parser.add_argument('--split', type=str, default='overlap', help="Split method")  # non-overlap
-parser.add_argument('--slide_step', type=int, default=12, help="Slide step for overlap split")
+parser.add_argument('--slide_step', type=int, default=2, help="Slide step for overlap split")
 parser.add_argument('--smoothing_value', type=float, default=0.0, help="Label smoothing value\n")
-parser.add_argument("--pretrained_model", type=str, default="output/emptyjudge5_checkpoint.bin", help="load pretrained model")
+parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_checkpoint_good.pth", help="load pretrained model")
+#parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_vit_checkpoint.pth", help="load pretrained model") #使用自定义VIT
 args = parser.parse_args()
 
 args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 args.nprocs = torch.cuda.device_count()
 
-# Prepare Model
+# 准备模型
 config = CONFIGS["ViT-B_16"]
 config.split = args.split
 config.slide_step = args.slide_step
 
-cls_dict = {}
-num_classes = 0
-if args.dataset == "emptyJudge5":
-    num_classes = 5
-    cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
-elif args.dataset == "emptyJudge4":
-    num_classes = 4
-    cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "stack"}
-elif args.dataset == "emptyJudge3":
-    num_classes = 3
-    cls_dict = {0: "noemp", 1: "yesemp", 2: "hard"}
-elif args.dataset == "emptyJudge2":
-    num_classes = 2
-    cls_dict = {0: "noemp", 1: "yesemp"}
+num_classes = 5
+cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
+
+
+
+
 model = VisionTransformer(config, args.img_size, zero_head=True, num_classes=num_classes, smoothing_value=args.smoothing_value)
+
 if args.pretrained_model is not None:
-    pretrained_model = torch.load(args.pretrained_model, map_location=torch.device('cpu'))['model']
-    model.load_state_dict(pretrained_model)
+    model = torch.load(args.pretrained_model) #自己预训练模型
 model.to(args.device)
 model.eval()
-# test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
-#                             transforms.CenterCrop((448, 448)),
-#                             transforms.ToTensor(),
-#                             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
-test_transform = transforms.Compose([transforms.Resize((448, 448), Image.BILINEAR),
+
+test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
+
+
+
+#自定义Vit模型
+# test_transform = transforms.Compose([transforms.Resize((320, 320), Image.BILINEAR),
+#                                      transforms.ToTensor(),
+#                                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
+
+
+
+#img = Image.open("img.jpg")
+
 img = Image.open("img.jpg")
+
+
 x = test_transform(img)
-part_logits = model(x.unsqueeze(0))
+
+
+startime = time.process_time()
+
+part_logits = model(x.unsqueeze(0).to(args.device))
 
 probs = torch.nn.Softmax(dim=-1)(part_logits)
 top5 = torch.argsort(probs, dim=-1, descending=True)
@@ -62,3 +73,7 @@ print("Prediction Label\n")
 for idx in top5[0, :5]:
     print(f'{probs[0, idx.item()]:.5f} : {cls_dict[idx.item()]}', end='\n')
 
+endtime = time.process_time()
+
+print("Time cost:"+ str(endtime - startime)) #评估一张图片耗时2.8秒
+