Upload 7 files

scripts/align_all_parallel.py

@@ -0,0 +1,215 @@
+"""
+brief: face alignment with FFHQ method (https://github.com/NVlabs/ffhq-dataset)
+author: lzhbrian (https://lzhbrian.me)
+date: 2020.1.5
+note: code is heavily borrowed from 
+    https://github.com/NVlabs/ffhq-dataset
+    http://dlib.net/face_landmark_detection.py.html
+
+requirements:
+    apt install cmake
+    conda install Pillow numpy scipy
+    pip install dlib
+    # download face landmark model from:
+    # http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
+"""
+from argparse import ArgumentParser
+import time
+import numpy as np
+import PIL
+import PIL.Image
+import os
+import scipy
+import scipy.ndimage
+import dlib
+import multiprocessing as mp
+import math
+
+from configs.paths_config import model_paths
+SHAPE_PREDICTOR_PATH = model_paths["shape_predictor"]
+
+
+def get_landmark(filepath, predictor):
+    """get landmark with dlib
+    :return: np.array shape=(68, 2)
+    """
+    detector = dlib.get_frontal_face_detector()
+    if type(filepath) == str:
+        img = dlib.load_rgb_image(filepath)
+    else:
+        img = filepath
+    dets = detector(img, 1)
+
+    if len(dets) == 0:
+        print('Error: no face detected! If you are sure there are faces in your input, you may rerun the code or change the image several times until the face is detected. Sometimes the detector is unstable.')
+        return None
+
+    shape = None
+    for k, d in enumerate(dets):
+        shape = predictor(img, d)
+
+    t = list(shape.parts())
+    a = []
+    for tt in t:
+        a.append([tt.x, tt.y])
+    lm = np.array(a)
+    return lm
+
+
+def align_face(filepath, predictor):
+    """
+    :param filepath: str
+    :return: PIL Image
+    """
+
+    lm = get_landmark(filepath, predictor)
+    if lm is None:
+        return None    
+    
+    lm_chin = lm[0: 17]  # left-right
+    lm_eyebrow_left = lm[17: 22]  # left-right
+    lm_eyebrow_right = lm[22: 27]  # left-right
+    lm_nose = lm[27: 31]  # top-down
+    lm_nostrils = lm[31: 36]  # top-down
+    lm_eye_left = lm[36: 42]  # left-clockwise
+    lm_eye_right = lm[42: 48]  # left-clockwise
+    lm_mouth_outer = lm[48: 60]  # left-clockwise
+    lm_mouth_inner = lm[60: 68]  # left-clockwise
+
+    # Calculate auxiliary vectors.
+    eye_left = np.mean(lm_eye_left, axis=0)
+    eye_right = np.mean(lm_eye_right, axis=0)
+    eye_avg = (eye_left + eye_right) * 0.5
+    eye_to_eye = eye_right - eye_left
+    mouth_left = lm_mouth_outer[0]
+    mouth_right = lm_mouth_outer[6]
+    mouth_avg = (mouth_left + mouth_right) * 0.5
+    eye_to_mouth = mouth_avg - eye_avg
+
+    # Choose oriented crop rectangle.
+    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+    x /= np.hypot(*x)
+    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+    y = np.flipud(x) * [-1, 1]
+    c = eye_avg + eye_to_mouth * 0.1
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    qsize = np.hypot(*x) * 2
+
+    # read image
+    if type(filepath) == str:
+        img = PIL.Image.open(filepath)
+    else:
+        img = PIL.Image.fromarray(filepath)
+
+    output_size = 256
+    transform_size = 256
+    enable_padding = True
+
+    # Shrink.
+    shrink = int(np.floor(qsize / output_size * 0.5))
+    if shrink > 1:
+        rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
+        img = img.resize(rsize, PIL.Image.ANTIALIAS)
+        quad /= shrink
+        qsize /= shrink
+
+    # Crop.
+    border = max(int(np.rint(qsize * 0.1)), 3)
+    crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+            int(np.ceil(max(quad[:, 1]))))
+    crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]),
+            min(crop[3] + border, img.size[1]))
+    if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+        img = img.crop(crop)
+        quad -= crop[0:2]
+
+    # Pad.
+    pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+           int(np.ceil(max(quad[:, 1]))))
+    pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0),
+           max(pad[3] - img.size[1] + border, 0))
+    if enable_padding and max(pad) > border - 4:
+        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+        h, w, _ = img.shape
+        y, x, _ = np.ogrid[:h, :w, :1]
+        mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+                          1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
+        blur = qsize * 0.02
+        img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
+        quad += pad[:2]
+
+    # Transform.
+    img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR)
+    if output_size < transform_size:
+        img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
+
+    # Save aligned image.
+    return img
+
+
+def chunks(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i:i + n]
+
+
+def extract_on_paths(file_paths):
+    predictor = dlib.shape_predictor(SHAPE_PREDICTOR_PATH)
+    pid = mp.current_process().name
+    print('\t{} is starting to extract on #{} images'.format(pid, len(file_paths)))
+    tot_count = len(file_paths)
+    count = 0
+    for file_path, res_path in file_paths:
+        count += 1
+        if count % 100 == 0:
+            print('{} done with {}/{}'.format(pid, count, tot_count))
+        try:
+            res = align_face(file_path, predictor)
+            res = res.convert('RGB')
+            os.makedirs(os.path.dirname(res_path), exist_ok=True)
+            res.save(res_path)
+        except Exception:
+            continue
+    print('\tDone!')
+
+
+def parse_args():
+    parser = ArgumentParser(add_help=False)
+    parser.add_argument('--num_threads', type=int, default=1)
+    parser.add_argument('--root_path', type=str, default='')
+    args = parser.parse_args()
+    return args
+
+
+def run(args):
+    root_path = args.root_path
+    out_crops_path = root_path + '_crops'
+    if not os.path.exists(out_crops_path):
+        os.makedirs(out_crops_path, exist_ok=True)
+
+    file_paths = []
+    for root, dirs, files in os.walk(root_path):
+        for file in files:
+            file_path = os.path.join(root, file)
+            fname = os.path.join(out_crops_path, os.path.relpath(file_path, root_path))
+            res_path = '{}.jpg'.format(os.path.splitext(fname)[0])
+            if os.path.splitext(file_path)[1] == '.txt' or os.path.exists(res_path):
+                continue
+            file_paths.append((file_path, res_path))
+
+    file_chunks = list(chunks(file_paths, int(math.ceil(len(file_paths) / args.num_threads))))
+    print(len(file_chunks))
+    pool = mp.Pool(args.num_threads)
+    print('Running on {} paths\nHere we goooo'.format(len(file_paths)))
+    tic = time.time()
+    pool.map(extract_on_paths, file_chunks)
+    toc = time.time()
+    print('Mischief managed in {}s'.format(toc - tic))
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    run(args)

