|
|
|
|
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from tqdm import tqdm |
|
from lpips import LPIPS |
|
import numpy as np |
|
from torch_utils.models import Generator as bodyGAN |
|
from torch_utils.models_face import Generator as FaceGAN |
|
import dlib |
|
from utils.face_alignment import align_face_for_insetgan |
|
from utils.util import visual,tensor_to_numpy, numpy_to_tensor |
|
import legacy |
|
import os |
|
import click |
|
|
|
|
|
class InsetGAN(torch.nn.Module): |
|
def __init__(self, stylebody_ckpt, styleface_ckpt): |
|
super().__init__() |
|
|
|
|
|
if not os.path.exists(stylebody_ckpt.replace('.pkl','.pth')): |
|
legacy.convert(stylebody_ckpt, stylebody_ckpt.replace('.pkl','.pth')) |
|
stylebody_ckpt = stylebody_ckpt.replace('.pkl','.pth') |
|
|
|
if not os.path.exists(styleface_ckpt.replace('.pkl','.pth')): |
|
legacy.convert(styleface_ckpt, styleface_ckpt.replace('.pkl','.pth')) |
|
styleface_ckpt = styleface_ckpt.replace('.pkl','.pth') |
|
|
|
|
|
config = {"latent" : 512, "n_mlp" : 8, "channel_multiplier": 2} |
|
self.body_generator = bodyGAN( |
|
size = 1024, |
|
style_dim=config["latent"], |
|
n_mlp=config["n_mlp"], |
|
channel_multiplier=config["channel_multiplier"] |
|
) |
|
self.body_generator.load_state_dict(torch.load(stylebody_ckpt)['g_ema']) |
|
self.body_generator.eval().requires_grad_(False).cuda() |
|
|
|
self.face_generator = FaceGAN( |
|
size = 1024, |
|
style_dim=config["latent"], |
|
n_mlp=config["n_mlp"], |
|
channel_multiplier=config["channel_multiplier"] |
|
) |
|
self.face_generator.load_state_dict(torch.load(styleface_ckpt)['g_ema']) |
|
self.face_generator.eval().requires_grad_(False).cuda() |
|
|
|
self.dlib_predictor = dlib.shape_predictor('./pretrained_models/shape_predictor_68_face_landmarks.dat') |
|
self.dlib_cnn_face_detector = dlib.cnn_face_detection_model_v1("pretrained_models/mmod_human_face_detector.dat") |
|
|
|
|
|
self.lpips_loss = LPIPS(net='alex').cuda().eval() |
|
self.l1_loss = torch.nn.L1Loss(reduction='mean') |
|
|
|
def loss_coarse(self, A_face, B, p1=500, p2=0.05): |
|
A_face = F.interpolate(A_face, size=(64, 64), mode='area') |
|
B = F.interpolate(B, size=(64, 64), mode='area') |
|
loss_l1 = p1 * self.l1_loss(A_face, B) |
|
loss_lpips = p2 * self.lpips_loss(A_face, B) |
|
return loss_l1 + loss_lpips |
|
|
|
@staticmethod |
|
def get_border_mask(A, x, spec): |
|
mask = torch.zeros_like(A) |
|
mask[:, :, :x, ] = 1 |
|
mask[:, :, -x:, ] = 1 |
|
mask[:, :, :, :x ] = 1 |
|
mask[:, :, :, -x:] = 1 |
|
return mask |
|
|
|
@staticmethod |
|
def get_body_mask(A, crop, padding=4): |
|
mask = torch.ones_like(A) |
|
mask[:, :, crop[1]-padding:crop[3]+padding, crop[0]-padding:crop[2]+padding] = 0 |
|
return mask |
|
|
|
def loss_border(self, A_face, B, p1=10000, p2=2, spec=None): |
|
mask = self.get_border_mask(A_face, 8, spec) |
|
loss_l1 = p1 * self.l1_loss(A_face*mask, B*mask) |
|
loss_lpips = p2 * self.lpips_loss(A_face*mask, B*mask) |
|
return loss_l1 + loss_lpips |
|
|
|
def loss_body(self, A, B, crop, p1=9000, p2=0.1): |
|
padding = int((crop[3] - crop[1]) / 20) |
|
mask = self.get_body_mask(A, crop, padding) |
|
loss_l1 = p1 * self.l1_loss(A*mask, B*mask) |
|
loss_lpips = p2 * self.lpips_loss(A*mask, B*mask) |
|
return loss_l1+loss_lpips |
|
|
|
def loss_face(self, A, B, crop, p1=5000, p2=1.75): |
|
mask = 1 - self.get_body_mask(A, crop) |
|
loss_l1 = p1 * self.l1_loss(A*mask, B*mask) |
|
loss_lpips = p2 * self.lpips_loss(A*mask, B*mask) |
|
return loss_l1+loss_lpips |
|
|
|
def loss_reg(self, w, w_mean, p1, w_plus_delta=None, p2=None): |
|
return p1 * torch.mean(((w - w_mean) ** 2)) + p2 * torch.mean(w_plus_delta ** 2) |
|
|
|
|
|
def detect_face_dlib(self, img): |
|
|
|
img = tensor_to_numpy(img) |
|
aligned_image, crop, rect = align_face_for_insetgan(img=img, |
|
detector=self.dlib_cnn_face_detector, |
|
predictor=self.dlib_predictor, |
|
output_size=256) |
|
|
|
aligned_image = np.array(aligned_image) |
|
aligned_image = numpy_to_tensor(aligned_image) |
|
return aligned_image, crop, rect |
|
|
|
|
|
def dual_optimizer(self, |
|
face_w, |
|
body_w, |
|
joint_steps=500, |
|
face_initial_learning_rate=0.02, |
|
body_initial_learning_rate=0.05, |
|
lr_rampdown_length=0.25, |
|
lr_rampup_length=0.05, |
|
seed=None, |
|
output_path=None, |
|
video=0): |
|
''' |
|
Given a face_w, optimize a body_w with suitable body pose & shape for face_w |
|
''' |
|
def visual_(path, synth_body, synth_face, body_crop, step, both=False, init_body_with_face=None): |
|
tmp = synth_body.clone().detach() |
|
tmp[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]] = synth_face |
|
if both: |
|
tmp = torch.cat([synth_body, tmp], dim=3) |
|
save_path = os.path.join(path, f"{step:04d}.jpg") |
|
visual(tmp, save_path) |
|
|
|
def forward(face_w_opt, |
|
body_w_opt, |
|
face_w_delta, |
|
body_w_delta, |
|
body_crop, |
|
update_crop=False |
|
): |
|
if face_w_opt.shape[1] != 18: |
|
face_ws = (face_w_opt).repeat([1, 18, 1]) |
|
else: |
|
face_ws = face_w_opt.clone() |
|
face_ws = face_ws + face_w_delta |
|
synth_face, _ = self.face_generator([face_ws], input_is_latent=True, randomize_noise=False) |
|
|
|
body_ws = (body_w_opt).repeat([1, 18, 1]) |
|
body_ws = body_ws + body_w_delta |
|
synth_body, _ = self.body_generator([body_ws], input_is_latent=True, randomize_noise=False) |
|
|
|
if update_crop: |
|
old_r = (body_crop[3]-body_crop[1]) // 2, (body_crop[2]-body_crop[0]) // 2 |
|
_, body_crop, _ = self.detect_face_dlib(synth_body) |
|
center = (body_crop[1] + body_crop[3]) // 2, (body_crop[0] + body_crop[2]) // 2 |
|
body_crop = (center[1] - old_r[1], center[0] - old_r[0], center[1] + old_r[1], center[0] + old_r[0]) |
|
|
|
synth_body_face = synth_body[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]] |
|
|
|
if synth_face.shape[2] > body_crop[3]-body_crop[1]: |
|
synth_face_resize = F.interpolate(synth_face, size=(body_crop[3]-body_crop[1], body_crop[2]-body_crop[0]), mode='area') |
|
|
|
return synth_body, synth_body_face, synth_face, synth_face_resize, body_crop |
|
|
|
def update_lr(init_lr, step, num_steps, lr_rampdown_length, lr_rampup_length): |
|
t = step / num_steps |
|
lr_ramp = min(1.0, (1.0 - t) / lr_rampdown_length) |
|
lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi) |
|
lr_ramp = lr_ramp * min(1.0, t / lr_rampup_length) |
|
lr = init_lr * lr_ramp |
|
return lr |
|
|
|
|
|
output_path = os.path.join(output_path, seed) |
|
os.makedirs(output_path, exist_ok=True) |
|
|
|
|
|
body_w_mean = self.body_generator.mean_latent(10000).detach() |
|
face_w_opt = face_w.clone().detach().requires_grad_(True) |
|
body_w_opt = body_w.clone().detach().requires_grad_(True) |
|
face_w_delta = torch.zeros_like(face_w.repeat([1, 18, 1])).requires_grad_(True) |
|
body_w_delta = torch.zeros_like(body_w.repeat([1, 18, 1])).requires_grad_(True) |
|
|
|
ref_body, _ = self.body_generator([body_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False) |
|
|
|
ref_face, _ = self.face_generator([face_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False) |
|
|
|
_, body_crop, _ = self.detect_face_dlib(ref_body) |
|
_, _, face_crop = self.detect_face_dlib(ref_face) |
|
|
|
face_optimizer = torch.optim.Adam([face_w_opt, face_w_delta], betas=(0.9, 0.999), lr=face_initial_learning_rate) |
|
body_optimizer = torch.optim.Adam([body_w_opt, body_w_delta], betas=(0.9, 0.999), lr=body_initial_learning_rate) |
|
|
|
global_step = 0 |
|
|
|
face_steps = 25 |
|
pbar = tqdm(range(face_steps)) |
|
for step in pbar: |
|
face_lr = update_lr(face_initial_learning_rate / 2, step, face_steps, lr_rampdown_length, lr_rampup_length) |
|
for param_group in face_optimizer.param_groups: |
|
param_group['lr'] =face_lr |
|
synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, |
|
body_w_opt, |
|
face_w_delta, |
|
body_w_delta, |
|
body_crop) |
|
loss_face = self.loss_face(synth_face_raw, ref_face, face_crop, 5000, 1.75) |
|
loss_coarse = self.loss_coarse(synth_face, synth_body_face, 50, 0.05) |
|
loss_border = self.loss_border(synth_face, synth_body_face, 1000, 0.1) |
|
loss = loss_coarse + loss_border + loss_face |
|
face_optimizer.zero_grad() |
|
loss.backward() |
|
face_optimizer.step() |
|
|
|
if video: |
|
visual_(output_path, synth_body, synth_face, body_crop, global_step) |
|
pbar.set_description( |
|
( |
|
f"face: {step:.4f}, lr: {face_lr}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};" |
|
f"loss_border: {loss_border.item():.2f}, loss_face: {loss_face.item():.2f};" |
|
) |
|
) |
|
global_step += 1 |
|
|
|
|
|
body_steps = 150 |
|
pbar = tqdm(range(body_steps)) |
|
for step in pbar: |
|
body_lr = update_lr(body_initial_learning_rate, step, body_steps, lr_rampdown_length, lr_rampup_length) |
|
update_crop = True if (step % 50 == 0) else False |
|
|
|
for param_group in body_optimizer.param_groups: |
|
param_group['lr'] =body_lr |
|
synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, |
|
body_w_opt, |
|
face_w_delta, |
|
body_w_delta, |
|
body_crop, |
|
update_crop=update_crop) |
|
loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05) |
|
loss_border = self.loss_border(synth_face, synth_body_face, 2500, 0) |
|
loss_body = self.loss_body(synth_body, ref_body, body_crop, 9000, 0.1) |
|
loss_reg = self.loss_reg(body_w_opt, body_w_mean, 15000, body_w_delta, 0) |
|
loss = loss_coarse + loss_border + loss_body + loss_reg |
|
body_optimizer.zero_grad() |
|
loss.backward() |
|
body_optimizer.step() |
|
|
|
|
|
if video: |
|
visual_(output_path, synth_body, synth_face, body_crop, global_step) |
|
pbar.set_description( |
|
( |
|
f"body: {step:.4f}, lr: {body_lr}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};" |
|
f"loss_border: {loss_border.item():.2f}, loss_body: {loss_body.item():.2f}, loss_reg: {loss_reg:.2f}" |
|
) |
|
) |
|
global_step += 1 |
|
|
|
|
|
interval = 50 |
|
joint_face_steps = joint_steps // 2 |
|
joint_body_steps = joint_steps // 2 |
|
face_step = 0 |
|
body_step = 0 |
|
pbar = tqdm(range(joint_steps)) |
|
flag = -1 |
|
for step in pbar: |
|
if step % interval == 0: flag += 1 |
|
text_flag = 'optimize_face' if flag % 2 == 0 else 'optimize_body' |
|
synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, |
|
body_w_opt, |
|
face_w_delta, |
|
body_w_delta, |
|
body_crop) |
|
if text_flag == 'optimize_face': |
|
face_lr = update_lr(face_initial_learning_rate, face_step, joint_face_steps, lr_rampdown_length, lr_rampup_length) |
|
for param_group in face_optimizer.param_groups: |
|
param_group['lr'] =face_lr |
|
loss_face = self.loss_face(synth_face_raw, ref_face, face_crop, 5000, 1.75) |
|
loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05) |
|
loss_border = self.loss_border(synth_face, synth_body_face, 25000, 0) |
|
loss = loss_coarse + loss_border + loss_face |
|
face_optimizer.zero_grad() |
|
loss.backward() |
|
face_optimizer.step() |
|
pbar.set_description( |
|
( |
|
f"face: {step}, lr: {face_lr:.4f}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};" |
|
f"loss_border: {loss_border.item():.2f}, loss_face: {loss_face.item():.2f};" |
|
) |
|
) |
|
face_step += 1 |
|
else: |
|
body_lr = update_lr(body_initial_learning_rate, body_step, joint_body_steps, lr_rampdown_length, lr_rampup_length) |
|
for param_group in body_optimizer.param_groups: |
|
param_group['lr'] =body_lr |
|
loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05) |
|
loss_border = self.loss_border(synth_face, synth_body_face, 2500, 0) |
|
loss_body = self.loss_body(synth_body, ref_body, body_crop, 9000, 0.1) |
|
loss_reg = self.loss_reg(body_w_opt, body_w_mean, 25000, body_w_delta, 0) |
|
loss = loss_coarse + loss_border + loss_body + loss_reg |
|
body_optimizer.zero_grad() |
|
loss.backward() |
|
body_optimizer.step() |
|
pbar.set_description( |
|
( |
|
f"body: {step}, lr: {body_lr:.4f}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};" |
|
f"loss_border: {loss_border.item():.2f}, loss_body: {loss_body.item():.2f}, loss_reg: {loss_reg:.2f}" |
|
) |
|
) |
|
body_step += 1 |
|
if video: |
|
visual_(output_path, synth_body, synth_face, body_crop, global_step) |
|
global_step += 1 |
|
return face_w_opt.repeat([1, 18, 1])+face_w_delta, body_w_opt.repeat([1, 18, 1])+body_w_delta, body_crop |
|
|
|
|
|
|
|
|
|
""" |
|
Jointly combine and optimize generated faces and bodies . |
|
Examples: |
|
|
|
\b |
|
# Combine the generate human full-body image from the provided StyleGAN-Human pre-trained model |
|
# and the generated face image from FFHQ model, optimize both latent codes to produce the coherent face-body image |
|
python insetgan.py --body_network=pretrained_models/stylegan_human_v2_1024.pkl --face_network=pretrained_models/ffhq.pkl \\ |
|
--body_seed=82 --face_seed=43 --trunc=0.6 --outdir=outputs/insetgan/ --video 1 |
|
""" |
|
|
|
@click.command() |
|
@click.pass_context |
|
@click.option('--face_network', default="./pretrained_models/ffhq.pkl", help='Network pickle filename', required=True) |
|
@click.option('--body_network', default='./pretrained_models/stylegan2_1024.pkl', help='Network pickle filename', required=True) |
|
@click.option('--face_seed', type=int, default=82, help='selected random seed') |
|
@click.option('--body_seed', type=int, default=43, help='selected random seed') |
|
@click.option('--joint_steps', type=int, default=500, help='num steps for joint optimization') |
|
@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=0.6, show_default=True) |
|
@click.option('--outdir', help='Where to save the output images', default= "outputs/insetgan/" , type=str, required=True, metavar='DIR') |
|
@click.option('--video', help="set to 1 if want to save video", type=int, default=0) |
|
def main( |
|
ctx: click.Context, |
|
face_network: str, |
|
body_network: str, |
|
face_seed: int, |
|
body_seed: int, |
|
joint_steps: int, |
|
truncation_psi: float, |
|
outdir: str, |
|
video: int): |
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
insgan = InsetGAN(body_network, face_network) |
|
os.makedirs(outdir, exist_ok=True) |
|
face_z = np.random.RandomState(face_seed).randn(1, 512).astype(np.float32) |
|
face_mean = insgan.face_generator.mean_latent(3000) |
|
face_w = insgan.face_generator.get_latent(torch.from_numpy(face_z).to(device)) |
|
face_w = truncation_psi * face_w + (1-truncation_psi) * face_mean |
|
face_img, _ = insgan.face_generator([face_w], input_is_latent=True) |
|
|
|
body_z = np.random.RandomState(body_seed).randn(1, 512).astype(np.float32) |
|
body_mean = insgan.body_generator.mean_latent(3000) |
|
body_w = insgan.body_generator.get_latent(torch.from_numpy(body_z).to(device)) |
|
body_w = truncation_psi * body_w + (1-truncation_psi) * body_mean |
|
body_img, _ = insgan.body_generator([body_w], input_is_latent=True) |
|
|
|
_, body_crop, _ = insgan.detect_face_dlib(body_img) |
|
face_img = F.interpolate(face_img, size=(body_crop[3]-body_crop[1], body_crop[2]-body_crop[0]), mode='area') |
|
cp_body = body_img.clone() |
|
cp_body[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]] = face_img |
|
|
|
optim_face_w, optim_body_w, crop = insgan.dual_optimizer( |
|
face_w, |
|
body_w, |
|
joint_steps=joint_steps, |
|
seed=f'{face_seed:04d}_{body_seed:04d}', |
|
output_path=outdir, |
|
video=video |
|
) |
|
|
|
if video: |
|
ffmpeg_cmd = f"ffmpeg -hide_banner -loglevel error -i ./{outdir}/{face_seed:04d}_{body_seed:04d}/%04d.jpg -c:v libx264 -vf fps=30 -pix_fmt yuv420p ./{outdir}/{face_seed:04d}_{body_seed:04d}.mp4" |
|
os.system(ffmpeg_cmd) |
|
new_face_img, _ = insgan.face_generator([optim_face_w], input_is_latent=True) |
|
new_shape = crop[3] - crop[1], crop[2] - crop[0] |
|
new_face_img_crop = F.interpolate(new_face_img, size=new_shape, mode='area') |
|
seamless_body, _ = insgan.body_generator([optim_body_w], input_is_latent=True) |
|
seamless_body[:, :, crop[1]:crop[3], crop[0]:crop[2]] = new_face_img_crop |
|
temp = torch.cat([cp_body, seamless_body], dim=3) |
|
visual(temp, f"{outdir}/{face_seed:04d}_{body_seed:04d}.png") |
|
|
|
if __name__ == "__main__": |
|
main() |
