update predict.py.

This commit is contained in:
Brainway
2022-11-01 02:48:14 +00:00
committed by Gitee
parent 4e640531da
commit 9363d22815

View File

@ -17,8 +17,8 @@ def parse_args():
parser.add_argument('--split', type=str, default='overlap', help="Split method") # non-overlap
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/ieemooempty_checkpoint_good.pth", help="load pretrained model")
parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/now/emptyjudge5_checkpoint.bin", help="load pretrained model")
parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_vit_checkpoint.pth", help="load pretrained model")
#parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/now/emptyjudge5_checkpoint.bin", help="load pretrained model")
#parser.add_argument("--pretrained_model", type=str, default="output/ieemooempty_vitgood_checkpoint.pth", help="load pretrained model")
return parser.parse_args()
@ -34,7 +34,7 @@ class Predictor(object):
self.num_classes = 0
self.model = None
self.prepare_model()
self.test_transform = transforms.Compose([transforms.Resize((320, 320), 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])])