mfrashad/CharacterGAN
10
1# Copyright 2020 Erik Härkönen. All rights reserved.2# This file is licensed to you under the Apache License, Version 2.0 (the "License");3# you may not use this file except in compliance with the License. You may obtain a copy4# of the License at http://www.apache.org/licenses/LICENSE-2.05 6# Unless required by applicable law or agreed to in writing, software distributed under7# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS8# OF ANY KIND, either express or implied. See the License for the specific language9# governing permissions and limitations under the License.10 11# Patch for broken CTRL+C handler12# https://github.com/ContinuumIO/anaconda-issues/issues/90513import os14os.environ['FOR_DISABLE_CONSOLE_CTRL_HANDLER'] = '1'15 16import torch, json, numpy as np17from types import SimpleNamespace18import matplotlib.pyplot as plt19from pathlib import Path20from os import makedirs21from PIL import Image22from netdissect import proggan, nethook, easydict, zdataset23from netdissect.modelconfig import create_instrumented_model24from estimators import get_estimator25from models import get_instrumented_model26from scipy.cluster.vq import kmeans27import re28import sys29import datetime30import argparse31from tqdm import trange32from config import Config33from decomposition import get_random_dirs, get_or_compute, get_max_batch_size, SEED_VISUALIZATION34from utils import pad_frames 35 36def x_closest(p):37 distances = np.sqrt(np.sum((X - p)**2, axis=-1))38 idx = np.argmin(distances)39 return distances[idx], X[idx]40 41def make_gif(imgs, duration_secs, outname):42 head, *tail = [Image.fromarray((x * 255).astype(np.uint8)) for x in imgs]43 ms_per_frame = 1000 * duration_secs / instances44 head.save(outname, format='GIF', append_images=tail, save_all=True, duration=ms_per_frame, loop=0)45 46def make_mp4(imgs, duration_secs, outname):47 import shutil48 import subprocess as sp49 50 FFMPEG_BIN = shutil.which("ffmpeg")51 assert FFMPEG_BIN is not None, 'ffmpeg not found, install with "conda install -c conda-forge ffmpeg"'52 assert len(imgs[0].shape) == 3, 'Invalid shape of frame data'53 54 resolution = imgs[0].shape[0:2]55 fps = int(len(imgs) / duration_secs)56 57 command = [ FFMPEG_BIN,58 '-y', # overwrite output file59 '-f', 'rawvideo',60 '-vcodec','rawvideo',61 '-s', f'{resolution[0]}x{resolution[1]}', # size of one frame62 '-pix_fmt', 'rgb24',63 '-r', f'{fps}',64 '-i', '-', # imput from pipe65 '-an', # no audio66 '-c:v', 'libx264',67 '-preset', 'slow',68 '-crf', '17',69 str(Path(outname).with_suffix('.mp4')) ]70 71 frame_data = np.concatenate([(x * 255).astype(np.uint8).reshape(-1) for x in imgs])72 with sp.Popen(command, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE) as p:73 ret = p.communicate(frame_data.tobytes())74 if p.returncode != 0:75 print(ret[1].decode("utf-8"))76 raise sp.CalledProcessError(p.returncode, command)77 78 79def make_grid(latent, lat_mean, lat_comp, lat_stdev, act_mean, act_comp, act_stdev, scale=1, n_rows=10, n_cols=5, make_plots=True, edit_type='latent'):80 from notebooks.notebook_utils import create_strip_centered81 82 inst.remove_edits()83 x_range = np.linspace(-scale, scale, n_cols, dtype=np.float32) # scale in sigmas84 85 rows = []86 for r in range(n_rows):87 curr_row = []88 out_batch = create_strip_centered(inst, edit_type, layer_key, [latent],89 act_comp[r], lat_comp[r], act_stdev[r], lat_stdev[r], act_mean, lat_mean, scale, 0, -1, n_cols)[0]90 for i, img in enumerate(out_batch):91 curr_row.append(('c{}_{:.2f}'.format(r, x_range[i]), img))92 93 rows.append(curr_row[:n_cols])94 95 inst.remove_edits()96 97 if make_plots:98 # If more rows than columns, make several blocks side by side99 n_blocks = 2 if n_rows > n_cols else 1100 101 for r, data in enumerate(rows):102 # Add white borders103 imgs = pad_frames([img for _, img in data]) 104 105 coord = ((r * n_blocks) % n_rows) + ((r * n_blocks) // n_rows)106 plt.subplot(n_rows//n_blocks, n_blocks, 1 + coord)107 plt.imshow(np.hstack(imgs))108 109 # Custom x-axis labels110 W = imgs[0].shape[1] # image width111 P = imgs[1].shape[1] # padding width112 locs = [(0.5*W + i*(W+P)) for i in range(n_cols)]113 plt.xticks(locs, ["{:.2f}".format(v) for v in x_range])114 plt.yticks([])115 plt.ylabel(f'C{r}')116 117 plt.tight_layout()118 plt.subplots_adjust(top=0.96) # make room for suptitle119 120 return [img for row in rows for img in row]121 122 123######################124### Visualize results125######################126 127if __name__ == '__main__':128 global max_batch, sample_shape, feature_shape, inst, args, layer_key, model129 130 args = Config().from_args()131 t_start = datetime.datetime.now()132 timestamp = lambda : datetime.datetime.now().strftime("%d.%m %H:%M")133 print(f'[{timestamp()}] {args.model}, {args.layer}, {args.estimator}')134 135 # Ensure reproducibility136 torch.manual_seed(0) # also sets cuda seeds137 np.random.seed(0)138 139 # Speed up backend140 torch.backends.cudnn.benchmark = True141 torch.autograd.set_grad_enabled(False)142 143 has_gpu = torch.cuda.is_available()144 device = torch.device('cuda' if has_gpu else 'cpu')145 layer_key = args.layer146 layer_name = layer_key #layer_key.lower().split('.')[-1]147 148 basedir = Path(__file__).parent.resolve()149 outdir = basedir / 'out'150 151 # Load model152 inst = get_instrumented_model(args.model, args.output_class, layer_key, device, use_w=args.use_w)153 model = inst.model154 feature_shape = inst.feature_shape[layer_key]155 latent_shape = model.get_latent_shape()156 print('Feature shape:', feature_shape)157 158 # Layout of activations159 if len(feature_shape) != 4: # non-spatial160 axis_mask = np.ones(len(feature_shape), dtype=np.int32)161 else:162 axis_mask = np.array([0, 1, 1, 1]) # only batch fixed => whole activation volume used163 164 # Shape of sample passed to PCA165 sample_shape = feature_shape*axis_mask166 sample_shape[sample_shape == 0] = 1167 168 # Load or compute components169 dump_name = get_or_compute(args, inst)170 data = np.load(dump_name, allow_pickle=False) # does not contain object arrays171 X_comp = data['act_comp']172 X_global_mean = data['act_mean']173 X_stdev = data['act_stdev']174 X_var_ratio = data['var_ratio']175 X_stdev_random = data['random_stdevs']176 Z_global_mean = data['lat_mean']177 Z_comp = data['lat_comp']178 Z_stdev = data['lat_stdev']179 n_comp = X_comp.shape[0]180 data.close()181 182 # Transfer components to device183 tensors = SimpleNamespace(184 X_comp = torch.from_numpy(X_comp).to(device).float(), #-1, 1, C, H, W185 X_global_mean = torch.from_numpy(X_global_mean).to(device).float(), # 1, C, H, W186 X_stdev = torch.from_numpy(X_stdev).to(device).float(),187 Z_comp = torch.from_numpy(Z_comp).to(device).float(),188 Z_stdev = torch.from_numpy(Z_stdev).to(device).float(),189 Z_global_mean = torch.from_numpy(Z_global_mean).to(device).float(),190 )191 192 transformer = get_estimator(args.estimator, n_comp, args.sparsity)193 tr_param_str = transformer.get_param_str()194 195 # Compute max batch size given VRAM usage196 max_batch = args.batch_size or (get_max_batch_size(inst, device) if has_gpu else 1)197 print('Batch size:', max_batch)198 199 def show():200 if args.batch_mode:201 plt.close('all')202 else:203 plt.show()204 205 print(f'[{timestamp()}] Creating visualizations')206 207 # Ensure visualization gets new samples208 torch.manual_seed(SEED_VISUALIZATION)209 np.random.seed(SEED_VISUALIZATION)210 211 # Make output directories212 est_id = f'spca_{args.sparsity}' if args.estimator == 'spca' else args.estimator213 outdir_comp = outdir/model.name/layer_key.lower()/est_id/'comp'214 outdir_inst = outdir/model.name/layer_key.lower()/est_id/'inst'215 outdir_summ = outdir/model.name/layer_key.lower()/est_id/'summ'216 makedirs(outdir_comp, exist_ok=True)217 makedirs(outdir_inst, exist_ok=True)218 makedirs(outdir_summ, exist_ok=True)219 220 # Measure component sparsity (!= activation sparsity)221 sparsity = np.mean(X_comp == 0) # percentage of zero values in components222 print(f'Sparsity: {sparsity:.2f}')223 224 def get_edit_name(mode):225 if mode == 'activation':226 is_stylegan = 'StyleGAN' in args.model227 is_w = layer_key in ['style', 'g_mapping']228 return 'W' if (is_stylegan and is_w) else 'ACT'229 elif mode == 'latent':230 return model.latent_space_name()231 elif mode == 'both':232 return 'BOTH'233 else:234 raise RuntimeError(f'Unknown edit mode {mode}')235 236 # Only visualize applicable edit modes237 if args.use_w and layer_key in ['style', 'g_mapping']:238 edit_modes = ['latent'] # activation edit is the same239 else:240 edit_modes = ['activation', 'latent']241 242 # Summary grid, real components243 for edit_mode in edit_modes:244 plt.figure(figsize = (14,12))245 plt.suptitle(f"{args.estimator.upper()}: {model.name} - {layer_name}, {get_edit_name(edit_mode)} edit", size=16)246 make_grid(tensors.Z_global_mean, tensors.Z_global_mean, tensors.Z_comp, tensors.Z_stdev, tensors.X_global_mean,247 tensors.X_comp, tensors.X_stdev, scale=args.sigma, edit_type=edit_mode, n_rows=14)248 plt.savefig(outdir_summ / f'components_{get_edit_name(edit_mode)}.jpg', dpi=300)249 show()250 251 if args.make_video:252 components = 15253 instances = 150254 255 # One reasonable, one over the top256 for sigma in [args.sigma, 3*args.sigma]:257 for c in range(components):258 for edit_mode in edit_modes:259 frames = make_grid(tensors.Z_global_mean, tensors.Z_global_mean, tensors.Z_comp[c:c+1, :, :], tensors.Z_stdev[c:c+1], tensors.X_global_mean,260 tensors.X_comp[c:c+1, :, :], tensors.X_stdev[c:c+1], n_rows=1, n_cols=instances, scale=sigma, make_plots=False, edit_type=edit_mode)261 plt.close('all')262 263 frames = [x for _, x in frames]264 frames = frames + frames[::-1]265 make_mp4(frames, 5, outdir_comp / f'{get_edit_name(edit_mode)}_sigma{sigma}_comp{c}.mp4')266 267 268 # Summary grid, random directions269 # Using the stdevs of the principal components for same norm270 random_dirs_act = torch.from_numpy(get_random_dirs(n_comp, np.prod(sample_shape)).reshape(-1, *sample_shape)).to(device)271 random_dirs_z = torch.from_numpy(get_random_dirs(n_comp, np.prod(inst.input_shape)).reshape(-1, *latent_shape)).to(device)272 273 for edit_mode in edit_modes:274 plt.figure(figsize = (14,12))275 plt.suptitle(f"{model.name} - {layer_name}, random directions w/ PC stdevs, {get_edit_name(edit_mode)} edit", size=16)276 make_grid(tensors.Z_global_mean, tensors.Z_global_mean, random_dirs_z, tensors.Z_stdev,277 tensors.X_global_mean, random_dirs_act, tensors.X_stdev, scale=args.sigma, edit_type=edit_mode, n_rows=14)278 plt.savefig(outdir_summ / f'random_dirs_{get_edit_name(edit_mode)}.jpg', dpi=300)279 show()280 281 # Random instances w/ components added282 n_random_imgs = 10283 latents = model.sample_latent(n_samples=n_random_imgs)284 285 for img_idx in trange(n_random_imgs, desc='Random images', ascii=True):286 #print(f'Creating visualizations for random image {img_idx+1}/{n_random_imgs}')287 z = latents[img_idx][None, ...]288 289 # Summary grid, real components290 for edit_mode in edit_modes:291 plt.figure(figsize = (14,12))292 plt.suptitle(f"{args.estimator.upper()}: {model.name} - {layer_name}, {get_edit_name(edit_mode)} edit", size=16)293 make_grid(z, tensors.Z_global_mean, tensors.Z_comp, tensors.Z_stdev,294 tensors.X_global_mean, tensors.X_comp, tensors.X_stdev, scale=args.sigma, edit_type=edit_mode, n_rows=14)295 plt.savefig(outdir_summ / f'samp{img_idx}_real_{get_edit_name(edit_mode)}.jpg', dpi=300)296 show()297 298 if args.make_video:299 components = 5300 instances = 150301 302 # One reasonable, one over the top303 for sigma in [args.sigma, 3*args.sigma]: #[2, 5]:304 for edit_mode in edit_modes:305 imgs = make_grid(z, tensors.Z_global_mean, tensors.Z_comp, tensors.Z_stdev, tensors.X_global_mean, tensors.X_comp, tensors.X_stdev,306 n_rows=components, n_cols=instances, scale=sigma, make_plots=False, edit_type=edit_mode)307 plt.close('all')308 309 for c in range(components):310 frames = [x for _, x in imgs[c*instances:(c+1)*instances]]311 frames = frames + frames[::-1]312 make_mp4(frames, 5, outdir_inst / f'{get_edit_name(edit_mode)}_sigma{sigma}_img{img_idx}_comp{c}.mp4')313 314 print('Done in', datetime.datetime.now() - t_start)