generate_conditional_mass.py

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Project given image to the latent space of pretrained network pickle."""

from typing import List
from utils import load_model
from interpolation import mapping
import os
from time import time
from training.training_loop import save_image_grid
import pathlib
from generate_grid import condition_list, num_range, ModifiedPath
import json

import click
import torch
from training.dataset import ImageFolderDataset


#----------------------------------------------------------------------------

#----------------------------------------------------------------------------

@click.command()
@click.option('--network', 'network_pkl', help='Network pickle filename')
@click.option('--num-samples',            type=num_range, help='Number of samples to use for drawing')
@click.option('--seeds',                  type=num_range, help='List of random seeds')
@click.option('--outdir',                 help='Where to save the output images', required=True, metavar='DIR')
@click.option('--massive-multi-domain-conditions', '--mmdc', help="Specify the directory where the MMDC annotations (generated by the `create_label_json.py` script) are stored."
              "If this option is not provided, MMDC is not used. If provided without a value, a default value is used.", is_flag=False, flag_value="./annotations/emotions-artist-style-genre", type=ModifiedPath(file_okay=False, exists=True, path_type=pathlib.Path), default=False, is_eager=True)
@click.option('--conditions',             help='Condition from', type=condition_list)
@click.option('--use-wandb', '--wandb',   help='run name in our wandb project')
def conditional_mass(
    network_pkl: str,
    outdir: str,
    seeds: List[int],
    num_samples: List[int],
    massive_multi_domain_conditions: ModifiedPath,
    conditions,
    use_wandb: str,
):
    G, device = load_model(network_pkl, use_wandb)

    # <condition_key>=<from_cond>,<to_cond>
    assert len(conditions) == 1

    with open(massive_multi_domain_conditions / "prepared_dataset.json") as f:
        data = json.load(f)
        label_shapes = data["shapes"]
    
    embedded_labels = ImageFolderDataset.transform_multidomain_conditions(conditions, label_shapes)

    os.makedirs(outdir, exist_ok=True)
    timestamp = int(time())
    imgs_outdir = pathlib.Path(outdir) / str(timestamp)
    os.makedirs(imgs_outdir)
    num_examples = 5

    for seed in seeds:
        mapping_kwargs = dict(G=G, seed=seed, device=device)
        ws = mapping(condition=embedded_labels[0], num=num_examples, **mapping_kwargs)
        ws = torch.from_numpy(ws).to(device)
        imgs = G.synthesis(ws, noise_mode="const")
        save_image_grid(imgs, imgs_outdir / f'samples-seed{seed:04d}.jpg', drange=[-1,1], grid_size=(num_examples,1))

        for num in num_samples:
            
            ws = mapping(condition=embedded_labels[0], num=num, **mapping_kwargs)
            print(ws.shape)

            ws = ws.mean(axis=0, keepdims=True)
            ws = torch.from_numpy(ws).to(device)
            img = G.synthesis(ws, noise_mode="const")
            save_image_grid(img, imgs_outdir / f'avg-seed{seed:04d}-{num}.jpg', drange=[-1,1], grid_size=(1,1))


#----------------------------------------------------------------------------

if __name__ == "__main__":
    conditional_mass() # pylint: disable=no-value-for-parameter

#----------------------------------------------------------------------------