File size: 5,706 Bytes
8c212a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""Miscellaneous utility functions."""
import os
import pickle
import numpy as np
import PIL.Image
import PIL.ImageFont
import dnnlib
#----------------------------------------------------------------------------
# Convenience wrappers for pickle that are able to load data produced by
# older versions of the code, and from external URLs.
def open_file_or_url(file_or_url):
if dnnlib.util.is_url(file_or_url):
return dnnlib.util.open_url(file_or_url, cache_dir='.stylegan2-cache')
return open(file_or_url, 'rb')
def load_pkl(file_or_url):
with open_file_or_url(file_or_url) as file:
return pickle.load(file, encoding='latin1')
def save_pkl(obj, filename):
with open(filename, 'wb') as file:
pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)
#----------------------------------------------------------------------------
# Image utils.
def adjust_dynamic_range(data, drange_in, drange_out):
if drange_in != drange_out:
scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (np.float32(drange_in[1]) - np.float32(drange_in[0]))
bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale)
data = data * scale + bias
return data
def create_image_grid(images, grid_size=None):
assert images.ndim == 3 or images.ndim == 4
num, img_w, img_h = images.shape[0], images.shape[-1], images.shape[-2]
if grid_size is not None:
grid_w, grid_h = tuple(grid_size)
else:
grid_w = max(int(np.ceil(np.sqrt(num))), 1)
grid_h = max((num - 1) // grid_w + 1, 1)
grid = np.zeros(list(images.shape[1:-2]) + [grid_h * img_h, grid_w * img_w], dtype=images.dtype)
for idx in range(num):
x = (idx % grid_w) * img_w
y = (idx // grid_w) * img_h
grid[..., y : y + img_h, x : x + img_w] = images[idx]
return grid
def convert_to_pil_image(image, drange=[0,1]):
assert image.ndim == 2 or image.ndim == 3
if image.ndim == 3:
if image.shape[0] == 1:
image = image[0] # grayscale CHW => HW
else:
image = image.transpose(1, 2, 0) # CHW -> HWC
image = adjust_dynamic_range(image, drange, [0,255])
image = np.rint(image).clip(0, 255).astype(np.uint8)
fmt = 'RGB' if image.ndim == 3 else 'L'
return PIL.Image.fromarray(image, fmt)
def save_image_grid(images, filename, drange=[0,1], grid_size=None):
convert_to_pil_image(create_image_grid(images, grid_size), drange).save(filename)
def apply_mirror_augment(minibatch):
mask = np.random.rand(minibatch.shape[0]) < 0.5
minibatch = np.array(minibatch)
minibatch[mask] = minibatch[mask, :, :, ::-1]
return minibatch
#----------------------------------------------------------------------------
# Loading data from previous training runs.
def parse_config_for_previous_run(run_dir):
with open(os.path.join(run_dir, 'submit_config.pkl'), 'rb') as f:
data = pickle.load(f)
data = data.get('run_func_kwargs', {})
return dict(train=data, dataset=data.get('dataset_args', {}))
#----------------------------------------------------------------------------
# Size and contents of the image snapshot grids that are exported
# periodically during training.
def setup_snapshot_image_grid(training_set,
size = '1080p', # '1080p' = to be viewed on 1080p display, '4k' = to be viewed on 4k display.
layout = 'random'): # 'random' = grid contents are selected randomly, 'row_per_class' = each row corresponds to one class label.
# Select size.
gw = 1; gh = 1
if size == '1080p':
gw = np.clip(1920 // training_set.shape[2], 3, 32)
gh = np.clip(1080 // training_set.shape[1], 2, 32)
if size == '4k':
gw = np.clip(3840 // training_set.shape[2], 7, 32)
gh = np.clip(2160 // training_set.shape[1], 4, 32)
if size == '8k':
gw = np.clip(7680 // training_set.shape[2], 7, 32)
gh = np.clip(4320 // training_set.shape[1], 4, 32)
# Initialize data arrays.
reals = np.zeros([gw * gh] + training_set.shape, dtype=training_set.dtype)
labels = np.zeros([gw * gh, training_set.label_size], dtype=training_set.label_dtype)
# Random layout.
if layout == 'random':
reals[:], labels[:] = training_set.get_minibatch_np(gw * gh)
# Class-conditional layouts.
class_layouts = dict(row_per_class=[gw,1], col_per_class=[1,gh], class4x4=[4,4])
if layout in class_layouts:
bw, bh = class_layouts[layout]
nw = (gw - 1) // bw + 1
nh = (gh - 1) // bh + 1
blocks = [[] for _i in range(nw * nh)]
for _iter in range(1000000):
real, label = training_set.get_minibatch_np(1)
idx = np.argmax(label[0])
while idx < len(blocks) and len(blocks[idx]) >= bw * bh:
idx += training_set.label_size
if idx < len(blocks):
blocks[idx].append((real, label))
if all(len(block) >= bw * bh for block in blocks):
break
for i, block in enumerate(blocks):
for j, (real, label) in enumerate(block):
x = (i % nw) * bw + j % bw
y = (i // nw) * bh + j // bw
if x < gw and y < gh:
reals[x + y * gw] = real[0]
labels[x + y * gw] = label[0]
return (gw, gh), reals, labels
#----------------------------------------------------------------------------
|