diff --git a/testsingle.py b/testsingle.py index 38b3a52..2bd6167 100755 --- a/testsingle.py +++ b/testsingle.py @@ -7,7 +7,7 @@ from torchvision import transforms import argparse from models.modeling import VisionTransformer, CONFIGS import time - +import lightrise #模型测试单张图片 @@ -17,11 +17,12 @@ 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('--smoothing_value', type=float, default=0.0, help="Label smoothing value\n") -#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") #使用自定义VIT +parser.add_argument("--pretrained_model", type=str, default="../module/ieemoo-ai-isempty/model/new/ieemooempty_vit_checkpoint.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.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +args.device = torch.device("cpu") args.nprocs = torch.cuda.device_count() # 准备模型 @@ -55,7 +56,7 @@ test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEA #img = Image.open("img.jpg") -img = Image.open("img.jpg") +img = Image.open("light.jpg") x = test_transform(img) @@ -74,12 +75,24 @@ for idx in top5[0, :5]: clas_ids = top5[0][0] clas_ids = 0 if 0==int(clas_ids) or 2 == int(clas_ids) or 3 == int(clas_ids) else 1 print("cur_img result: class id: %d, score: %0.3f" % (clas_ids, probs[0, clas_ids].item())) + result={} result["success"] = "true" result["rst_cls"] = str(clas_ids) + +riseresult = lightrise.riseempty(Image.open("light.jpg")) +if(int(result["rst_cls"])==1): + if(int(riseresult["rst_cls"])==1): + result = {} + result["success"] = "true" + result["rst_cls"] = 1 + else: + result = {} + result["success"] = "true" + result["rst_cls"] = 0 + print(result) endtime = time.process_time() -print("Time cost:"+ str(endtime - startime)) #评估一张图片耗时2.8秒 - +print("Time cost:"+ str(endtime - startime)) #评估一张图片耗时2.8秒 \ No newline at end of file