def run_network(image, mask=None):
global latent
image = np.transpose(image, (2,0,1))
real = misc.adjust_dynamic_range(image, [0, 255], [-1, 1])
real = np.array([real])
if mask is None:
mask = np.ones((1, 1, 512, 512), np.uint8)
#mask[0, 0, 270:500, 128:384] = 0 # mouth
mask[0, 0, 128:384, 0:256] = 0 # eye
else:
mask = np.array([[mask]])
fake = network.run(latent, [1], real, mask, truncation_psi=None)
return misc.adjust_dynamic_range(fake, [-1, 1], [0, 255]).clip(0, 255).astype(np.uint8)
def process_weight_tf(weight, training_set, testing_set, queue):
dnnlib.tflib.init_tf()
print("running: " + weight)
G, D, Gs = pickle.load(open(weight, "rb"))
train_pred = []
train_lab = []
train_data = "/".join(training_set.split("/")[:-1])
train_dir = training_set.split("/")[-1]
test_data = "/".join(testing_set.split("/")[:-1])
test_dir = testing_set.split("/")[-1]
train = dataset.load_dataset(data_dir=train_data,
tfrecord_dir=train_dir,
max_label_size=1,
repeat=False,
shuffle_mb=0)
test = dataset.load_dataset(data_dir=test_data,
tfrecord_dir=test_dir,
max_label_size=1,
repeat=False,
shuffle_mb=0)
for x in range(train._np_labels.shape[0] - 1):
try:
image, label = train.get_minibatch_np(1)
image = misc.adjust_dynamic_range(image, [0, 255], [-1, 1])
except:
break
train_pred.append(D.run(image, None)[0][0])
train_lab.append(label)
print("done train")
test_pred = []
test_lab = []
for x in range(test._np_labels.shape[0] - 1):
try:
image, label = test.get_minibatch_np(1)
image = misc.adjust_dynamic_range(image, [0, 255], [-1, 1])
except:
break
test_pred.append(D.run(image, None)[0][0])
test_lab.append(label)
hter = calculate_metrics(train_pred, train_lab, test_pred, test_lab)
queue.put((weight.split("-")[-1].split(".")[0], hter))
def _evaluate(self, Gs, E, Inv, num_gpus):
minibatch_size = num_gpus * self.minibatch_per_gpu
resolution = Gs.components.synthesis.output_shape[2]
placeholder_portraits = tf.placeholder(tf.float32, [self.minibatch_per_gpu, 3, resolution, resolution], name='placeholder_portraits')
placeholder_landmarks = tf.placeholder(tf.float32, [self.minibatch_per_gpu, 3, resolution, resolution], name='placeholder_landmarks')
placeholder_keypoints = tf.placeholder(tf.float32, [self.minibatch_per_gpu, 136], name='placeholder_landmarks')
fake_X_val = self.run_image_manipulation(E, Gs, Inv, placeholder_portraits, placeholder_landmarks, placeholder_keypoints, num_gpus)
hd_sum = 0.0
failed_counter = 0
for idx, data in enumerate(self._iterate_reals(minibatch_size=minibatch_size)):
image_data = data[0]
batch_portraits = image_data[:,0,:,:,:]
batch_landmarks = np.roll(image_data[:,1,:,:,:], shift=1, axis=0)
keypoints = np.roll(data[1], shift=1, axis=0)
batch_portraits = misc.adjust_dynamic_range(batch_portraits.astype(np.float32), [0, 255], [-1., 1.])
batch_landmarks = misc.adjust_dynamic_range(batch_landmarks.astype(np.float32), [0, 255], [-1., 1.])
begin = idx * minibatch_size
end = min(begin + minibatch_size, self.num_images)
samples_manipulated = tflib.run(fake_X_val, feed_dict={placeholder_portraits: batch_portraits, placeholder_landmarks: batch_landmarks, placeholder_keypoints: keypoints})
samples_manipulated = misc.adjust_dynamic_range(samples_manipulated.astype(np.float32), [-1., 1.], [0, 255])
samples_manipulated = np.transpose(samples_manipulated, [0, 2, 3, 1])
batch_landmarks = misc.adjust_dynamic_range(batch_landmarks.astype(np.float32), [-1., 1.], [0, 255])
for i in range(minibatch_size):
try:
ground_truth_lm = batch_landmarks[i]
generated_lm, _ = self.landmark_extractor.generate_landmark_image(source_path_or_image=samples_manipulated[i], resolution=resolution)
generated_lm = generated_lm.cpu().detach().numpy()
ground_truth_lm = np.transpose(ground_truth_lm, [1, 2, 0])
hd_sum += self.calculate_landmark_hausdorff(ground_truth_lm, generated_lm)
except Exception as e:
print('Error: Landmark couldnt be extracted from generated output. Skipping the sample')
failed_counter += 1
continue
if end == self.num_images:
break
avg_hd_dist = hd_sum/(self.num_images - failed_counter)
self._report_result(np.real(avg_hd_dist), suffix='/Average Landmark Hausdorff Distance')
self._report_result(failed_counter, suffix='/Number of failed landmark extractions')
def create_from_images(checkpoint, image, mask, output):
real = np.asarray(PIL.Image.open(image)).transpose([2, 0, 1])
real = misc.adjust_dynamic_range(real, [0, 255], [-1, 1])
mask = np.asarray(PIL.Image.open(mask).convert('1'),
dtype=np.float32)[np.newaxis]
tflib.init_tf()
_, _, Gs = misc.load_pkl(checkpoint)
latent = np.random.randn(1, *Gs.input_shape[1:])
fake = Gs.run(latent, None, real[np.newaxis], mask[np.newaxis])[0]
fake = misc.adjust_dynamic_range(fake, [-1, 1], [0, 255])
fake = fake.clip(0, 255).astype(np.uint8).transpose([1, 2, 0])
fake = PIL.Image.fromarray(fake)
fake.save(output)
def process_reals(x, lod, mirror_augment, drange_data, drange_net):
with tf.name_scope('ProcessReals'):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
s = tf.shape(x)
mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
mask = tf.tile(mask, [1, s[1], s[2], s[3]])
x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
with tf.name_scope(
'FadeLOD'
): # Smooth crossfade between consecutive levels-of-detail.
s = tf.shape(x)
y = tf.reshape(x, [-1, s[1], s[2] // 2, 2, s[3] // 2, 2])
y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
y = tf.tile(y, [1, 1, 1, 2, 1, 2])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
x = tflib.lerp(x, y, lod - tf.floor(lod))
with tf.name_scope(
'UpscaleLOD'
): # Upscale to match the expected input/output size of the networks.
s = tf.shape(x)
factor = tf.cast(2**tf.floor(lod), tf.int32)
x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
x = tf.tile(x, [1, 1, 1, factor, 1, factor])
x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
return x
def process_reals(x, labels, lod, mirror_augment, drange_data, drange_net):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
random_vector = tf.random_uniform([tf.shape(x)[0]]) < 0.5
x = tf.where(random_vector, x, tf.reverse(x, [3]))
rotation_offset = 108
indices_first = tf.range(rotation_offset)
swaps = tf.constant([0, 7, 6, 5, 4, 3, 2, 1]) + rotation_offset
indices_last = tf.range(rotation_offset + 8, tf.shape(labels)[1])
indices = tf.concat([indices_first, swaps, indices_last], axis=0)
mirrored_labels = tf.gather(labels, indices, axis=1)
labels = tf.where(random_vector, labels, mirrored_labels)
with tf.name_scope(
'FadeLOD'
): # Smooth crossfade between consecutive levels-of-detail.
s = tf.shape(x)
y = tf.reshape(x, [-1, s[1], s[2] // 2, 2, s[3] // 2, 2])
y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
y = tf.tile(y, [1, 1, 1, 2, 1, 2])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
x = tflib.lerp(x, y, lod - tf.floor(lod))
with tf.name_scope(
'UpscaleLOD'
): # Upscale to match the expected input/output size of the networks.
s = tf.shape(x)
factor = tf.cast(2**tf.floor(lod), tf.int32)
x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
x = tf.tile(x, [1, 1, 1, factor, 1, factor])
x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
return x, labels
def project_real_images(network_pkl, dataset_name, data_dir, num_images,
num_steps, num_snapshots, save_every_dlatent,
save_final_dlatent):
assert num_snapshots <= num_steps, "Can't have more snapshots than number of steps taken!"
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
proj = projector.Projector(num_steps=num_steps)
proj.set_network(Gs)
print('Loading images from "%s"...' % dataset_name)
dataset_obj = dataset.load_dataset(data_dir=data_dir,
tfrecord_dir=dataset_name,
max_label_size=0,
repeat=False,
shuffle_mb=0)
assert dataset_obj.shape == Gs.output_shape[1:]
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx, num_images))
try:
images, _labels = dataset_obj.get_minibatch_np(1)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
project_image(proj,
targets=images,
png_prefix=dnnlib.make_run_dir_path('image%04d-' %
image_idx),
num_snapshots=num_snapshots,
save_every_dlatent=save_every_dlatent,
save_final_dlatent=save_final_dlatent)
except tf.errors.OutOfRangeError:
print(
f'Error! There are only {image_idx} images in {data_dir}{dataset_name}!'
)
sys.exit(1)
def process_reals(x, lod, mirror_augment, drange_data, drange_net):
with tf.name_scope('ProcessReals'):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
s = tf.shape(x)
mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
mask = tf.tile(mask, [1, s[1], s[2], s[3]])
x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
with tf.name_scope('FadeLOD'): # 连续细节级别之间的平滑淡入淡出。
s = tf.shape(x)
y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
y = tf.tile(y, [1, 1, 1, 2, 1, 2])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
x = tflib.lerp(x, y, lod - tf.floor(lod))
with tf.name_scope('UpscaleLOD'): # 缩放以匹配网络的预期输入/输出大小。
s = tf.shape(x)
factor = tf.cast(2 ** tf.floor(lod), tf.int32)
x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
x = tf.tile(x, [1, 1, 1, factor, 1, factor])
x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
return x
def generate_images(network_pkl, seeds, create_new_G, new_func_name):
tflib.init_tf()
print('Loading networks from "%s"...' % network_pkl)
# _G, _D, Gs = pretrained_networks.load_networks(network_pkl)
_G, _D, I, Gs = misc.load_pkl(network_pkl)
if create_new_G:
Gs = Gs.convert(new_func_name=new_func_name)
# noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
Gs_kwargs = dnnlib.EasyDict()
# Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs.randomize_noise = True
for seed_idx, seed in enumerate(seeds):
print('Generating image for seed %d (%d/%d) ...' %
(seed, seed_idx, len(seeds)))
rnd = np.random.RandomState(seed)
z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
# tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
images, _ = Gs.run(z, None,
**Gs_kwargs) # [minibatch, height, width, channel]
images = misc.adjust_dynamic_range(images, [-1, 1], [0, 255])
# np.clip(images, 0, 255, out=images)
images = np.transpose(images, [0, 2, 3, 1])
images = np.rint(images).clip(0, 255).astype(np.uint8)
PIL.Image.fromarray(images[0], 'RGB').save(
dnnlib.make_run_dir_path('seed%04d.png' % seed))
def generate_domain_shift(network_pkl, seeds, domain_dim):
tflib.init_tf()
print('Loading networks from "%s"...' % network_pkl)
# _G, _D, Gs = pretrained_networks.load_networks(network_pkl)
_G, _D, I, Gs = misc.load_pkl(network_pkl)
Gs_kwargs = dnnlib.EasyDict()
# Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs.randomize_noise = True
for seed_idx, seed in enumerate(seeds):
print('Generating image for seed %d (%d/%d) ...' %
(seed, seed_idx, len(seeds)))
rnd = np.random.RandomState(seed)
z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component]
# tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]
images_1, _ = Gs.run(z, None,
**Gs_kwargs) # [minibatch, c, height, width]
z[:, domain_dim] = -z[:, domain_dim]
images_2, _ = Gs.run(z, None,
**Gs_kwargs) # [minibatch, c, height, width]
images = np.concatenate((images_1, images_2), axis=3)
images = misc.adjust_dynamic_range(images, [-1, 1], [0, 255])
# np.clip(images, 0, 255, out=images)
images = np.transpose(images, [0, 2, 3, 1])
images = np.rint(images).clip(0, 255).astype(np.uint8)
PIL.Image.fromarray(images[0], 'RGB').save(
dnnlib.make_run_dir_path('seed%04d.png' % seed))
def process_reals(x, mirror_augment, drange_data, drange_net, depth=9):
def downsample2x(in_x, factor=2):
assert isinstance(factor,
int) and factor >= 1, "factor can't be negative or 0"
if factor == 1:
return in_x
with tf.variable_scope("Downscale2D"):
ksize = [1, 1, factor, factor]
return tf.nn.avg_pool(
in_x,
ksize=ksize,
strides=ksize,
padding="VALID",
data_format="NCHW"
) # NOTE: requires tf_config['graph_options.place_pruned_graph'] = True
return in_x
with tf.name_scope("ProcessReals"):
with tf.name_scope("DynamicRange"):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope("MirrorAugment"):
s = tf.shape(x)
mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
mask = tf.tile(mask, [1, s[1], s[2], s[3]])
x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
x_s = [x]
for _ in range(depth - 1):
x = downsample2x(x)
x_s.append(x)
return list(reversed(x_s))
def project_real_images(network_pkl, dataset_name, data_dir, num_images,
num_snapshots):
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
proj = projector.Projector()
proj.set_network(Gs)
print('Loading images from "%s"...' % dataset_name)
dataset_obj = dataset.load_dataset(data_dir=data_dir,
tfrecord_dir=dataset_name,
max_label_size=0,
repeat=False,
shuffle_mb=0)
assert dataset_obj.shape == Gs.output_shape[1:]
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx, num_images))
images, _labels = dataset_obj.get_minibatch_np(1)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
print("read image from database:", images.shape)
import cv2
cv2.imwrite("test_project_real_images_image_from_db.png", images)
project_image(proj,
targets=images,
png_prefix=dnnlib.make_run_dir_path('image%04d-' %
image_idx),
num_snapshots=num_snapshots)
def project_real_images(network_pkl, dataset_name, data_dir, num_images,
num_snapshots):
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
proj = projector.Projector()
proj.set_network(Gs)
print('Loading images from "%s"...' % dataset_name)
dataset_obj = dataset.load_dataset(data_dir=data_dir,
tfrecord_dir=dataset_name,
max_label_size=0,
repeat=False,
shuffle_mb=0)
assert dataset_obj.shape == Gs.output_shape[
1:], "%sexpected shape %s, got %s%s" % (
dnnlib.util.Col.RB, Gs.output_shape[1:], dataset_obj.shape,
dnnlib.util.Col.AU)
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx, num_images))
images, _labels = dataset_obj.get_minibatch_np(1)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
project_image(proj,
targets=images,
png_prefix=dnnlib.make_run_dir_path('image%04d-' %
image_idx),
num_snapshots=num_snapshots)
def project_real_images(submit_config, network_pkl, dataset_name, data_dir,
num_images, num_snapshots):
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
proj = projector.Projector()
proj.verbose = submit_config.verbose
proj.set_network(Gs)
print('Loading images from "%s"...' % dataset_name)
dataset_obj = dataset.load_dataset(data_dir=data_dir,
tfrecord_dir=dataset_name,
max_label_size=0,
repeat=False,
shuffle_mb=0)
print('dso shape: ' + str(dataset_obj.shape) + ' vs gs shape: ' +
str(Gs.output_shape[1:]))
assert dataset_obj.shape == Gs.output_shape[1:]
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx, num_images))
images, _labels = dataset_obj.get_minibatch_np(1)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
project_image(proj,
targets=images,
png_prefix=dnnlib.make_run_dir_path('image%04d-' %
image_idx),
num_snapshots=num_snapshots)
def project_image(proj, src_file, dst_dir, tmp_dir, video=False):
data_dir = '%s/dataset' % tmp_dir
if os.path.exists(data_dir):
shutil.rmtree(data_dir)
image_dir = '%s/images' % data_dir
tfrecord_dir = '%s/tfrecords' % data_dir
os.makedirs(image_dir, exist_ok=True)
shutil.copy(src_file, image_dir + '/')
dataset_tool.create_from_images_raw(tfrecord_dir, image_dir, shuffle=0)
dataset_obj = dataset.load_dataset(
data_dir=data_dir, tfrecord_dir='tfrecords',
max_label_size=0, repeat=False, shuffle_mb=0
)
print('Projecting image "%s"...' % os.path.basename(src_file))
images, _labels = dataset_obj.get_minibatch_np(1)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
proj.start(images)
if video:
video_dir = '%s/video' % tmp_dir
os.makedirs(video_dir, exist_ok=True)
while proj.get_cur_step() < proj.num_steps:
print('\r%d / %d ... ' % (proj.get_cur_step(), proj.num_steps), end='', flush=True)
proj.step()
if video:
filename = '%s/%08d.png' % (video_dir, proj.get_cur_step())
misc.save_image_grid(proj.get_images(), filename, drange=[-1,1])
print('\r%-30s\r' % '', end='', flush=True)
os.makedirs(dst_dir, exist_ok=True)
filename = os.path.join(dst_dir, os.path.basename(src_file)[:-4] + '.png')
misc.save_image_grid(proj.get_images(), filename, drange=[-1,1])
filename = os.path.join(dst_dir, os.path.basename(src_file)[:-4] + '.npy')
np.save(filename, proj.get_dlatents()[0])
def prepro_imgs(frames, drange=[0, 255], nchw_to_nhwc=True):
"""Converts NCHW frame tensor in float [-1, 1] to NHWC tensor in uint8 [0, 255]"""
frames = adjust_dynamic_range(frames, [-1, 1], drange)
frames = np.rint(frames).clip(drange[0],
drange[1]) # The rounding is important!
if nchw_to_nhwc:
frames = np.transpose(frames, (0, 2, 3, 1)).astype(np.uint8)
return frames
def process_reals(x, labels, mirror_augment, drange_data, drange_net):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
x = tf.where(tf.random_uniform([tf.shape(x)[0]]) < 0.5, x, tf.reverse(x, [3]))
return x, labels
def generate_batch_factor_code(self, ground_truth_data,
representation_function, num_points,
random_state, batch_size):
"""Sample a single training sample based on a mini-batch of ground-truth data.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observation as input and
outputs a representation.
num_points: Number of points to sample.
random_state: Numpy random state used for randomness.
batch_size: Batchsize to sample points.
Returns:
representations: Codes (num_codes, num_points)-np array.
factors: Factors generating the codes (num_factors, num_points)-np array.
"""
representations = None
factors = None
i = 0
while i < num_points:
num_points_iter = min(num_points - i, batch_size)
current_factors, current_observations = \
ground_truth_data.sample(num_points_iter, random_state)
current_observations = misc.adjust_dynamic_range(
current_observations, [-1., 1.], self.drange_net)
if i == 0:
factors = current_factors
# representations = representation_function(current_observations)
if self.has_label_place:
representations = get_return_v(
representation_function.run(
current_observations,
np.zeros([current_observations.shape[0], 0]),
is_validation=True), 1)
else:
representations = get_return_v(
representation_function.run(current_observations,
is_validation=True), 1)
else:
factors = np.vstack((factors, current_factors))
if self.has_label_place:
representations_i = get_return_v(
representation_function.run(
current_observations,
np.zeros([current_observations.shape[0], 0]),
is_validation=True), 1)
else:
representations_i = get_return_v(
representation_function.run(current_observations,
is_validation=True), 1)
# representations = np.vstack((representations,
# representation_function(
# current_observations)))
representations = np.vstack(
(representations, representations_i))
i += num_points_iter
return np.transpose(representations), np.transpose(factors)
def process_reals(x, drange_data, drange_net, mirror_augment):
with tf.name_scope("DynamicRange"):
x = tf.cast(x, tf.float32)
x.set_shape([None, 3, None, None])
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope("MirrorAugment"):
x = tf.where(tf.random_uniform([tf.shape(x)[0]]) < 0.5, x, tf.reverse(x, [3]))
return x
def sweep_fwd_bcw(start_idx, end_idx, opt_args, forward=True, prefix='', latest_img=None):
if forward:
assert start_idx <= end_idx
num_images = end_idx - start_idx + 1
traveral_order = np.arange(start_idx, end_idx + 1)
else:
assert start_idx >= end_idx
num_images = start_idx - end_idx + 1
traversal_order = np.arange(start_idx, end_idx - 1, -1)
for image_idx in traversal_order:
print('Projecting image %s %d/%d ...' % (prefix, image_idx, num_images))
images = skimage.img_as_float(skimage.io.imread(img_files[image_idx]))
images = np.expand_dims(np.transpose(images, (2, 0, 1)), 0)
images = misc.adjust_dynamic_range(images, [0, 1], [-1, 1])
if opt_args.encourage_advected_last_frame and latest_img is not None:
exrfile = OpenEXR.InputFile(os.path.join(data_dir, dataset_name, 'raw', '%04d.exr' % (image_idx + 1)))
if forward:
speed_vectors = channels_to_ndarray(exrfile, ['vector_x.V', 'vector_y.V'])
else:
speed_vectors = -channels_to_ndarray(exrfile, ['vector_z.V', 'vector_a.V'])
if speed_vectors.shape[0] != latest_img.shape[2] or speed_vectors.shape[1] != latest_img.shape[3]:
if exr_pl is None:
exr_pl = np.zeros((latest_img.shape[2], latest_img.shape[3], 2))
if xv is None or yv is None:
xv, yv = np.meshgrid(np.arange(latest_img.shape[2]), np.arange(latest_img.shape[3]), indexing='ij')
h_diff = latest_img.shape[2] - speed_vectors.shape[0]
w_diff = latest_img.shape[3] - speed_vectors.shape[1]
exr_pl[h_diff // 2: h_diff // 2 + speed_vectors.shape[0], w_diff // 2: w_diff // 2 + speed_vectors.shape[1], :] = speed_vectors[:, :, :]
speed_vectors = exr_pl
advected_img, _ = speed_vector.advect_img(latest_img, -speed_vectors[:, :, 0], -speed_vectors[:, :, 1], xv, yv, is_tensor=True)
if opt_args.zero_constrain_mask_boundary:
nonzero_speed = (speed_vectors[:, :, 0] != 0) * (speed_vectors[:, :, 1] != 0)
dilation_r = np.ceil(np.max(np.abs(speed_vectors))) + 1
selem = skimage.morphology.disk(dilation_r)
dilated_speed = skimage.morphology.dilation(nonzero_speed, selem)
zero_constrain_mask = dilated_speed
zero_constrain_mask[nonzero_speed] = 0
if opt_args.mask_speed_vector:
advect_speed_mask = np.ones((speed_vectors.shape[0], speed_vectors.shape[1]))
advect_speed_zero_idx = (speed_vectors[:, :, 0] == 0) * (speed_vectors[:, :, 1] == 0)
advect_speed_mask[advect_speed_zero_idx] = opt_args.advected_still_scale
if zero_constrain_mask_boundary:
advect_speed_mask[zero_constrain_mask] = 0
advect_speed_mask = np.expand_dims(np.expand_dims(advect_speed_mask, 0), 0)
def project():
steps = int(flask.request.args.get("steps", 1000))
yieldInterval = int(flask.request.args.get("yieldInterval", 10))
# regularizeNoiseWeight = float(flask.request.args.get('regularizeNoiseWeight', 1e5))
imageFile = flask.request.files.get("image")
if not imageFile:
flask.abort(400, "image field is requested.")
global model_name
fetched_model_name = flask.request.args.get("model_name", "ffhq")
print(fetched_model_name, model_name)
model_name = fetched_model_name
Gs, _ = loadGs()
image = PIL.Image.open(imageFile.stream).resize(
(Gs.output_shape[2], Gs.output_shape[3]), PIL.Image.ANTIALIAS)
image_array = np.array(image)[:, :, :3].swapaxes(0, 2).swapaxes(1, 2)
image_array = misc.adjust_dynamic_range(image_array, [0, 255], [-1, 1])
# print('shape:', image_array.shape)
def gen():
proj = loadProjector()
# proj.regularize_noise_weight = regularizeNoiseWeight
proj.start([image_array])
for step in proj.runSteps(steps):
print("\rProjecting: %d / %d" % (step, steps), end="", flush=True)
if step % yieldInterval == 0:
dlatents = proj.get_dlatents()
images = proj.get_images()
pilImage = misc.convert_to_pil_image(
misc.create_image_grid(images), drange=[-1, 1])
fp = io.BytesIO()
pilImage.save(fp, PIL.Image.registered_extensions()[".png"])
imgUrl = "data:image/png;base64,%s" % base64.b64encode(
fp.getvalue()).decode("ascii")
# latentsList = list(dlatents.reshape((-1, dlatents.shape[2])))
# latentCodes = list(map(lambda latents: latentCode.encodeFloat32(latents).decode('ascii'), latentsList))
latentCodes = latentCode.encodeFixed16(
dlatents.flatten()).decode("ascii")
yield json.dumps(
dict(step=step, img=imgUrl,
latentCodes=latentCodes)) + "\n\n"
print("\rProjecting finished.%s" % (" " * 8))
return flask.Response(gen(), mimetype="text/plain")
def process_reals(x, mirror_augment, drange_data, drange_net):
with tf.name_scope('ProcessReals'):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
s = tf.shape(x)
mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
mask = tf.tile(mask, [1, s[1], s[2], s[3]])
x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
return x
def process_reals(x, labels, mirror_augment, drange_data, drange_net):
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
random_vector = tf.random_uniform([tf.shape(x)[0]]) < 0.5
x = tf.where(random_vector, x, tf.reverse(x, [3]))
swaps = tf.constant([0, 7, 6, 5, 4, 3, 2, 1])
mirrored_labels = tf.gather(labels, swaps, axis=1)
labels = tf.where(random_vector, labels, mirrored_labels)
return x, labels
def projectImage(Gs, projc, img_fn):
img = PIL.Image.open(img_fn).convert("RGB")
img = misc.pad_min_square(img)
img = img.resize((Gs.output_shape[2], Gs.output_shape[3]), PIL.Image.ANTIALIAS) # resize mode?
img = np.asarray(img)
channels = img.shape[1] if img.ndim == 3 else 1
img = img[np.newaxis, :, :] if channels == 1 else img.transpose([2, 0, 1]) # HW => CHW # HWC => CHW
assert img.shape == tuple(Gs.output_shape[1:])
img = misc.adjust_dynamic_range(img, [0, 255], [-1, 1])[np.newaxis]
return run_projector.project_img(projc, img) # w.append()
def process_reals(x, labels, lod, mirror_augment, drange_data, drange_net):
rotation_offset = 108
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
random_vector = tf.random_uniform([tf.shape(x)[0]]) < 0.5
x = tf.where(random_vector, x, tf.reverse(x, [3]))
rotation_cos = tf.expand_dims(labels[:, rotation_offset], axis=-1)
rotation_sin = tf.expand_dims(labels[:, rotation_offset + 1],
axis=-1)
angle = tf.atan2(rotation_sin, rotation_cos)
new_rotation_cos = tf.cos(angle)
new_rotation_sin = tf.sin(angle) * -1
mirrored_labels = tf.concat([
labels[:, :rotation_offset], new_rotation_cos,
new_rotation_sin, labels[:, rotation_offset + 2:]
],
axis=1)
labels = tf.where(random_vector, labels, mirrored_labels)
with tf.name_scope('BalanceLabels'):
zero_rotation = tf.expand_dims(tf.zeros([tf.shape(x)[0]]), axis=-1)
removed_labels = tf.concat([
labels[:, :rotation_offset], zero_rotation, zero_rotation,
labels[:, rotation_offset + 2:]
],
axis=1)
condition = tf.equal(labels[:, 108], 0.7071)
random_vector = tf.random_uniform([tf.shape(x)[0]]) < 0.5
remove_condition = tf.logical_and(condition, random_vector)
labels = tf.where(remove_condition, removed_labels, labels)
with tf.name_scope(
'FadeLOD'
): # Smooth crossfade between consecutive levels-of-detail.
s = tf.shape(x)
y = tf.reshape(x, [-1, s[1], s[2] // 2, 2, s[3] // 2, 2])
y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
y = tf.tile(y, [1, 1, 1, 2, 1, 2])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
x = tflib.lerp(x, y, lod - tf.floor(lod))
with tf.name_scope(
'UpscaleLOD'
): # Upscale to match the expected input/output size of the networks.
s = tf.shape(x)
factor = tf.cast(2**tf.floor(lod), tf.int32)
x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
x = tf.tile(x, [1, 1, 1, factor, 1, factor])
x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
return x, labels
def process_reals(x, labels, lod, mirror_augment, drange_data, drange_net):
rotation_offset = 108
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
if mirror_augment:
with tf.name_scope('MirrorAugment'):
random_vector = tf.random_uniform([tf.shape(x)[0]]) < 0.5
x = tf.where(random_vector, x, tf.reverse(x, [3]))
indices_first = tf.range(rotation_offset)
swaps = tf.constant([0, 7, 6, 5, 4, 3, 2, 1]) + rotation_offset
indices_last = tf.range(rotation_offset + 8, tf.shape(labels)[1])
indices = tf.concat([indices_first, swaps, indices_last], axis=0)
mirrored_labels = tf.gather(labels, indices, axis=1)
labels = tf.where(random_vector, labels, mirrored_labels)
with tf.name_scope('FadeLOD'): # Smooth crossfade between consecutive levels-of-detail.
s = tf.shape(x)
y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
y = tf.tile(y, [1, 1, 1, 2, 1, 2])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
x = tflib.lerp(x, y, lod - tf.floor(lod))
with tf.name_scope('UpscaleLOD'): # Upscale to match the expected input/output size of the networks.
s = tf.shape(x)
factor = tf.cast(2 ** tf.floor(lod), tf.int32)
x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
x = tf.tile(x, [1, 1, 1, factor, 1, factor])
x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
with tf.name_scope('InterpolateLabels'):
batch_size = 4
rotations = labels[:, rotation_offset:rotation_offset + 8]
rotation_index = tf.cast(tf.argmax(rotations, axis=1), dtype=tf.int32)
rotation_index_45 = tf.mod(rotation_index, 2)
ones = tf.ones([batch_size])
rotation_shift = tf.cast(
tf.where(tf.random_uniform(shape=[batch_size], minval=-1, maxval=1) > 0, ones, ones * -1), dtype=tf.int32)
rotation_shift_45 = rotation_shift * rotation_index_45
new_rotation_index = tf.mod(rotation_index + rotation_shift_45, 8)
new_rotation = tf.cast(tf.one_hot(new_rotation_index, 8), dtype=tf.int32)
new_rotation = new_rotation * tf.cast(
tf.reduce_max(labels[:, rotation_offset:rotation_offset + 8], axis=-1, keepdims=True), dtype=tf.int32)
new_rotation = tf.cast(new_rotation, dtype=tf.float32)
labels_copy = tf.identity(labels)
labels_interpolate = tf.concat(
[labels_copy[:, :rotation_offset], new_rotation, labels_copy[:, rotation_offset + 8:]], axis=-1)
interpolation_mag = tf.random_uniform(shape=[batch_size, 1])
labels = labels * interpolation_mag + labels_interpolate * (1 - interpolation_mag)
return x, labels
def project_real_images(network_pkl, dataset_name, data_dir, num_images,
num_snapshots, num_steps, save_snapshots=False, save_latents=False,
save_umap=False, save_tiles=False):
print('Loading networks from "%s"...' % network_pkl)
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
proj = projector.Projector()
proj.set_network(Gs)
proj.num_steps = num_steps
print('Loading images from "%s"...' % dataset_name)
dataset_obj = dataset.load_dataset(data_dir=data_dir, tfrecord_dir=dataset_name, max_label_size=0, repeat=False, shuffle_mb=0)
assert dataset_obj.shape == Gs.output_shape[1:]
latents = np.zeros((num_images, Gs.input_shape[1]), dtype=np.float32)
tiles = [None] * num_images
for image_idx in range(num_images):
print('Projecting image %d/%d ...' % (image_idx, num_images))
images, _labels = dataset_obj.get_minibatch_np(1)
tiles[image_idx] = images[0, ...].transpose(1, 2, 0)
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
latents[image_idx, ...] = project_image(proj, targets=images,
png_prefix=dnnlib.make_run_dir_path('image%04d-' % image_idx),
num_snapshots=num_snapshots, save_snapshots=save_snapshots)
if save_latents:
filename = dnnlib.make_run_dir_path('real_image_latent_{:06d}'.format(image_idx))
np.save(filename, latents[image_idx, ...])
if save_latents:
filename = dnnlib.make_run_dir_path('real_image_latents.npy')
np.save(filename, latents)
if save_umap:
reducer = umap.UMAP()
embeddings = reducer.fit_transform(latents)
filename = dnnlib.make_run_dir_path('real_image_umap.json')
with open(filename, 'w', encoding='utf-8') as f:
json.dump(embeddings.tolist(), f, ensure_ascii=False)
if save_tiles:
tiles_prefix = dnnlib.make_run_dir_path('real_tile_solid')
misc.save_texture_grid(tiles, tiles_prefix)
textures_prefix = dnnlib.make_run_dir_path('real_texture_solid')
textures = [misc.make_white_square() for _ in range(len(tiles))]
misc.save_texture_grid(textures, textures_prefix)
filename = dnnlib.make_run_dir_path('labels.json')
with open(filename, 'w', encoding='utf-8') as f:
json.dump([0.0] * len(tiles), f, ensure_ascii=False)
def get_reals(training_set, drange_data, drange_net, label_start, label_end,
minibatch_size):
x, labels = training_set.get_minibatch_tf()
labels = labels[:, label_start:label_end]
sort_index = tf.cast(tf.argsort(tf.reduce_sum(labels, axis=-1),
axis=0,
direction='DESCENDING'),
dtype=tf.int32)
labels = tf.gather(labels, sort_index, axis=0)[:minibatch_size]
x = tf.gather(x, sort_index, axis=0)[:minibatch_size]
with tf.name_scope('DynamicRange'):
x = tf.cast(x, tf.float32)
x = misc.adjust_dynamic_range(x, drange_data, drange_net)
return x, labels
def preprocess(file_path):
# print(file_path)
img = np.asarray(PIL.Image.open(file_path))
# Preprocessing from dataset_tool.create_from_images
img = img.transpose([2, 0, 1]) # HWC => CHW
# img = np.expand_dims(img, axis=0)
img = img.reshape((1, 3, 512, 512))
# Preprocessing from training_loop.process_reals
img = adjust_dynamic_range(data=img,
drange_in=[0, 255],
drange_out=[-1.0, 1.0])
return img
def project_image(proj, src_file, dst_dir, tmp_dir, video=False):
data_dir = '%s/dataset' % tmp_dir # ./stylegan2-tmp/dataset
if os.path.exists(data_dir):
shutil.rmtree(data_dir)
image_dir = '%s/images' % data_dir # ./stylegan2-tmp/dataset/images
tfrecord_dir = '%s/tfrecords' % data_dir # ./stylegan2-tmp/dataset/tfrecords
os.makedirs(image_dir, exist_ok=True)
# 将源图片文件copy到./stylegan2-tmp/dataset/images下
shutil.copy(src_file, image_dir + '/')
# 在./stylegan2-tmp/dataset/tfrecords下生成tfrecord临时文件
# tfrecord临时文件序列化存储了不同lod下的图像的shape和数据
# 举例，如果图像是1024x1024，则tfr_file命名从10--2，如：tfrecords-r10.tfrecords...tfrecords-r05.tfrecords...
dataset_tool.create_from_images(tfrecord_dir, image_dir, shuffle=0)
# TFRecordDataset类在“dataset.py”中定义，从一组.tfrecords文件中加载数据集到dataset_obj
# load_dataset是个helper函数，用于构建dataset对象（在TFRecordDataset类创建对象实例时完成）
dataset_obj = dataset.load_dataset(
data_dir=data_dir, tfrecord_dir='tfrecords',
max_label_size=0, repeat=False, shuffle_mb=0
)
# 生成用于优化迭代的目标图像（组）
print('Projecting image "%s"...' % os.path.basename(src_file))
# 取下一个minibatch=1作为Numpy数组
images, _labels = dataset_obj.get_minibatch_np(1)
# 把images的取值从[0. 255]区间调整到[-1, 1]区间
images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1])
# Projector初始化：start
proj.start(images)
if video:
video_dir = '%s/video' % tmp_dir
os.makedirs(video_dir, exist_ok=True)
while proj.get_cur_step() < proj.num_steps:
print('\r%d / %d ... ' % (proj.get_cur_step(), proj.num_steps), end='', flush=True)
# Projector优化迭代：step
proj.step()
# 如果配置了video选项，将优化过程图像存入./ stylegan2 - tmp / video
if video:
filename = '%s/%08d.png' % (video_dir, proj.get_cur_step())
misc.save_image_grid(proj.get_images(), filename, drange=[-1,1])
print('\r%-30s\r' % '', end='', flush=True)
# 在目的地目录中保存图像，保存dlatents文件
os.makedirs(dst_dir, exist_ok=True)
filename = os.path.join(dst_dir, os.path.basename(src_file)[:-4] + '.png')
misc.save_image_grid(proj.get_images(), filename, drange=[-1,1])
filename = os.path.join(dst_dir, os.path.basename(src_file)[:-4] + '.npy')
np.save(filename, proj.get_dlatents()[0])