scripts/calc_id_loss_parallel.py

@@ -0,0 +1,119 @@
+from argparse import ArgumentParser
+import time
+import numpy as np
+import os
+import json
+import sys
+from PIL import Image
+import multiprocessing as mp
+import math
+import torch
+import torchvision.transforms as trans
+
+sys.path.append(".")
+sys.path.append("..")
+
+from models.mtcnn.mtcnn import MTCNN
+from models.encoders.model_irse import IR_101
+from configs.paths_config import model_paths
+CIRCULAR_FACE_PATH = model_paths['circular_face']
+
+
+def chunks(lst, n):
+	"""Yield successive n-sized chunks from lst."""
+	for i in range(0, len(lst), n):
+		yield lst[i:i + n]
+
+
+def extract_on_paths(file_paths):
+	facenet = IR_101(input_size=112)
+	facenet.load_state_dict(torch.load(CIRCULAR_FACE_PATH))
+	facenet.cuda()
+	facenet.eval()
+	mtcnn = MTCNN()
+	id_transform = trans.Compose([
+		trans.ToTensor(),
+		trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+	])
+
+	pid = mp.current_process().name
+	print('\t{} is starting to extract on {} images'.format(pid, len(file_paths)))
+	tot_count = len(file_paths)
+	count = 0
+
+	scores_dict = {}
+	for res_path, gt_path in file_paths:
+		count += 1
+		if count % 100 == 0:
+			print('{} done with {}/{}'.format(pid, count, tot_count))
+		if True:
+			input_im = Image.open(res_path)
+			input_im, _ = mtcnn.align(input_im)
+			if input_im is None:
+				print('{} skipping {}'.format(pid, res_path))
+				continue
+
+			input_id = facenet(id_transform(input_im).unsqueeze(0).cuda())[0]
+
+			result_im = Image.open(gt_path)
+			result_im, _ = mtcnn.align(result_im)
+			if result_im is None:
+				print('{} skipping {}'.format(pid, gt_path))
+				continue
+
+			result_id = facenet(id_transform(result_im).unsqueeze(0).cuda())[0]
+			score = float(input_id.dot(result_id))
+			scores_dict[os.path.basename(gt_path)] = score
+
+	return scores_dict
+
+
+def parse_args():
+	parser = ArgumentParser(add_help=False)
+	parser.add_argument('--num_threads', type=int, default=4)
+	parser.add_argument('--data_path', type=str,  default='results')
+	parser.add_argument('--gt_path', type=str, default='gt_images')
+	args = parser.parse_args()
+	return args
+
+
+def run(args):
+	file_paths = []
+	for f in os.listdir(args.data_path):
+		image_path = os.path.join(args.data_path, f)
+		gt_path = os.path.join(args.gt_path, f)
+		if f.endswith(".jpg") or f.endswith('.png'):
+			file_paths.append([image_path, gt_path.replace('.png','.jpg')])
+
+	file_chunks = list(chunks(file_paths, int(math.ceil(len(file_paths) / args.num_threads))))
+	pool = mp.Pool(args.num_threads)
+	print('Running on {} paths\nHere we goooo'.format(len(file_paths)))
+
+	tic = time.time()
+	results = pool.map(extract_on_paths, file_chunks)
+	scores_dict = {}
+	for d in results:
+		scores_dict.update(d)
+
+	all_scores = list(scores_dict.values())
+	mean = np.mean(all_scores)
+	std = np.std(all_scores)
+	result_str = 'New Average score is {:.2f}+-{:.2f}'.format(mean, std)
+	print(result_str)
+
+	out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics')
+	if not os.path.exists(out_path):
+		os.makedirs(out_path)
+
+	with open(os.path.join(out_path, 'stat_id.txt'), 'w') as f:
+		f.write(result_str)
+	with open(os.path.join(out_path, 'scores_id.json'), 'w') as f:
+		json.dump(scores_dict, f)
+
+	toc = time.time()
+	print('Mischief managed in {}s'.format(toc - tic))
+
+
+if __name__ == '__main__':
+	args = parse_args()
+	run(args)

