|
import argparse |
|
import os |
|
|
|
import cv2 |
|
import numpy as np |
|
import torch |
|
|
|
from model import Generator |
|
from psp_encoder.psp_encoders import PSPEncoder |
|
from utils import ten2cv, cv2ten |
|
import glob |
|
import random |
|
|
|
seed = 0 |
|
|
|
random.seed(seed) |
|
np.random.seed(seed) |
|
torch.manual_seed(seed) |
|
torch.cuda.manual_seed_all(seed) |
|
|
|
|
|
if __name__ == '__main__': |
|
device = 'cpu' |
|
|
|
parser = argparse.ArgumentParser() |
|
|
|
parser.add_argument('--size', type=int, default=1024) |
|
|
|
parser.add_argument('--ckpt', type=str, default='', help='path to BlendGAN checkpoint') |
|
parser.add_argument('--psp_encoder_ckpt', type=str, default='', help='path to psp_encoder checkpoint') |
|
|
|
parser.add_argument('--style_img_path', type=str, default=None, help='path to style image') |
|
parser.add_argument('--input_img_path', type=str, default=None, help='path to input image') |
|
parser.add_argument('--add_weight_index', type=int, default=6) |
|
|
|
parser.add_argument('--channel_multiplier', type=int, default=2) |
|
parser.add_argument('--outdir', type=str, default="") |
|
|
|
args = parser.parse_args() |
|
|
|
|
|
|
|
args.latent = 512 |
|
args.n_mlp = 8 |
|
|
|
checkpoint = torch.load(args.ckpt) |
|
model_dict = checkpoint['g_ema'] |
|
print('ckpt: ', args.ckpt) |
|
|
|
g_ema = Generator( |
|
args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier |
|
).to(device) |
|
g_ema.load_state_dict(model_dict) |
|
g_ema.eval() |
|
|
|
psp_encoder = PSPEncoder(args.psp_encoder_ckpt, output_size=args.size).to(device) |
|
psp_encoder.eval() |
|
|
|
num = 0 |
|
|
|
|
|
print(num) |
|
num += 1 |
|
|
|
img_in = cv2.imread(args.input_img_path) |
|
img_in_ten = cv2ten(img_in, device) |
|
img_in = cv2.resize(img_in, (args.size, args.size)) |
|
|
|
|
|
img_style = cv2.imread(args.style_img_path) |
|
img_style_ten = cv2ten(img_style, device) |
|
img_style = cv2.resize(img_style, (args.size, args.size)) |
|
|
|
with torch.no_grad(): |
|
sample_style = g_ema.get_z_embed(img_style_ten) |
|
sample_in = psp_encoder(img_in_ten) |
|
img_out_ten, _ = g_ema([sample_in], z_embed=sample_style, add_weight_index=args.add_weight_index, |
|
input_is_latent=True, return_latents=False, randomize_noise=False) |
|
img_out = ten2cv(img_out_ten) |
|
|
|
cv2.imwrite('out.jpg', img_out) |
|
|
|
print('Done!') |
|
|
|
|
