def test():
data_loader = DataLoader(dataset=HellenDataset(False, 224),
batch_size=1,
shuffle=True,
num_workers=1)
device = torch.device("cpu")
model = CnnAlign().to(device)
state = torch.load(MODEL_FACE_ALIGN)
model.load_state_dict(state['net'])
to_pil_img = tfs.ToPILImage()
model.eval()
for img, label in data_loader:
start = time.time()
output = model(img)
end = time.time()
cost = (end - start)
print("cost : " + str(cost))
pil_img = to_pil_img(img[0].cpu())
ann = output[0].cpu().detach().numpy()
ann = np.resize(ann, (194, 2))
draw_ann(pil_img, ann.tolist(), font1, font_size)
plt.figure(num=1, figsize=(15, 8), dpi=80) # 开启一个窗口，同时设置大小，分辨率
plt.imshow(pil_img)
plt.show()
plt.close()
image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
return image
def get_pytorch_model(path):
device = torch.device("cpu")
model = CnnAlign().to(device)
state = torch.load(path, map_location='cpu')
model.load_state_dict(state['net'])
model.eval()
return model
if __name__ == "__main__":
face_align_model = get_pytorch_model(MODEL_FACE_ALIGN)
to_tensor = tfs.ToTensor()
cap = cv2.VideoCapture(0)
while(1):
# get a frame
ret, img = cap.read()
image = cv2image(img)
image = image.resize((224, 224))
output = face_align_model(to_tensor(image).view(-1, 3, 224, 224))
ann = output[0].cpu().detach().numpy()
ann = np.resize(ann, (194, 2))
draw_ann(image, ann, font1, font_size)
img = image2cv(image)
cv2.imshow("test", img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
def train(args):
start_epoch = 0
data_loader = DataLoader(dataset=HellenDataset(True, 224),
batch_size=args.batch,
shuffle=True,
num_workers=16)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model = CnnAlignHard()
print("add graph")
writer.add_graph(model, torch.zeros((1, 3, 224, 224)))
print("add graph over")
if args.pretrained and os.path.exists(MODEL_SAVE_PATH):
print("loading ...")
state = torch.load(MODEL_SAVE_PATH)
model.load_state_dict(state['net'])
start_epoch = state['epoch']
print("loading over, start_epoch: " + str(start_epoch))
model = torch.nn.DataParallel(model, device_ids=[0, 1]) # multi-GPU
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
scheduler = StepLR(optimizer, step_size=args.step, gamma=args.gama)
train_loss = 0
to_pil_img = tfs.ToPILImage()
to_tensor = tfs.ToTensor()
model.train()
for epoch in range(start_epoch, start_epoch + args.epoes):
model.train()
prefetcher = DataPrefetcher(data_loader)
img_tensor, label_tensor = prefetcher.next()
last_img_tensor = img_tensor
last_label_tensor = label_tensor
optimizer.zero_grad()
i_batch = 0
while img_tensor is not None:
last_img_tensor = img_tensor
last_label_tensor = label_tensor
output = model(img_tensor)
loss = torch.nn.functional.smooth_l1_loss(
output, label_tensor.view(-1, output.size(1)))
if loss is None:
img_tensor, label_tensor = prefetcher.next()
continue
loss.backward()
if i_batch % args.mini_batch == 0:
optimizer.step()
optimizer.zero_grad()
train_loss = loss.item()
global_step = epoch * len(data_loader) + i_batch
progress_bar(i_batch, len(data_loader),
'loss: %f, epeche: %d' % (train_loss, epoch))
writer.add_scalar("loss", train_loss, global_step=global_step)
img_tensor, label_tensor = prefetcher.next()
i_batch += 1
#save one pic and output
pil_img = to_pil_img(last_img_tensor[0].cpu())
ann = output[0].cpu().detach().numpy()
ann = np.resize(ann, (194, 2))
draw_ann(pil_img, ann.tolist(), font1, font_size)
writer.add_image("img: " + str(epoch), to_tensor(pil_img))
scheduler.step()
if epoch % 10 == 0:
print('Saving..')
state = {
'net': model.module.state_dict(),
'epoch': epoch,
}
torch.save(state, "./output/face_align_hard" + str(epoch) + ".pt")
if not os.path.isdir('data'):
os.mkdir('data')
print('Saving..')
state = {
'net': model.module.state_dict(),
'epoch': epoch,
}
torch.save(state, MODEL_SAVE_PATH)
writer.close()