def predict_img(args):
checkpoint_path = args.checkpoint
if torch.cuda.is_available():
checkpoint = torch.load(checkpoint_path)
else:
checkpoint = torch.load(checkpoint_path,map_location='cpu')
net_name=checkpoint['pretrained_net']
if "dense" in net_name:
flag =1
else:
flag=2
model = fnc.load_checkpoint(checkpoint_path,flag)
if args.gpu:
device = 'cuda'
else:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
image_path = args.input
a=fnc.process_image(image_path)
a=a.to(device)
model.eval()
if args.top_k is None:
topk = 1
else:
topk = args.top_k
with torch.no_grad():
top_k_prob , idx=torch.exp(model(a)).topk(k=topk , dim=1)
idx_to_class=dict(map(reversed, model.class_to_idx.items()))
top_k_class = []
for i in idx[0]:
top_k_class.append(idx_to_class[i.item()])
top_k_prob=top_k_prob.cpu().numpy().ravel()
if args.category_names is None:
print(top_k_class)
print(list(map('{:.3f}'.format,top_k_prob)))
else:
with open(args.category_names, 'r') as f:
cat_to_name = json.load(f)
class_name = [cat_to_name[x] for x in top_k_class]
print(class_name)
print(list(map('{:.3f}'.format,top_k_prob)))
image = results.image_path
top_k = results.topk
gpu_mode = results.gpu
cat_names = results.cat_name_dir
with open(cat_names, 'r') as f:
cat_to_name = json.load(f)
# Establish model template
pre_tr_model = results.pretrained_model
model = getattr(models, pre_tr_model)(pretrained=True)
# Load model
loaded_model = load_checkpoint(model, save_dir, gpu_mode)
# Preprocess image - assumes jpeg format
processed_image = process_image(image)
if gpu_mode == True:
processed_image = processed_image.to('cuda')
else:
pass
# Carry out prediction
probs, classes = predict(processed_image, loaded_model, top_k, gpu_mode)
# Print probabilities and predicted classes
print(probs)
print(classes)
names = []
for i in classes: