98 lines
3.3 KiB
Python
Executable File
98 lines
3.3 KiB
Python
Executable File
# coding=utf-8
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
import argparse
|
|
from models.modeling import VisionTransformer, CONFIGS
|
|
import time
|
|
import lightrise
|
|
|
|
|
|
#模型测试单张图片
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset", choices=["emptyJudge5"], default="emptyJudge5", help="Which dataset.")
|
|
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/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("cpu")
|
|
args.nprocs = torch.cuda.device_count()
|
|
|
|
# 准备模型
|
|
config = CONFIGS["ViT-B_16"]
|
|
config.split = args.split
|
|
config.slide_step = args.slide_step
|
|
|
|
num_classes = 5
|
|
cls_dict = {0: "noemp", 1: "yesemp", 2: "hard", 3: "fly", 4: "stack"}
|
|
|
|
model = None
|
|
#model = VisionTransformer(config, args.img_size, zero_head=True, num_classes=num_classes, smoothing_value=args.smoothing_value)
|
|
|
|
if args.pretrained_model is not None:
|
|
model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
|
|
model.to(args.device)
|
|
model.eval()
|
|
|
|
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])])
|
|
|
|
|
|
|
|
#自定义Vit模型
|
|
# 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])])
|
|
|
|
|
|
|
|
#img = Image.open("img.jpg")
|
|
|
|
img = Image.open("light.jpg")
|
|
|
|
|
|
x = test_transform(img)
|
|
|
|
|
|
startime = time.process_time()
|
|
|
|
part_logits = model(x.unsqueeze(0).to(args.device))
|
|
|
|
probs = torch.nn.Softmax(dim=-1)(part_logits)
|
|
top5 = torch.argsort(probs, dim=-1, descending=True)
|
|
print("Prediction Label\n")
|
|
for idx in top5[0, :5]:
|
|
print(f'{probs[0, idx.item()]:.5f} : {cls_dict[idx.item()]}', end='\n')
|
|
|
|
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秒 |