From e558921a6bb03b3a67dddc5d1c4dabf014a91bde Mon Sep 17 00:00:00 2001
From: Brainway <brainway@protonmail.com>
Date: Wed, 26 Oct 2022 01:32:33 +0000
Subject: [PATCH] update predict.py.

---
 predict.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/predict.py b/predict.py
index caafcfb..8dc946d 100755
--- a/predict.py
+++ b/predict.py
@@ -32,7 +32,7 @@ class Predictor(object):
         self.num_classes = 0
         self.model = None
         self.prepare_model()
-        self.test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEAR),
+        self.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])])
 
@@ -41,7 +41,7 @@ class Predictor(object):
         config.split = self.args.split
         config.slide_step = self.args.slide_step
         self.num_classes = 5
-        self.cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
+        self.cls_dict = {0: "noemp", 1: "yesemp"}
 
         self.model = VisionTransformer(config, self.args.img_size, zero_head=True, num_classes=self.num_classes, smoothing_value=self.args.smoothing_value)
         
@@ -79,10 +79,15 @@ if __name__ == "__main__":
 
     y_true = []
     y_pred = []
-    # test_dir = "./emptyJudge5/images/"
+    #test_dir = "./emptyJudge5/images/"
     # dir_dict = {"noemp":"0", "yesemp":"1", "hard": "2", "fly": "3", "stack": "4"}
-    test_dir = "../emptyJudge2/"
-    dir_dict = {"noempty":"0", "empty":"1"}
+    
+    # test_dir = "../emptyJudge2/images"
+    # dir_dict = {"noempty":"0", "empty":"1"}
+
+    test_dir = "../emptyJudge5/images"
+    dir_dict = {"noemp":"0", "yesemp":"1"}
+
     total = 0
     num = 0
     t0 = time.time()