scripts/calc_losses_on_images.py

@@ -0,0 +1,84 @@
+from argparse import ArgumentParser
+import os
+import json
+import sys
+from tqdm import tqdm
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+
+sys.path.append(".")
+sys.path.append("..")
+
+from criteria.lpips.lpips import LPIPS
+from datasets.gt_res_dataset import GTResDataset
+
+
+def parse_args():
+	parser = ArgumentParser(add_help=False)
+	parser.add_argument('--mode', type=str, default='lpips', choices=['lpips', 'l2'])
+	parser.add_argument('--data_path', type=str, default='results')
+	parser.add_argument('--gt_path', type=str, default='gt_images')
+	parser.add_argument('--workers', type=int, default=4)
+	parser.add_argument('--batch_size', type=int, default=4)
+	args = parser.parse_args()
+	return args
+
+
+def run(args):
+
+	transform = transforms.Compose([transforms.Resize((256, 256)),
+									transforms.ToTensor(),
+									transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+
+	print('Loading dataset')
+	dataset = GTResDataset(root_path=args.data_path,
+	                       gt_dir=args.gt_path,
+						   transform=transform)
+
+	dataloader = DataLoader(dataset,
+	                        batch_size=args.batch_size,
+	                        shuffle=False,
+	                        num_workers=int(args.workers),
+	                        drop_last=True)
+
+	if args.mode == 'lpips':
+		loss_func = LPIPS(net_type='alex')
+	elif args.mode == 'l2':
+		loss_func = torch.nn.MSELoss()
+	else:
+		raise Exception('Not a valid mode!')
+	loss_func.cuda()
+
+	global_i = 0
+	scores_dict = {}
+	all_scores = []
+	for result_batch, gt_batch in tqdm(dataloader):
+		for i in range(args.batch_size):
+			loss = float(loss_func(result_batch[i:i+1].cuda(), gt_batch[i:i+1].cuda()))
+			all_scores.append(loss)
+			im_path = dataset.pairs[global_i][0]
+			scores_dict[os.path.basename(im_path)] = loss
+			global_i += 1
+
+	all_scores = list(scores_dict.values())
+	mean = np.mean(all_scores)
+	std = np.std(all_scores)
+	result_str = 'Average loss is {:.2f}+-{:.2f}'.format(mean, std)
+	print('Finished with ', args.data_path)
+	print(result_str)
+
+	out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics')
+	if not os.path.exists(out_path):
+		os.makedirs(out_path)
+
+	with open(os.path.join(out_path, 'stat_{}.txt'.format(args.mode)), 'w') as f:
+		f.write(result_str)
+	with open(os.path.join(out_path, 'scores_{}.json'.format(args.mode)), 'w') as f:
+		json.dump(scores_dict, f)
+
+
+if __name__ == '__main__':
+	args = parse_args()
+	run(args)

scripts/generate_sketch_data.py

@@ -0,0 +1,62 @@
+from torchvision import transforms
+from torchvision.utils import save_image
+from torch.utils.serialization import load_lua
+import os
+import cv2
+import numpy as np
+
+"""
+NOTE!: Must have torch==0.4.1 and torchvision==0.2.1
+The sketch simplification model (sketch_gan.t7) from Simo Serra et al. can be downloaded from their official implementation: 
+    https://github.com/bobbens/sketch_simplification
+"""
+
+
+def sobel(img):
+    opImgx = cv2.Sobel(img, cv2.CV_8U, 0, 1, ksize=3)
+    opImgy = cv2.Sobel(img, cv2.CV_8U, 1, 0, ksize=3)
+    return cv2.bitwise_or(opImgx, opImgy)
+
+
+def sketch(frame):
+    frame = cv2.GaussianBlur(frame, (3, 3), 0)
+    invImg = 255 - frame
+    edgImg0 = sobel(frame)
+    edgImg1 = sobel(invImg)
+    edgImg = cv2.addWeighted(edgImg0, 0.75, edgImg1, 0.75, 0)
+    opImg = 255 - edgImg
+    return opImg
+
+
+def get_sketch_image(image_path):
+    original = cv2.imread(image_path)
+    original = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
+    sketch_image = sketch(original)
+    return sketch_image[:, :, np.newaxis]
+
+
+use_cuda = True
+
+cache = load_lua("/path/to/sketch_gan.t7")
+model = cache.model
+immean = cache.mean
+imstd = cache.std
+model.evaluate()
+
+data_path = "/path/to/data/imgs"
+images = [os.path.join(data_path, f) for f in os.listdir(data_path)]
+
+output_dir = "/path/to/data/edges"
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+for idx, image_path in enumerate(images):
+    if idx % 50 == 0:
+        print("{} out of {}".format(idx, len(images)))
+    data = get_sketch_image(image_path)
+    data = ((transforms.ToTensor()(data) - immean) / imstd).unsqueeze(0)
+    if use_cuda:
+        pred = model.cuda().forward(data.cuda()).float()
+    else:
+        pred = model.forward(data)
+    save_image(pred[0], os.path.join(output_dir, "{}_edges.jpg".format(image_path.split("/")[-1].split('.')[0])))

scripts/inference.py

@@ -0,0 +1,136 @@
+import os
+from argparse import Namespace
+
+from tqdm import tqdm
+import time
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+import sys
+
+sys.path.append(".")
+sys.path.append("..")
+
+from configs import data_configs
+from datasets.inference_dataset import InferenceDataset
+from utils.common import tensor2im, log_input_image
+from options.test_options import TestOptions
+from models.psp import pSp
+
+
+def run():
+    test_opts = TestOptions().parse()
+
+    if test_opts.resize_factors is not None:
+        assert len(
+            test_opts.resize_factors.split(',')) == 1, "When running inference, provide a single downsampling factor!"
+        out_path_results = os.path.join(test_opts.exp_dir, 'inference_results',
+                                        'downsampling_{}'.format(test_opts.resize_factors))
+        out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled',
+                                        'downsampling_{}'.format(test_opts.resize_factors))
+    else:
+        out_path_results = os.path.join(test_opts.exp_dir, 'inference_results')
+        out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled')
+
+    os.makedirs(out_path_results, exist_ok=True)
+    os.makedirs(out_path_coupled, exist_ok=True)
+
+    # update test options with options used during training
+    ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
+    opts = ckpt['opts']
+    opts.update(vars(test_opts))
+    if 'learn_in_w' not in opts:
+        opts['learn_in_w'] = False
+    if 'output_size' not in opts:
+        opts['output_size'] = 1024
+    opts = Namespace(**opts)
+
+    net = pSp(opts)
+    net.eval()
+    net.cuda()
+
+    print('Loading dataset for {}'.format(opts.dataset_type))
+    dataset_args = data_configs.DATASETS[opts.dataset_type]
+    transforms_dict = dataset_args['transforms'](opts).get_transforms()
+    dataset = InferenceDataset(root=opts.data_path,
+                               transform=transforms_dict['transform_inference'],
+                               opts=opts)
+    dataloader = DataLoader(dataset,
+                            batch_size=opts.test_batch_size,
+                            shuffle=False,
+                            num_workers=int(opts.test_workers),
+                            drop_last=True)
+
+    if opts.n_images is None:
+        opts.n_images = len(dataset)
+
+    global_i = 0
+    global_time = []
+    for input_batch in tqdm(dataloader):
+        if global_i >= opts.n_images:
+            break
+        with torch.no_grad():
+            input_cuda = input_batch.cuda().float()
+            tic = time.time()
+            result_batch = run_on_batch(input_cuda, net, opts)
+            toc = time.time()
+            global_time.append(toc - tic)
+
+        for i in range(opts.test_batch_size):
+            result = tensor2im(result_batch[i])
+            im_path = dataset.paths[global_i]
+
+            if opts.couple_outputs or global_i % 100 == 0:
+                input_im = log_input_image(input_batch[i], opts)
+                resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)
+                if opts.resize_factors is not None:
+                    # for super resolution, save the original, down-sampled, and output
+                    source = Image.open(im_path)
+                    res = np.concatenate([np.array(source.resize(resize_amount)),
+                                          np.array(input_im.resize(resize_amount, resample=Image.NEAREST)),
+                                          np.array(result.resize(resize_amount))], axis=1)
+                else:
+                    # otherwise, save the original and output
+                    res = np.concatenate([np.array(input_im.resize(resize_amount)),
+                                          np.array(result.resize(resize_amount))], axis=1)
+                Image.fromarray(res).save(os.path.join(out_path_coupled, os.path.basename(im_path)))
+
+            im_save_path = os.path.join(out_path_results, os.path.basename(im_path))
+            Image.fromarray(np.array(result)).save(im_save_path)
+
+            global_i += 1
+
+    stats_path = os.path.join(opts.exp_dir, 'stats.txt')
+    result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time))
+    print(result_str)
+
+    with open(stats_path, 'w') as f:
+        f.write(result_str)
+
+
+def run_on_batch(inputs, net, opts):
+    if opts.latent_mask is None:
+        result_batch = net(inputs, randomize_noise=False, resize=opts.resize_outputs)
+    else:
+        latent_mask = [int(l) for l in opts.latent_mask.split(",")]
+        result_batch = []
+        for image_idx, input_image in enumerate(inputs):
+            # get latent vector to inject into our input image
+            vec_to_inject = np.random.randn(1, 512).astype('float32')
+            _, latent_to_inject = net(torch.from_numpy(vec_to_inject).to("cuda"),
+                                      input_code=True,
+                                      return_latents=True)
+            # get output image with injected style vector
+            res = net(input_image.unsqueeze(0).to("cuda").float(),
+                      latent_mask=latent_mask,
+                      inject_latent=latent_to_inject,
+                      alpha=opts.mix_alpha,
+                      resize=opts.resize_outputs)
+            result_batch.append(res)
+        result_batch = torch.cat(result_batch, dim=0)
+    return result_batch
+
+
+if __name__ == '__main__':
+    run()

scripts/style_mixing.py

@@ -0,0 +1,101 @@
+import os
+from argparse import Namespace
+
+from tqdm import tqdm
+import numpy as np
+from PIL import Image
+import torch
+from torch.utils.data import DataLoader
+import sys
+
+sys.path.append(".")
+sys.path.append("..")
+
+from configs import data_configs
+from datasets.inference_dataset import InferenceDataset
+from utils.common import tensor2im, log_input_image
+from options.test_options import TestOptions
+from models.psp import pSp
+
+
+def run():
+	test_opts = TestOptions().parse()
+
+	if test_opts.resize_factors is not None:
+		factors = test_opts.resize_factors.split(',')
+		assert len(factors) == 1, "When running inference, please provide a single downsampling factor!"
+		mixed_path_results = os.path.join(test_opts.exp_dir, 'style_mixing',
+		                                  'downsampling_{}'.format(test_opts.resize_factors))
+	else:
+		mixed_path_results = os.path.join(test_opts.exp_dir, 'style_mixing')
+	os.makedirs(mixed_path_results, exist_ok=True)
+
+	# update test options with options used during training
+	ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
+	opts = ckpt['opts']
+	opts.update(vars(test_opts))
+	if 'learn_in_w' not in opts:
+		opts['learn_in_w'] = False
+	if 'output_size' not in opts:
+		opts['output_size'] = 1024
+	opts = Namespace(**opts)
+
+	net = pSp(opts)
+	net.eval()
+	net.cuda()
+
+	print('Loading dataset for {}'.format(opts.dataset_type))
+	dataset_args = data_configs.DATASETS[opts.dataset_type]
+	transforms_dict = dataset_args['transforms'](opts).get_transforms()
+	dataset = InferenceDataset(root=opts.data_path,
+	                           transform=transforms_dict['transform_inference'],
+	                           opts=opts)
+	dataloader = DataLoader(dataset,
+	                        batch_size=opts.test_batch_size,
+	                        shuffle=False,
+	                        num_workers=int(opts.test_workers),
+	                        drop_last=True)
+
+	latent_mask = [int(l) for l in opts.latent_mask.split(",")]
+	if opts.n_images is None:
+		opts.n_images = len(dataset)
+
+	global_i = 0
+	for input_batch in tqdm(dataloader):
+		if global_i >= opts.n_images:
+			break
+		with torch.no_grad():
+			input_batch = input_batch.cuda()
+			for image_idx, input_image in enumerate(input_batch):
+				# generate random vectors to inject into input image
+				vecs_to_inject = np.random.randn(opts.n_outputs_to_generate, 512).astype('float32')
+				multi_modal_outputs = []
+				for vec_to_inject in vecs_to_inject:
+					cur_vec = torch.from_numpy(vec_to_inject).unsqueeze(0).to("cuda")
+					# get latent vector to inject into our input image
+					_, latent_to_inject = net(cur_vec,
+					                          input_code=True,
+					                          return_latents=True)
+					# get output image with injected style vector
+					res = net(input_image.unsqueeze(0).to("cuda").float(),
+					          latent_mask=latent_mask,
+					          inject_latent=latent_to_inject,
+					          alpha=opts.mix_alpha,
+							  resize=opts.resize_outputs)
+					multi_modal_outputs.append(res[0])
+
+				# visualize multi modal outputs
+				input_im_path = dataset.paths[global_i]
+				image = input_batch[image_idx]
+				input_image = log_input_image(image, opts)
+				resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)
+				res = np.array(input_image.resize(resize_amount))
+				for output in multi_modal_outputs:
+					output = tensor2im(output)
+					res = np.concatenate([res, np.array(output.resize(resize_amount))], axis=1)
+				Image.fromarray(res).save(os.path.join(mixed_path_results, os.path.basename(input_im_path)))
+				global_i += 1
+
+
+if __name__ == '__main__':
+	run()

scripts/train.py

@@ -0,0 +1,32 @@
+"""
+This file runs the main training/val loop
+"""
+import os
+import json
+import sys
+import pprint
+
+sys.path.append(".")
+sys.path.append("..")
+
+from options.train_options import TrainOptions
+from training.coach import Coach
+
+
+def main():
+	opts = TrainOptions().parse()
+	if os.path.exists(opts.exp_dir):
+		raise Exception('Oops... {} already exists'.format(opts.exp_dir))
+	os.makedirs(opts.exp_dir)
+
+	opts_dict = vars(opts)
+	pprint.pprint(opts_dict)
+	with open(os.path.join(opts.exp_dir, 'opt.json'), 'w') as f:
+		json.dump(opts_dict, f, indent=4, sort_keys=True)
+
+	coach = Coach(opts)
+	coach.train()
+
+
+if __name__ == '__main__':
+	main()