def main():
print('Loading networks from "%s"...' % a.model)
sess = tflib.init_tf()
with open(a.model, 'rb') as file:
network = pickle.load(file, encoding='latin1')
try: _, _, Gs = network
except: Gs = network
resolution = tuple(Gs.output_shape[2:])
proj = projector.Projector(a.steps)
proj.set_network(Gs)
img_files = img_list(a.in_dir)
num_images = len(img_files)
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx+1, num_images))
images = img_read(img_files[image_idx])
images = np.expand_dims(np.transpose(images, [2,0,1]), 0)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
work_dir = osp.join(a.out_dir, basename(img_files[image_idx]))
os.makedirs(work_dir, exist_ok=True)
project_image(proj, images, work_dir, resolution, a.num_snapshots)
def generate():
os.makedirs(a.out_dir, exist_ok=True)
np.random.seed(seed=696)
device = torch.device('cuda')
# setup generator
Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.verbose = a.verbose
Gs_kwargs.size = a.size
Gs_kwargs.scale_type = a.scale_type
# mask/blend latents with external latmask or by splitting the frame
if a.latmask is None:
nHW = [int(s) for s in a.nXY.split('-')][::-1]
assert len(nHW) == 2, ' Wrong count nXY: %d (must be 2)' % len(nHW)
n_mult = nHW[0] * nHW[1]
if a.verbose is True and n_mult > 1:
print(' Latent blending w/split frame %d x %d' % (nHW[1], nHW[0]))
lmask = np.tile(np.asarray([[[[1]]]]), (1, n_mult, 1, 1))
Gs_kwargs.countHW = nHW
Gs_kwargs.splitfine = a.splitfine
else:
if a.verbose is True: print(' Latent blending with mask', a.latmask)
n_mult = 2
if os.path.isfile(a.latmask): # single file
lmask = np.asarray([[img_read(a.latmask)[:, :, 0] / 255.]
]) # [h,w]
elif os.path.isdir(a.latmask): # directory with frame sequence
lmask = np.asarray([[
img_read(f)[:, :, 0] / 255. for f in img_list(a.latmask)
]]) # [h,w]
else:
print(' !! Blending mask not found:', a.latmask)
exit(1)
lmask = np.concatenate((lmask, 1 - lmask), 1) # [frm,2,h,w]
lmask = torch.from_numpy(lmask).to(device)
# load base or custom network
pkl_name = osp.splitext(a.model)[0]
if '.pkl' in a.model.lower():
custom = False
print(' .. Gs from pkl ..', basename(a.model))
else:
custom = True
print(' .. Gs custom ..', basename(a.model))
with dnnlib.util.open_url(pkl_name + '.pkl') as f:
Gs = legacy.load_network_pkl(f,
custom=custom, **Gs_kwargs)['G_ema'].to(
device) # type: ignore
if a.verbose is True: print(' out shape', Gs.output_shape[1:])
if a.verbose is True: print(' making timeline..')
lats = [] # list of [frm,1,512]
for i in range(n_mult):
lat_tmp = latent_anima((1, Gs.z_dim),
a.frames,
a.fstep,
cubic=a.cubic,
gauss=a.gauss,
verbose=False) # [frm,1,512]
lats.append(lat_tmp) # list of [frm,1,512]
latents = np.concatenate(lats, 1) # [frm,X,512]
print(' latents', latents.shape)
latents = torch.from_numpy(latents).to(device)
frame_count = latents.shape[0]
# distort image by tweaking initial const layer
if a.digress > 0:
try:
init_res = Gs.init_res
except:
init_res = (4, 4) # default initial layer size
dconst = []
for i in range(n_mult):
dc_tmp = a.digress * latent_anima([1, Gs.z_dim, *init_res],
a.frames,
a.fstep,
cubic=True,
verbose=False)
dconst.append(dc_tmp)
dconst = np.concatenate(dconst, 1)
else:
dconst = np.zeros([frame_count, 1, 1, 1, 1])
dconst = torch.from_numpy(dconst).to(device)
# labels / conditions
label_size = Gs.c_dim
if label_size > 0:
labels = torch.zeros((frame_count, n_mult, label_size),
device=device) # [frm,X,lbl]
if a.labels is None:
label_ids = []
for i in range(n_mult):
label_ids.append(random.randint(0, label_size - 1))
else:
label_ids = [int(x) for x in a.labels.split('-')]
label_ids = label_ids[:n_mult] # ensure we have enough labels
for i, l in enumerate(label_ids):
labels[:, i, l] = 1
else:
labels = [None]
# generate images from latent timeline
pbar = ProgressBar(frame_count)
for i in range(frame_count):
latent = latents[i] # [X,512]
label = labels[i % len(labels)]
latmask = lmask[i %
len(lmask)] if lmask is not None else [None] # [X,h,w]
dc = dconst[i % len(dconst)] # [X,512,4,4]
# generate multi-latent result
if custom:
output = Gs(latent,
label,
latmask,
dc,
truncation_psi=a.trunc,
noise_mode='const')
else:
output = Gs(latent,
label,
truncation_psi=a.trunc,
noise_mode='const')
output = (output.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(
torch.uint8).cpu().numpy()
# save image
ext = 'png' if output.shape[3] == 4 else 'jpg'
filename = osp.join(a.out_dir, "%06d.%s" % (i, ext))
imsave(filename, output[0])
pbar.upd()
# convert latents to dlatents, save them
if a.save_lat is True:
latents = latents.squeeze(1) # [frm,512]
dlatents = Gs.mapping(latents, label) # [frm,18,512]
if a.size is None: a.size = [''] * 2
filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1],
a.size[0])
filename = osp.join(osp.dirname(a.out_dir), filename)
dlatents = dlatents.cpu().numpy()
np.save(filename, dlatents)
print('saved dlatents', dlatents.shape, 'to', filename)
def main():
os.makedirs(a.out_dir, exist_ok=True)
np.random.seed(seed=696)
# setup generator
fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.func_name = 'training.stylegan2_multi.G_main'
Gs_kwargs.verbose = a.verbose
Gs_kwargs.size = a.size
Gs_kwargs.scale_type = a.scale_type
Gs_kwargs.impl = a.ops
# mask/blend latents with external latmask or by splitting the frame
if a.latmask is None:
nHW = [int(s) for s in a.nXY.split('-')][::-1]
assert len(nHW)==2, ' Wrong count nXY: %d (must be 2)' % len(nHW)
n_mult = nHW[0] * nHW[1]
if a.verbose is True and n_mult > 1: print(' Latent blending w/split frame %d x %d' % (nHW[1], nHW[0]))
lmask = np.tile(np.asarray([[[[None]]]]), (1,n_mult,1,1))
Gs_kwargs.countHW = nHW
Gs_kwargs.splitfine = a.splitfine
else:
if a.verbose is True: print(' Latent blending with mask', a.latmask)
n_mult = 2
if os.path.isfile(a.latmask): # single file
lmask = np.asarray([[img_read(a.latmask)[:,:,0] / 255.]]) # [h,w]
elif os.path.isdir(a.latmask): # directory with frame sequence
lmask = np.asarray([[img_read(f)[:,:,0] / 255. for f in img_list(a.latmask)]]) # [h,w]
else:
print(' !! Blending mask not found:', a.latmask); exit(1)
lmask = np.concatenate((lmask, 1 - lmask), 1) # [frm,2,h,w]
Gs_kwargs.latmask_res = lmask.shape[2:]
# load model with arguments
sess = tflib.init_tf({'allow_soft_placement':True})
pkl_name = osp.splitext(a.model)[0]
with open(pkl_name + '.pkl', 'rb') as file:
network = pickle.load(file, encoding='latin1')
try: _, _, network = network
except: pass
for k in list(network.static_kwargs.keys()):
Gs_kwargs[k] = network.static_kwargs[k]
# reload custom network, if needed
if '.pkl' in a.model.lower():
print(' .. Gs from pkl ..', basename(a.model))
Gs = network
else: # reconstruct network
print(' .. Gs custom ..', basename(a.model))
# print(Gs_kwargs)
Gs = tflib.Network('Gs', **Gs_kwargs)
Gs.copy_vars_from(network)
if a.verbose is True: print('kwargs:', ['%s: %s'%(kv[0],kv[1]) for kv in sorted(Gs.static_kwargs.items())])
if a.verbose is True: print(' out shape', Gs.output_shape[1:])
if a.size is None: a.size = Gs.output_shape[2:]
if a.verbose is True: print(' making timeline..')
lats = [] # list of [frm,1,512]
for i in range(n_mult):
lat_tmp = latent_anima((1, Gs.input_shape[1]), a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, verbose=False) # [frm,1,512]
lats.append(lat_tmp) # list of [frm,1,512]
latents = np.concatenate(lats, 1) # [frm,X,512]
print(' latents', latents.shape)
frame_count = latents.shape[0]
# distort image by tweaking initial const layer
if a.digress > 0:
try: latent_size = Gs.static_kwargs['latent_size']
except: latent_size = 512 # default latent size
try: init_res = Gs.static_kwargs['init_res']
except: init_res = (4,4) # default initial layer size
dconst = []
for i in range(n_mult):
dc_tmp = a.digress * latent_anima([1, latent_size, *init_res], a.frames, a.fstep, cubic=True, verbose=False)
dconst.append(dc_tmp)
dconst = np.concatenate(dconst, 1)
else:
dconst = np.zeros([frame_count, 1, 1, 1, 1])
# labels / conditions
try:
label_size = Gs_kwargs.label_size
except:
label_size = 0
if label_size > 0:
labels = np.zeros((frame_count, n_mult, label_size)) # [frm,X,lbl]
if a.labels is None:
label_ids = []
for i in range(n_mult):
label_ids.append(random.randint(0, label_size-1))
else:
label_ids = [int(x) for x in a.labels.split('-')]
label_ids = label_ids[:n_mult] # ensure we have enough labels
for i, l in enumerate(label_ids):
labels[:,i,l] = 1
else:
labels = [None]
# generate images from latent timeline
pbar = ProgressBar(frame_count)
for i in range(frame_count):
latent = latents[i] # [X,512]
label = labels[i % len(labels)]
latmask = lmask[i % len(lmask)] if lmask is not None else [None] # [X,h,w]
dc = dconst[i % len(dconst)] # [X,512,4,4]
# generate multi-latent result
if Gs.num_inputs == 2:
output = Gs.run(latent, label, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt)
else:
output = Gs.run(latent, label, latmask, dc, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt)
# save image
ext = 'png' if output.shape[3]==4 else 'jpg'
filename = osp.join(a.out_dir, "%06d.%s" % (i,ext))
imsave(filename, output[0])
pbar.upd()
# convert latents to dlatents, save them
if a.save_lat is True:
latents = latents.squeeze(1) # [frm,512]
dlatents = Gs.components.mapping.run(latents, label, dtype='float16') # [frm,18,512]
filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0])
filename = osp.join(osp.dirname(a.out_dir), filename)
np.save(filename, dlatents)
print('saved dlatents', dlatents.shape, 'to', filename)