def train():
# define model, dataloader, 3dmm eigenvectors, optimization method
model = densenet
model.load_state_dict(torch.load("model/chkpoint_000.pt"))
model.cuda()
#optimizer2 = torch.optim.Adam(model.parameters(),lr=1e-4)
adjustmentnet.load_state_dict(torch.load("model/chkpoint_adj_000.pt"))
adjustmentnet.cuda()
adjustmentnet.train()
loader = dataloader.BatchLoader
face3dmm = dataloader.Face3DMM()
optimizer = torch.optim.Adam(list(adjustmentnet.parameters()) +
list(model.parameters()),
lr=1e-5)
# main training loop
for epoch in itertools.count():
for i, batch in enumerate(loader):
optimizer.zero_grad()
x = batch['image'].cuda()
y = batch['lm2d'].cuda()
y_pred = model(x)
batchsize = x.shape[0]
alphas = y_pred[:, :199]
betas = y_pred[:, 199:228]
s = y_pred[:, 228]
t = y_pred[:, 229:231]
t = torch.tanh(t)
r = y_pred[:, 231:235]
r = torch.tanh(r) * (3.14 / 4)
# apply 3DMM model from predicted parameters
alpha_matrix = alphas.unsqueeze(2).expand(
*alphas.size(), alphas.size(1)) * torch.eye(
alphas.size(1)).cuda()
beta_matrix = betas.unsqueeze(2).expand(*betas.size(),
betas.size(1)) * torch.eye(
betas.size(1)).cuda()
shape_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.shape_eigenvec]),
alpha_matrix)
exp_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.exp_eigenvec]), beta_matrix)
shape_cov = shape_cov.sum(2)
exp_cov = exp_cov.sum(2)
# alignment
shape = (face3dmm.mu_shape.unsqueeze(0) + shape_cov.view(
(batchsize, 53215, 3))) + exp_cov.view((batchsize, 53215, 3))
lm = shape[:, face3dmm.lm, :]
R = util.R(r).cuda()
scaledshape = s.unsqueeze(1).unsqueeze(1) * torch.bmm(lm, R)
alignedshape = t.unsqueeze(1) + scaledshape[:, :, :2]
# adjustment network applied onto the landmarks of 3dMM taking input image
adjustment = torch.tanh(adjustmentnet(x).view(-1, 68, 2))
pred = (alignedshape + adjustment) * 112 + 112
gt = y * 112 + 112
# weight update
loss = torch.mean(torch.norm(gt - pred, p=2, dim=2))
loss.backward()
optimizer.step()
gt = gt[0].cpu().data.numpy()
pred = pred[0].cpu().data.numpy()
sample = x[0].cpu().permute(1, 2, 0).data.numpy()
sample = sample * 255
util.viewLM(sample.astype(np.uint8), pred)
#io.imsave(f"example_{i:04d}.png",sample)
print(f"epoch/batch {epoch}/{i} | Loss: {loss:.4f}")
print("saving!")
torch.save(model.state_dict(), f"model/chkpoint_{epoch:03d}.pt")
torch.save(adjustmentnet.state_dict(),
f"model/chkpoint_adj_{epoch:03d}.pt")
import itertools
import argparse
import os
import numpy as np
import torch
import cv2
from skimage import io, transform
import dataloader
import util
from model import densenet
from model import adjustmentnet
# read the 3dmm eigenvectors
face3dmm = dataloader.Face3DMM()
# define model, dataloader, 3dmm eigenvectors, optimization method
model = densenet
model.load_state_dict(torch.load("model/chkpoint_000.pt"))
model.cuda()
adjustmentnet.load_state_dict(torch.load("model/chkpoint_adj_000.pt"))
adjustmentnet.cuda()
loader = dataloader.W300Loader()
#loader = dataloader.LFWLoader()
#loader = dataloader.BatchLoader
face3dmm = dataloader.Face3DMM()
results = []
def train():
# define model, dataloader, 3dmm eigenvectors, optimization method
model = densenet
model.cuda()
adjustmentnet.cuda()
loader = dataloader.BatchLoader
face3dmm = dataloader.Face3DMM()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# main training loop
for epoch in itertools.count():
for i, batch in enumerate(loader):
optimizer.zero_grad()
x = batch['image'].cuda()
y = batch['lm2d'].cuda()
y_pred = model(x)
batchsize = x.shape[0]
alphas = y_pred[:, :199]
betas = y_pred[:, 199:228]
s = y_pred[:, 228]
t = y_pred[:, 229:231]
t = torch.tanh(t)
r = y_pred[:, 231:235]
r = torch.tanh(r) * (3.14 / 4)
# apply 3DMM model from predicted parameters
alpha_matrix = alphas.unsqueeze(2).expand(
*alphas.size(), alphas.size(1)) * torch.eye(
alphas.size(1)).cuda()
beta_matrix = betas.unsqueeze(2).expand(*betas.size(),
betas.size(1)) * torch.eye(
betas.size(1)).cuda()
shape_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.shape_eigenvec]),
alpha_matrix)
exp_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.exp_eigenvec]), beta_matrix)
shape_cov = shape_cov.sum(2)
exp_cov = exp_cov.sum(2)
# alignment
shape = (face3dmm.mu_shape.unsqueeze(0) + shape_cov.view(
(batchsize, 53215, 3))) + exp_cov.view((batchsize, 53215, 3))
lm = shape[:, face3dmm.lm, :]
R = util.R(r).cuda()
scaledshape = s.unsqueeze(1).unsqueeze(1) * torch.bmm(lm, R)
alignedshape = t.unsqueeze(1) + scaledshape[:, :, :2]
# adjustment network applied onto the landmarks of 3dMM taking input image
# adjustment = adjustmentnet(x).view(-1,68,2)
# weight update
loss = torch.mean(torch.abs(y - (alignedshape)))
loss.backward()
optimizer.step()
print(f"epoch/batch {epoch}/{i} | Loss: {loss:.4f}")
print("saving!")
torch.save(model.state_dict(), f"model/chkpoint_{epoch:03d}.pt")
for p in lm2d:
x = int(p[0])
y = int(p[1])
cv2.circle(myimg, (x, y), 2, [0, 0, 255], -1)
cv2.imshow('img', (myimg).astype(np.uint8))
cv2.waitKey(1)
return cv2.cvtColor(myimg, cv2.COLOR_BGR2RGB)
if __name__ == '__main__':
import dataloader
facemodel = dataloader.Face3DMM()
mu_s = facemodel.mu_shape
mu_exp = facemodel.mu_exp
s_eigen = facemodel.shape_eigenvec
exp_eigen = facemodel.exp_eigenvec
lm = facemodel.lm
alphas = torch.matmul(torch.randn(199), torch.eye(199)).float().cuda()
betas = torch.matmul(torch.randn(29), torch.eye(29)).float().cuda()
euler = np.random.rand(3)
R = torch.Tensor(euler2rotm(euler)).float().cuda()
T = torch.randn((1, 3)).float().cuda() * 10
s = torch.randn(1).float().cuda()