From 1d4997bd42892b9c8764c9fadd41cf2c0a5745bb Mon Sep 17 00:00:00 2001
From: Brainway <brainway@protonmail.com>
Date: Tue, 27 Sep 2022 02:30:54 +0000
Subject: [PATCH] update predict.py.

---
 predict.py | 50 ++++++++++++++++++++++++++------------------------
 1 file changed, 26 insertions(+), 24 deletions(-)

diff --git a/predict.py b/predict.py
index f14096f..a9612f2 100755
--- a/predict.py
+++ b/predict.py
@@ -10,14 +10,14 @@ from PIL import Image
 from torchvision import transforms
 from models.modeling import VisionTransformer, CONFIGS
 
-
+#模型预测
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--img_size", default=448, type=int, help="Resolution size")
+    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")
     return parser.parse_args()
 
 
@@ -32,7 +32,7 @@ class Predictor(object):
         self.num_classes = 0
         self.model = None
         self.prepare_model()
-        self.test_transform = transforms.Compose([transforms.Resize((448, 448), Image.BILINEAR),
+        self.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])])
 
@@ -40,28 +40,18 @@ class Predictor(object):
         config = CONFIGS["ViT-B_16"]
         config.split = self.args.split
         config.slide_step = self.args.slide_step
-        model_name = os.path.basename(self.args.pretrained_model).replace("_checkpoint.bin", "")
-        print("use model_name: ", model_name)
-        if model_name.lower() == "emptyJudge5".lower():
-            self.num_classes = 5
-            self.cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
-        elif model_name.lower() == "emptyJudge4".lower():
-            self.num_classes = 4
-            self.cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "stack"}
-        elif model_name.lower() == "emptyJudge3".lower():
-            self.num_classes = 3
-            self.cls_dict = {0: "noemp", 1: "yesemp", 2: "hard"}
-        elif model_name.lower() == "emptyJudge2".lower():
-            self.num_classes = 2
-            self.cls_dict = {0: "noemp", 1: "yesemp"}
+        self.num_classes = 5
+        self.cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
+
         self.model = VisionTransformer(config, self.args.img_size, zero_head=True, num_classes=self.num_classes, smoothing_value=self.args.smoothing_value)
+        
         if self.args.pretrained_model is not None:
             if not torch.cuda.is_available():
-                pretrained_model = torch.load(self.args.pretrained_model, map_location=torch.device('cpu'))['model']
-                self.model.load_state_dict(pretrained_model)
+                self.model = torch.load(self.args.pretrained_model)
             else:
-                pretrained_model = torch.load(self.args.pretrained_model)['model']
-                self.model.load_state_dict(pretrained_model)
+                self.model = torch.load(self.args.pretrained_model,map_location='cpu')
+                
+        
         self.model.to(self.args.device)
         self.model.eval()
 
@@ -89,7 +79,7 @@ if __name__ == "__main__":
 
     y_true = []
     y_pred = []
-    test_dir = "/data/pfc/fineGrained/test_5cls"
+    test_dir = "./emptyJudge5/images/"
     dir_dict = {"noemp":"0", "yesemp":"1", "hard": "2", "fly": "3", "stack": "4"}
     total = 0
     num = 0
@@ -125,6 +115,18 @@ if __name__ == "__main__":
     print(rst_C)
     print(rst_f1)
 
+'''
+所有数据集
+
+The cast of time is :160.738966 seconds
+The classification accuracy is 0.986836
+[[4923   58]
+ [  34 1974]]
+0.9839851634589902
+
+'''
+
+
 '''
 test_imgs: yesemp=145, noemp=453  大图