def predict(model, loader, device):
model.eval()
predictions = np.zeros((len(loader.dataset), NUM_PTS, 2))
for i, batch in enumerate(tqdm.tqdm(loader, total=len(loader), desc='test prediction...')):
images = batch['image'].to(device)
with torch.no_grad():
pred_landmarks = model(images).cpu()
pred_landmarks = pred_landmarks.numpy().reshape((len(pred_landmarks), NUM_PTS, 2)) # B x NUM_PTS x 2
fs = batch['scale_coef'].numpy() # B
margins_x = batch['crop_margin_x'].numpy() # B
margins_y = batch['crop_margin_y'].numpy() # B
prediction = restore_landmarks_batch(pred_landmarks, fs, margins_x, margins_y) # B x NUM_PTS x 2
predictions[i * loader.batch_size: (i + 1) * loader.batch_size] = prediction
return predictions
def predict(model, loader, device):
model.eval()
predictions = np.zeros((len(loader.dataset), NUM_PTS, 2))
with tqdm.tqdm(total=len(loader), position=0, leave=True) as pbar:
for i, batch in enumerate(tqdm.tqdm(loader, total=len(loader), desc="test prediction...", position=0, leave=True)):
images = batch["image"].to(device)
with torch.no_grad():
pred_landmarks = model(images).cpu()
pred_landmarks = pred_landmarks.numpy().reshape((len(pred_landmarks), NUM_PTS, 2)) # B x NUM_PTS x 2
fs = batch["scale_coef"].numpy() # B
margins_x = batch["crop_margin_x"].numpy() # B
margins_y = batch["crop_margin_y"].numpy() # B
prediction = restore_landmarks_batch(pred_landmarks, fs, margins_x, margins_y) # B x NUM_PTS x 2
predictions[i * loader.batch_size: (i + 1) * loader.batch_size] = prediction
pbar.update()
return predictions