def getLandmarksFromFrame(self,image,detected_faces):
if len(detected_faces) == 0:
return []
centers = []
scales = []
for i, d in enumerate(detected_faces):
center = torch.FloatTensor(
[d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
center[1] = center[1] - (d[3] - d[1]) * 0.12
scale = (d[2] - d[0] + d[3] - d[1]) / 195
centers.append(center)
scales.append(scale)
inp = crop(image, center, scale)
inp = torch.from_numpy(inp.transpose(
(2, 0, 1))).float()
inp = inp.to(self.device)
inp.div_(255.0).unsqueeze_(0)
if i == 0:
imgs = inp
else:
imgs = torch.cat((imgs,inp), dim=0)
out = self.forward(imgs)
out = out[-1].cpu()
pts, pts_img = get_preds_fromhm(out, centers, scales)
#pts, pts_img = pts.view(68, 2) * 4, pts_img.view(68, 2)
#landmarks.append(pts_img.numpy())
return pts_img.numpy().tolist()
def train_model(model,
train_dataloaders,
train_dataset_sizes,
valid_dataloaders,
valid_dataset_sizes,
batch_size,
optimizer,
criterion,
lr_scheduler,
writer_dict,
logger,
use_gpu=True,
epoches=5,
dataset='train',
save_path='./',
num_landmarks=68):
global_nme = 0
best_nme = 100
for epoch in range(epoches):
print("this is epoch {}".format(epoch))
model.train()
running_loss = 0
step = 0
total_nme = 0
total_count = 0
fail_count = 0
nmes = []
# running_corrects = 0
# Iterate over data.
for data in train_dataloaders[dataset]:
total_runtime = 0
run_count = 0
step_start = time.time()
step += 1
# get the inputs
inputs = data['image'].type(torch.FloatTensor)
labels_heatmap = data['heatmap'].type(torch.FloatTensor)
labels_boundary = data['boundary'].type(torch.FloatTensor)
landmarks = data['landmarks'].type(torch.FloatTensor)
loss_weight_map = data['weight_map'].type(torch.FloatTensor)
# wrap them in Variable
if use_gpu:
inputs = inputs.to(device)
labels_heatmap = labels_heatmap.to(device)
labels_boundary = labels_boundary.to(device)
loss_weight_map = loss_weight_map.to(device)
else:
inputs, labels_heatmap = Variable(inputs), Variable(
labels_heatmap)
labels_boundary = Variable(labels_boundary)
labels = torch.cat((labels_heatmap, labels_boundary), 1)
single_start = time.time()
outputs, boundary_channels = model(inputs)
##反向传播
loss = criterion(outputs[-1], labels, loss_weight_map)
optimizer.zero_grad()
loss.backward()
optimizer.step()
single_end = time.time()
total_runtime += single_end - single_start
run_count += 1
step_end = time.time()
for i in range(inputs.shape[0]):
#对batch的image处理
img = inputs[i]
img = img.cpu().numpy()
img = img.transpose((1, 2, 0)) * 255.0
img = img.astype(np.uint8)
img = Image.fromarray(img)
# pred_heatmap = outputs[-1][i].detach().cpu()[:-1, :, :] detach()和.data作用基本相同 都是防止反传改变梯度
pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0))
pred_landmarks = pred_landmarks.squeeze().numpy()
gt_landmarks = data['landmarks'][i].numpy()
if num_landmarks == 68:
left_eye = np.average(gt_landmarks[36:42], axis=0)
right_eye = np.average(gt_landmarks[42:48], axis=0)
norm_factor = np.linalg.norm(left_eye - right_eye)
# norm_factor = np.linalg.norm(gt_landmarks[36]- gt_landmarks[45])
elif num_landmarks == 98:
norm_factor = np.linalg.norm(gt_landmarks[60] -
gt_landmarks[72])
elif num_landmarks == 19:
left, top = gt_landmarks[-2, :]
right, bottom = gt_landmarks[-1, :]
norm_factor = math.sqrt(
abs(right - left) * abs(top - bottom))
gt_landmarks = gt_landmarks[:-2, :]
elif num_landmarks == 29:
# norm_factor = np.linalg.norm(gt_landmarks[8]- gt_landmarks[9])
norm_factor = np.linalg.norm(gt_landmarks[16] -
gt_landmarks[17])
single_nme = (np.sum(
np.linalg.norm(pred_landmarks * 4 - gt_landmarks, axis=1))
/ pred_landmarks.shape[0]) / norm_factor
nmes.append(single_nme)
total_count += 1
if single_nme > 0.1:
fail_count += 1
if step % 50 == 0:
# print('Step {} Time: {:.6f} Input Mean: {:.6f} Output Mean: {:.6f}'.format(
# step, step_end - step_start,
# torch.mean(labels),
# torch.mean(outputs[0])))
msg = "epoch {} Step [{} / {}] step Time: {:.6f} loss: {:.4f} ".format(
epoch, step,
int(train_dataset_sizes['train'] / batch_size),
step_end - step_start, loss.item())
logger.info(msg)
#################写入train_loss###################
if writer_dict:
writer = writer_dict["writer"]
global_steps = writer_dict["train_global_steps"]
writer.add_scalar("train_loss", loss.item(), global_steps)
writer_dict["train_global_steps"] += 1
# gt_landmarks = landmarks.numpy()
# pred_heatmap = outputs[-1].to('cpu').numpy()
gt_landmarks = landmarks
batch_nme = fan_NME(outputs[-1][:, :-1, :, :].detach().cpu(),
gt_landmarks, num_landmarks)
# batch_nme = 0
total_nme += batch_nme
lr_scheduler.step()
epoch_nme = total_nme / train_dataset_sizes['train']
global_nme += epoch_nme
#每个epoch检验一次
nme, predictions = eval_model(model, valid_dataloaders,
valid_dataset_sizes, writer_dict, logger)
is_best = epoch_nme < best_nme
best_nme = min(epoch_nme, best_nme)
#储存模型
save_checkpoint(
{
"state_dict": model.module.state_dict(),
"epoch": epoch + 1,
"best_nme": best_nme,
"optimizer": optimizer.state_dict(),
}, predictions, is_best, save_path,
'checkpoint_{}.pth'.format(epoch))
nme_save_path = os.path.join(save_path, 'nme_log.npy')
np.save(nme_save_path, np.array(nmes))
msg = 'epoch {} TRAIN NME: {:.6f} Failure Rate: {:.6f} Total Count: {:.6f} Fail Count: {:.6f}'.format(
epoch, epoch_nme, fail_count / total_count, total_count,
fail_count)
logger.info(msg)
print("is_best", is_best)
msg = 'training done! Average NME: {:.6f}'.format(global_nme / epoches)
logger.info(msg)
msg = 'Everage runtime for a single batch: {:.6f}'.format(total_runtime /
run_count)
logger.info(msg)
return model
def eval_model(model,
dataloaders,
dataset_sizes,
writer_dict,
logger,
use_gpu=True,
dataset='val',
save_path='./',
num_landmarks=68):
model.eval()
running_loss = 0
global_nme = 0
step = 0
total_nme = 0
total_count = 0
fail_count_008 = 0
fail_count_010 = 0
nmes = []
preds = np.zeros((dataset_sizes[dataset], num_landmarks, 2))
# running_corrects = 0
# Iterate over data.
with torch.no_grad():
for data in dataloaders[dataset]:
total_runtime = 0
run_count = 0
step_start = time.time()
step += 1
# get the inputs
inputs = data['image'].type(torch.FloatTensor)
labels_heatmap = data['heatmap'].type(torch.FloatTensor)
labels_boundary = data['boundary'].type(torch.FloatTensor)
landmarks = data['landmarks'].type(torch.FloatTensor)
loss_weight_map = data['weight_map'].type(torch.FloatTensor)
# wrap them in Variable
if use_gpu:
inputs = inputs.to(device)
labels_heatmap = labels_heatmap.to(device)
labels_boundary = labels_boundary.to(device)
loss_weight_map = loss_weight_map.to(device)
else:
inputs, labels_heatmap = Variable(inputs), Variable(
labels_heatmap)
labels_boundary = Variable(labels_boundary)
labels = torch.cat((labels_heatmap, labels_boundary), 1)
single_start = time.time()
outputs, boundary_channels = model(inputs)
single_end = time.time()
total_runtime += single_end - single_start
run_count += 1
step_end = time.time()
for i in range(inputs.shape[0]):
img = inputs[i]
img = img.cpu().numpy()
img = img.transpose((1, 2, 0)) * 255.0
img = img.astype(np.uint8)
img = Image.fromarray(img)
# pred_heatmap = outputs[-1][i].detach().cpu()[:-1, :, :] detach()和.data作用基本相同 都是防止反传改变梯度
pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
pred_landmarks, pred_landmarks_oripic = get_preds_fromhm(
pred_heatmap.unsqueeze(0))
pred_landmarks = pred_landmarks.squeeze().numpy()
preds[data['index'][i]] = pred_landmarks_oripic
gt_landmarks = data['landmarks'][i].numpy()
if num_landmarks == 68:
left_eye = np.average(gt_landmarks[36:42], axis=0)
right_eye = np.average(gt_landmarks[42:48], axis=0)
norm_factor = np.linalg.norm(left_eye - right_eye)
# norm_factor = np.linalg.norm(gt_landmarks[36]- gt_landmarks[45])
elif num_landmarks == 98:
norm_factor = np.linalg.norm(gt_landmarks[60] -
gt_landmarks[72])
elif num_landmarks == 19:
left, top = gt_landmarks[-2, :]
right, bottom = gt_landmarks[-1, :]
norm_factor = math.sqrt(
abs(right - left) * abs(top - bottom))
gt_landmarks = gt_landmarks[:-2, :]
elif num_landmarks == 29:
# norm_factor = np.linalg.norm(gt_landmarks[8]- gt_landmarks[9])
norm_factor = np.linalg.norm(gt_landmarks[16] -
gt_landmarks[17])
single_nme = (np.sum(
np.linalg.norm(pred_landmarks * 4 - gt_landmarks, axis=1))
/ pred_landmarks.shape[0]) / norm_factor
nmes.append(single_nme)
total_count += 1
if single_nme > 0.08:
fail_count_008 += 1
if single_nme > 0.1:
fail_count_010 += 1
# gt_landmarks = landmarks.numpy()
# pred_heatmap = outputs[-1].to('cpu').numpy()
gt_landmarks = landmarks
batch_nme = fan_NME(outputs[-1][:, :-1, :, :].detach().cpu(),
gt_landmarks, num_landmarks)
# batch_nme = 0
total_nme += batch_nme
global_nme = total_nme / dataset_sizes['val']
msg = 'val_NME: {:.6f} Failure( > 0.08) Rate: {:.6f} Failure( > 0.10) Rate: {:.6f} Total Count: {:.6f} '.format(
global_nme, fail_count_008 / total_count,
fail_count_010 / total_count, total_count)
logger.info(msg)
if writer_dict:
writer = writer_dict["writer"]
global_steps = writer_dict["valid_global_steps"]
writer.add_scalar("valid_nme", global_nme, global_steps)
writer_dict["valid_global_steps"] += 1
print('Evaluation done! Average NME: {:.6f}'.format(global_nme))
print('Everage runtime for a single batch: {:.6f}'.format(total_runtime /
run_count))
return global_nme, preds
def get_landmakrs(img_folder, lmks_folder, img_listfile=None):
if not os.path.exists(lmks_folder):
os.mkdir(lmks_folder)
imgs = []
if img_listfile is None:
for fname in os.listdir(img_folder):
ext = fname.split('.')[-1]
if ext not in ["png", "jpg"]:
continue
src_img_path = os.path.join(img_folder, fname)
imgs.append((fname, src_img_path))
else:
with open(img_listfile, "r") as f:
names = f.readlines()
names = [l.strip("\n") for l in names]
for fname in names:
ext = fname.split('.')[-1]
if ext not in ["png", "jpg"]:
continue
src_img_path = os.path.join(img_folder, fname)
imgs.append((fname, src_img_path))
prefix = "/".join(fname.split("/")[:-1])
if not os.path.exists(os.path.join(lmks_folder, prefix)):
os.makedirs(os.path.join(lmks_folder, prefix))
imgs.sort()
args = SimpleNamespace()
args.val_img_dir = './dataset/WFLW_test/images/'
args.val_landmarks_dir = './dataset/WFLW_test/landmarks/'
args.ckpt_save_path = './experiments/eval_iccv_0620'
args.hg_blocks = 4
args.pretrained_weights = './ckpt/WFLW_4HG.pth'
args.num_landmarks = 98
args.end_relu = 'False'
args.batch_size = 1
args.gray_scale = False
args.save_vis = True
VAL_IMG_DIR = args.val_img_dir
VAL_LANDMARKS_DIR = args.val_landmarks_dir
CKPT_SAVE_PATH = args.ckpt_save_path
BATCH_SIZE = args.batch_size
PRETRAINED_WEIGHTS = args.pretrained_weights
GRAY_SCALE = False if args.gray_scale == 'False' else True
HG_BLOCKS = args.hg_blocks
END_RELU = False if args.end_relu == 'False' else True
NUM_LANDMARKS = args.num_landmarks
device = "cuda"
use_gpu = torch.cuda.is_available()
print("use_gpu", use_gpu)
model_ft = models.FAN(HG_BLOCKS, END_RELU, GRAY_SCALE, NUM_LANDMARKS)
if PRETRAINED_WEIGHTS != "None":
checkpoint = torch.load(PRETRAINED_WEIGHTS,
map_location=torch.device('cpu'))
if 'state_dict' not in checkpoint:
model_ft.load_state_dict(checkpoint)
else:
pretrained_weights = checkpoint['state_dict']
model_weights = model_ft.state_dict()
pretrained_weights = {k: v for k, v in pretrained_weights.items() \
if k in model_weights}
model_weights.update(pretrained_weights)
model_ft.load_state_dict(model_weights)
to_tensor = transforms.Compose(
[transforms.Resize(256), transforms.ToTensor()])
model_ft = model_ft.to(device)
with torch.no_grad():
for fname, src_img_path in imgs:
image = Image.open(src_img_path)
image = image.convert('RGB')
w, h = image.size
# padding if neccesary
desired_size = 640
ratio = desired_size / float(min(w, h))
w2, h2 = int(w * ratio), int(h * ratio)
image = image.resize((w2, h2))
image_raw = image
old_size = image.size # old_size[0] is in (width, height) format
new_im = Image.new("RGB", (desired_size, desired_size))
x0 = (desired_size - old_size[0]) // 2
y0 = (desired_size - old_size[1]) // 2
new_im.paste(image, (x0, y0))
image = new_im
isz = image.size[0]
inputs = to_tensor(image).unsqueeze(0)
inputs = inputs.to(device)
outputs, boundary_channels = model_ft(inputs)
pred_heatmap = outputs[-1][:, :-1, :, :][0].detach().cpu()
pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0))
pred_landmarks = pred_landmarks.squeeze() * 4 / 256.0 * isz
lmks68 = torch.zeros([68, 2])
for i98, i68 in WFLW_98_TO_DLIB_68_IDX_MAPPING.items():
lmks68[i68, :] = pred_landmarks[i98, :]
#lmks68 = pred_landmarks[:-2]
lmks68[:, 0] -= x0
lmks68[:, 1] -= y0
lmks68 /= ratio
ext = fname.split('.')[-1]
save_npy_path = os.path.join(lmks_folder,
fname.replace(ext, "npy"))
print(save_npy_path)
np.save(save_npy_path, lmks68)
image = image_raw.resize((w, h))
if args.save_vis:
plt.figure()
plt.imshow(image)
plt.scatter(lmks68[:, 0], lmks68[:, 1], s=5, marker='.', c='r')
save_vis_path = os.path.join(lmks_folder, fname)
plt.savefig(save_vis_path)
print(save_vis_path)
def eval_model(model, dataloaders, dataset_sizes,
writer, use_gpu=True, epoches=5, dataset='val',
save_path='./', num_landmarks=68):
global_nme = 0
model.eval()
for epoch in range(epoches):
running_loss = 0
step = 0
total_nme = 0
total_count = 0
fail_count = 0
nmes = []
# running_corrects = 0
# Iterate over data.
with torch.no_grad():
for data in dataloaders[dataset]:
total_runtime = 0
run_count = 0
step_start = time.time()
step += 1
# get the inputs
inputs = data['image'].type(torch.FloatTensor)
labels_heatmap = data['heatmap'].type(torch.FloatTensor)
labels_boundary = data['boundary'].type(torch.FloatTensor)
landmarks = data['landmarks'].type(torch.FloatTensor)
loss_weight_map = data['weight_map'].type(torch.FloatTensor)
# wrap them in Variable
if use_gpu:
inputs = inputs.to(device)
labels_heatmap = labels_heatmap.to(device)
labels_boundary = labels_boundary.to(device)
loss_weight_map = loss_weight_map.to(device)
else:
inputs, labels_heatmap = Variable(inputs), Variable(labels_heatmap)
labels_boundary = Variable(labels_boundary)
labels = torch.cat((labels_heatmap, labels_boundary), 1)
single_start = time.time()
outputs, boundary_channels = model(inputs)
single_end = time.time()
total_runtime += time.time() - single_start
run_count += 1
step_end = time.time()
for i in range(inputs.shape[0]):
print(inputs.shape)
img = inputs[i]
img = img.cpu().numpy()
img = img.transpose((1, 2, 0)) #*255.0
# img = img.astype(np.uint8)
# img = Image.fromarray(img)
# pred_heatmap = outputs[-1][i].detach().cpu()[:-1, :, :]
pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
pred_landmarks, _ = get_preds_fromhm(pred_heatmap.unsqueeze(0))
pred_landmarks = pred_landmarks.squeeze().numpy()
gt_landmarks = data['landmarks'][i].numpy()
print(pred_landmarks, gt_landmarks)
import cv2
while(True):
imgshow = vis_landmark_on_img(cv2.UMat(img), pred_landmarks*4)
cv2.imshow('img', imgshow)
if(cv2.waitKey(10) == ord('q')):
break
if num_landmarks == 68:
left_eye = np.average(gt_landmarks[36:42], axis=0)
right_eye = np.average(gt_landmarks[42:48], axis=0)
norm_factor = np.linalg.norm(left_eye - right_eye)
# norm_factor = np.linalg.norm(gt_landmarks[36]- gt_landmarks[45])
elif num_landmarks == 98:
norm_factor = np.linalg.norm(gt_landmarks[60]- gt_landmarks[72])
elif num_landmarks == 19:
left, top = gt_landmarks[-2, :]
right, bottom = gt_landmarks[-1, :]
norm_factor = math.sqrt(abs(right - left)*abs(top-bottom))
gt_landmarks = gt_landmarks[:-2, :]
elif num_landmarks == 29:
# norm_factor = np.linalg.norm(gt_landmarks[8]- gt_landmarks[9])
norm_factor = np.linalg.norm(gt_landmarks[16]- gt_landmarks[17])
single_nme = (np.sum(np.linalg.norm(pred_landmarks*4 - gt_landmarks, axis=1)) / pred_landmarks.shape[0]) / norm_factor
nmes.append(single_nme)
total_count += 1
if single_nme > 0.1:
fail_count += 1
if step % 10 == 0:
print('Step {} Time: {:.6f} Input Mean: {:.6f} Output Mean: {:.6f}'.format(
step, step_end - step_start,
torch.mean(labels),
torch.mean(outputs[0])))
# gt_landmarks = landmarks.numpy()
# pred_heatmap = outputs[-1].to('cpu').numpy()
gt_landmarks = landmarks
batch_nme = fan_NME(outputs[-1][:, :-1, :, :].detach().cpu(), gt_landmarks, num_landmarks)
# batch_nme = 0
total_nme += batch_nme
epoch_nme = total_nme / dataset_sizes['val']
global_nme += epoch_nme
nme_save_path = os.path.join(save_path, 'nme_log.npy')
np.save(nme_save_path, np.array(nmes))
print('NME: {:.6f} Failure Rate: {:.6f} Total Count: {:.6f} Fail Count: {:.6f}'.format(epoch_nme, fail_count/total_count, total_count, fail_count))
print('Evaluation done! Average NME: {:.6f}'.format(global_nme/epoches))
print('Everage runtime for a single batch: {:.6f}'.format(total_runtime/run_count))
return model