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

update testsingle.py.

Browse Source
This commit is contained in:
Brainway
2022-11-09 07:03:01 +00:00
committed by Gitee
parent d75f905648
commit 611db26c8a
1 changed files with 20 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
testsingle.py
View File

@ -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秒