aiden
/
DragGAN
mirror of https://github.com/XingangPan/DragGAN
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
main
Branches Tags
${ item.name }
Create tag ${ searchTerm }
Create branch ${ searchTerm }
from 'main'
${ noResults }
DragGAN/stylegan_human/insetgan.py

398 lines
20 KiB
Python
Raw Permalink Blame History

# Copyright (c) SenseTime Research. All rights reserved.
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__()
## convert pkl to pth
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')
# dual generator
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()
# crop function
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")
# criterion
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)
# FFHQ type
def detect_face_dlib(self, img):
# tensor to numpy array rgb uint8
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
# joint optimization
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
# update output_path
output_path = os.path.join(output_path, seed)
os.makedirs(output_path, exist_ok=True)
# define optimized params
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)
# generate ref face & body
ref_body, _ = self.body_generator([body_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False)
# for inversion
ref_face, _ = self.face_generator([face_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False)
# get initilized crop
_, body_crop, _ = self.detect_face_dlib(ref_body)
_, _, face_crop = self.detect_face_dlib(ref_face) # NOTE: this is face rect only. no FFHQ type.
# create optimizer
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
# Stage1: remove background of face image
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()
# visualization
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
# Stage2: find a suitable body
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
# update_crop = 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()
# visualization
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
# Stage3: joint optimization
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)) # [N, L, C]
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)) # [N, L, C]
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()
Reference in New Issue View Git Blame Copy Permalink