def generator_demo_example_lips(img_path):
name = img_path.split('/')[-1]
landmark_path = os.path.join('../image/', name.replace('jpg', 'npy'))
region_path = os.path.join('../image/',
name.replace('.jpg', '_region.jpg'))
roi, landmark = crop_image(img_path)
if np.sum(landmark[37:39, 1] - landmark[40:42, 1]) < -9:
# pts2 = np.float32(np.array([template[36],template[45],template[30]]))
template = np.load('../basics/base_68.npy')
else:
template = np.load('../basics/base_68_close.npy')
# pts2 = np.float32(np.vstack((template[27:36,:], template[39,:],template[42,:],template[45,:])))
pts2 = np.float32(template[27:45, :])
# pts2 = np.float32(template[17:35,:])
# pts1 = np.vstack((landmark[27:36,:], landmark[39,:],landmark[42,:],landmark[45,:]))
pts1 = np.float32(landmark[27:45, :])
# pts1 = np.float32(landmark[17:35,:])
tform = tf.SimilarityTransform()
tform.estimate(pts2, pts1)
dst = tf.warp(roi, tform, output_shape=(163, 163))
dst = np.array(dst * 255, dtype=np.uint8)
dst = dst[1:129, 1:129, :]
cv2.imwrite(region_path, dst)
gray = cv2.cvtColor(dst, cv2.COLOR_BGR2GRAY)
# detect faces in the grayscale image
rects = detector(gray, 1)
for (i, rect) in enumerate(rects):
shape = predictor(gray, rect)
shape = utils.shape_to_np(shape)
shape, _, _ = normLmarks(shape)
np.save(landmark_path, shape)
lmark = shape.reshape(68, 2)
name = img_path[:-4] + 'lmark.png'
utils.plot_flmarks(lmark,
name, (-0.2, 0.2), (-0.2, 0.2),
'x',
'y',
figsize=(10, 10))
return dst, lmark
def test():
os.environ["CUDA_VISIBLE_DEVICES"] = config.device_ids
config.is_train = 'test'
if config.lstm == True:
generator = AT_net()
else:
generator = AT_single()
if config.dataset == 'grid':
pca = torch.FloatTensor(np.load('../basics/U_grid.npy')[:, :6]).cuda()
mean = torch.FloatTensor(np.load('../basics/mean_grid.npy')).cuda()
elif config.dataset == 'lrw':
pca = torch.FloatTensor(np.load('../basics/U_lrw1.npy')[:, :6]).cuda()
mean = torch.FloatTensor(np.load('../basics/mean_lrw1.npy')).cuda()
else:
raise Exception('wrong key word for the dataset input')
xLim = (-1.0, 1.0)
yLim = (-1.0, 1.0)
xLab = 'x'
yLab = 'y'
state_dict = multi2single(config.model_name, 1)
generator.load_state_dict(state_dict)
print('load pretrained [{}]'.format(config.model_name))
if config.lstm:
if config.dataset == 'lrw':
dataset = LRW_1D_lstm_landmark_pca(config.dataset_dir,
train=config.is_train)
else:
dataset = GRID_1D_lstm_landmark_pca(config.dataset_dir,
train=config.is_train)
else:
if config.dataset == 'lrw':
dataset = LRW_1D_single_landmark_pca(config.dataset_dir,
train=config.is_train)
else:
dataset = GRID_1D_single_landmark_pca(config.dataset_dir,
train=config.is_train)
data_loader = DataLoader(dataset,
batch_size=config.batch_size,
num_workers=config.num_thread,
shuffle=False,
drop_last=True)
data_iter = iter(data_loader)
data_iter.next()
if not os.path.exists(config.sample_dir):
os.mkdir(config.sample_dir)
if not os.path.exists(os.path.join(config.sample_dir, 'fake')):
os.mkdir(os.path.join(config.sample_dir, 'fake'))
if not os.path.exists(os.path.join(config.sample_dir, 'real')):
os.mkdir(os.path.join(config.sample_dir, 'real'))
if config.cuda:
generator = generator.cuda()
generator.eval()
for step, (example_landmark, example_audio, lmark,
audio) in enumerate(data_loader):
with torch.no_grad():
print(step)
if step == 5:
break
if config.cuda:
example_audio = Variable(example_audio.float()).cuda()
lmark = Variable(lmark.float()).cuda()
audio = Variable(audio.float()).cuda()
example_landmark = Variable(example_landmark.float()).cuda()
if config.lstm:
fake_lmark = generator(example_landmark, audio)
fake_lmark = fake_lmark.view(
fake_lmark.size(0) * fake_lmark.size(1), 6)
fake_lmark[:, 1:6] *= 2 * torch.FloatTensor(
np.array([1.2, 1.4, 1.6, 1.8, 2.0])).cuda()
fake_lmark = torch.mm(fake_lmark, pca.t())
fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
fake_lmark = fake_lmark.view(config.batch_size, 16, 136)
fake_lmark = fake_lmark.data.cpu().numpy()
lmark = lmark.view(lmark.size(0) * lmark.size(1), 6)
lmark = torch.mm(lmark, pca.t())
lmark = lmark + mean.expand_as(lmark)
lmark = lmark.view(config.batch_size, 16, 136)
lmark = lmark.data.cpu().numpy()
for indx in range(config.batch_size):
for jj in range(16):
name = "{}real_{}_{}_{}.png".format(
config.sample_dir, step, indx, jj)
utils.plot_flmarks(lmark[indx, jj],
name,
xLim,
yLim,
xLab,
yLab,
figsize=(10, 10))
name = "{}fake_{}_{}_{}.png".format(
config.sample_dir, step, indx, jj)
utils.plot_flmarks(fake_lmark[indx, jj],
name,
xLim,
yLim,
xLab,
yLab,
figsize=(10, 10))
else:
fake_lmark = generator(example_landmark, audio)
fake_lmark[:, 1:6] *= 2 * torch.FloatTensor(
np.array([1.2, 1.4, 1.6, 1.8, 2.0])).cuda()
fake_lmark = torch.mm(fake_lmark, pca.t())
fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
fake_lmark = fake_lmark.data.cpu().numpy()
lmark = torch.mm(lmark, pca.t())
lmark = lmark + mean.expand_as(lmark)
lmark = lmark.data.cpu().numpy()
for indx in range(config.batch_size):
name = '{}real/real_{}.png'.format(
config.sample_dir, step * config.batch_size + indx)
utils.plot_flmarks(lmark[indx],
name,
xLim,
yLim,
xLab,
yLab,
figsize=(10, 10))
name = '{}fake/fake_{}.png'.format(
config.sample_dir, step * config.batch_size + indx)
utils.plot_flmarks(fake_lmark[indx],
name,
xLim,
yLim,
xLab,
yLab,
figsize=(10, 10))
def fit(self):
config = self.config
if config.dataset == 'grid':
pca = torch.FloatTensor( np.load('basics/pca.npy')[:,:6]).cuda()
mean =torch.FloatTensor( np.load('basics/mean.npy')).cuda()
elif config.dataset == 'lrw':
pca = torch.FloatTensor(np.load('basics/pca.npy')[:,:6] ).cuda()
mean = torch.FloatTensor(np.load('basics/mean.npy')).cuda()
else:
raise Exception('wrong key word for the dataset input')
num_steps_per_epoch = len(self.data_loader)
cc = 0
t0 = time.time()
xLim=(-1.0, 1.0)
yLim=(-1.0, 1.0)
xLab = 'x'
yLab = 'y'
for epoch in range(self.start_epoch, config.max_epochs):
for step, (example_landmark, example_audio, lmark, audio) in enumerate(self.data_loader):
t1 = time.time()
if config.cuda:
lmark = Variable(lmark.float()).cuda()
audio = Variable(audio.float()).cuda()
example_landmark = Variable(example_landmark.float()).cuda()
else:
lmark = Variable(lmark.float())
audio = Variable(audio.float())
example_landmark = Variable(example_landmark.float())
fake_lmark= self.generator( example_landmark, audio)
loss = self.mse_loss_fn(fake_lmark , lmark)
loss.backward()
self.opt_g.step()
self._reset_gradients()
if (step+1) % 10 == 0 or (step+1) == num_steps_per_epoch:
steps_remain = num_steps_per_epoch-step+1 + \
(config.max_epochs-epoch+1)*num_steps_per_epoch
print("[{}/{}][{}/{}] loss1: {:.8f},data time: {:.4f}, model time: {} second"
.format(epoch+1, config.max_epochs,
step+1, num_steps_per_epoch, loss, t1-t0, time.time() - t1))
if (step) % (int(num_steps_per_epoch / 10 )) == 0 and step != 0:
if config.lstm:
lmark = lmark.view(lmark.size(0) * lmark.size(1) , 6)
lmark = torch.mm( lmark, pca.t() )
lmark = lmark + mean.expand_as(lmark)
fake_lmark = fake_lmark.view(fake_lmark.size(0) * fake_lmark.size(1) , 6)
fake_lmark[:, 1:6] *= 2*torch.FloatTensor(np.array([1.1, 1.2, 1.3, 1.4, 1.5])).cuda()
fake_lmark = torch.mm( fake_lmark, pca.t() )
fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
lmark = lmark.view(config.batch_size, 16, 136)
lmark = lmark.data.cpu().numpy()
fake_lmark = fake_lmark.view(config.batch_size, 16, 136)
fake_lmark = fake_lmark.data.cpu().numpy()
for indx in range(3):
for jj in range(16):
name = "{}real_{}_{}_{}.png".format(config.sample_dir,cc, indx,jj)
utils.plot_flmarks(lmark[indx,jj], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
name = "{}fake_{}_{}_{}.png".format(config.sample_dir,cc, indx,jj)
utils.plot_flmarks(fake_lmark[indx,jj], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
torch.save(self.generator.state_dict(),
"{}/atnet_lstm_{}.pth"
.format(config.model_dir,cc))
else:
lmark = torch.mm( lmark, pca.t() )
lmark = lmark + mean.expand_as(lmark)
fake_lmark[ 1:6] *= 2*torch.FloatTensor(np.array([1.1, 1.2, 1.3, 1.4, 1.5])).cuda()
fake_lmark = torch.mm( fake_lmark, pca.t() )
fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
lmark = lmark.data.cpu().numpy()
fake_lmark = fake_lmark.data.cpu().numpy()
for indx in range(0,config.batch_size, 8):
name = "{}single_real_{}_{}.png".format(config.sample_dir,cc, indx)
utils.plot_flmarks(lmark[indx], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
name = "{}single_fake_{}_{}.png".format(config.sample_dir,cc, indx)
utils.plot_flmarks(fake_lmark[indx], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
torch.save(self.generator.state_dict(),
"{}anet_single_{}.pth"
.format(config.model_dir,cc))
cc += 1
t0 = time.time()