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