lichen/ieemoo-ai-isempty
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update predict.py.

Browse Source
This commit is contained in:
Brainway
2022-10-26 01:32:33 +00:00
committed by Gitee
parent 5a8c6a5d2e
commit e558921a6b
1 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
predict.py
View File

@ -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()