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

update lightrise.py.

Browse Source
This commit is contained in:
Brainway
2022-11-09 03:17:26 +00:00
committed by Gitee
parent a07331a414
commit f5dc4b6af4
1 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
lightrise.py
View File

@ -15,11 +15,6 @@ 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.device = torch.device("cpu")
# 准备模型
model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
model.to(args.device)
model.eval()
num_classes = 2 num_classes = 2
cls_dict = {0: "noemp", 1: "yesemp"} cls_dict = {0: "noemp", 1: "yesemp"}
@ -30,6 +25,10 @@ test_transform = transforms.Compose([transforms.Resize((600, 600), Image.BILINEA
#复杂场景小模型测试单张图片 #复杂场景小模型测试单张图片
def riseempty(imgdata): def riseempty(imgdata):
# 准备模型
model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
model.to(args.device)
model.eval()
x = test_transform(imgdata) x = test_transform(imgdata)
part_logits = model(x.unsqueeze(0).to(args.device)) part_logits = model(x.unsqueeze(0).to(args.device))
probs = torch.nn.Softmax(dim=-1)(part_logits) probs = torch.nn.Softmax(dim=-1)(part_logits)