Пример #1
0
def create_from_image_folders(datadir, jpg=False, shuffle=True, size=None):
    assert os.path.isdir(datadir)
    imgs = []
    labels = []
    for root, subdirs, files in os.walk(datadir):
        for i, subdir in enumerate(subdirs):
            tmp_list = img_list(os.path.join(root, subdir))
            imgs = imgs + tmp_list
            labels = labels + [i] * len(tmp_list)
    labels = np.array(labels)
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.

    assert len(imgs) > 0, ' No input images found!'
    sample_img = np.asarray(PIL.Image.open(imgs[0]))
    sample_shape = sample_img.shape
    channels = sample_shape[2] if sample_img.ndim == 3 else 1
    assert channels in [1, 3, 4], ' Weird color dim: %d' % channels
    print(' Making dataset ..', datadir, sample_shape,
          '%d labels' % (np.max(labels) + 1))
    if jpg is True: print(' Loading JPG as is!')

    with TFRecordExporter(datadir, len(imgs)) as tfr:
        order = tfr.choose_shuffled_order() if shuffle else np.arange(
            len(imgs))
        pbar = ProgressBar(order.size)
        for idx in range(order.size):
            img_path = imgs[order[idx]]
            tfr.add_image(img_path, jpg=jpg, size=size)
            pbar.upd()
        tfr.add_labels(onehot[order])
    return tfr.tfr_file, len(imgs)
Пример #2
0
def make_transit(lat1, lat2, fcount, start_frame=0):
    pbar = ProgressBar(fcount)
    for i in range(fcount):
        img = render_latent_mix(lat1, lat2, i/fcount)
        fname = os.path.join(a.out_dir, 'ttt', "%05d.jpg" % (i+start_frame))
        if i%2==0:
            cv2.imshow('latent', img[:,:,::-1])
            cv2.waitKey(10)
        imsave(fname, img)
        pbar.upd()
Пример #3
0
def make_loop(base_latent, direction, lrange, fcount, start_frame=0):
    coeffs = get_coeffs_dir(lrange, fcount//2)
    pbar = ProgressBar(fcount)
    for i in range(fcount):
        img = render_latent_dir(base_latent, direction, coeffs[i])
        fname1 = os.path.join(a.out_dir, 'ttt', "%05d.jpg" % (i+start_frame))
        if i%2==0:
            cv2.imshow('latent', img[:,:,::-1])
            cv2.waitKey(10)
        imsave(fname1, img)
        pbar.upd()
Пример #4
0
def create_from_images(dataset, jpg=False, shuffle=True, size=None):
    assert os.path.isdir(dataset)
    image_filenames = sorted(img_list(dataset))
    assert len(image_filenames) > 0, ' No input images found!'

    sample_img = np.asarray(PIL.Image.open(image_filenames[0]))
    sample_shape = sample_img.shape
    channels = sample_shape[2] if sample_img.ndim == 3 else 1
    assert channels in [1,3,4], ' Weird color dim: %d' % channels
    print(' Making dataset', dataset, sample_shape)
    if jpg is True: print(' Loading JPG as is!')

    with TFRecordExporter(dataset, len(image_filenames)) as tfr:
        order = tfr.choose_shuffled_order() if shuffle else np.arange(len(image_filenames))
        pbar = ProgressBar(order.size)
        for idx in range(order.size):
            img_path = image_filenames[order[idx]]
            tfr.add_image(img_path, jpg=jpg, size=size)
            pbar.upd()
    return tfr.tfr_file, len(image_filenames)
Пример #5
0
def project_image(proj, targets, work_dir, resolution, num_snapshots):
    filename = osp.join(work_dir, basename(work_dir))
    video_out = cv2.VideoWriter(filename + '.avi',
                                cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 24,
                                resolution)

    snapshot_steps = set(proj.num_steps - np.linspace(
        0, proj.num_steps, num_snapshots, endpoint=False, dtype=int))
    misc.save_image_grid(targets, filename + '.jpg', drange=[-1, 1])
    proj.start(targets)
    pbar = ProgressBar(proj.num_steps)
    while proj.get_cur_step() < proj.num_steps:
        proj.step()
        write_video_frame(proj, video_out)
        if proj.get_cur_step() in snapshot_steps:
            misc.save_image_grid(proj.get_images(),
                                 filename + '-%04d.jpg' % proj.get_cur_step(),
                                 drange=[-1, 1])
        pbar.upd()

    dlats = proj.get_dlatents()
    np.save(filename + '-%04d.npy' % proj.get_cur_step(), dlats)
    video_out.release()
Пример #6
0
def main():
    os.makedirs(a.out_dir, exist_ok=True)

    # setup generator
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.func_name = 'training.stylegan2_multi.G_main'
    Gs_kwargs.verbose = a.verbose
    Gs_kwargs.size = a.size
    Gs_kwargs.scale_type = a.scale_type
    Gs_kwargs.impl = a.ops

    # load model with arguments
    sess = tflib.init_tf({'allow_soft_placement': True})
    pkl_name = osp.splitext(a.model)[0]
    with open(pkl_name + '.pkl', 'rb') as file:
        network = pickle.load(file, encoding='latin1')
    try:
        _, _, network = network
    except:
        pass
    for k in list(network.static_kwargs.keys()):
        Gs_kwargs[k] = network.static_kwargs[k]

    # reload custom network, if needed
    if '.pkl' in a.model.lower():
        print(' .. Gs from pkl ..', basename(a.model))
        Gs = network
    else:  # reconstruct network
        print(' .. Gs custom ..', basename(a.model))
        Gs = tflib.Network('Gs', **Gs_kwargs)
        Gs.copy_vars_from(network)

    z_dim = Gs.input_shape[1]
    dz_dim = 512  # dlatent_size
    try:
        dl_dim = 2 * (int(np.floor(np.log2(Gs_kwargs.resolution))) - 1)
    except:
        print(' Resave model, no resolution kwarg found!')
        exit(1)
    dlat_shape = (1, dl_dim, dz_dim)  # [1,18,512]

    # read saved latents
    if a.dlatents is not None and osp.isfile(a.dlatents):
        key_dlatents = load_latents(a.dlatents)
        if len(key_dlatents.shape) == 2:
            key_dlatents = np.expand_dims(key_dlatents, 0)
    elif a.dlatents is not None and osp.isdir(a.dlatents):
        # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1]
        key_dlatents = []
        npy_list = file_list(a.dlatents, 'npy')
        for npy in npy_list:
            key_dlatent = load_latents(npy)
            if len(key_dlatent.shape) == 2:
                key_dlatent = np.expand_dims(key_dlatent, 0)
            key_dlatents.append(key_dlatent)
        key_dlatents = np.concatenate(key_dlatents)  # [frm,18,512]
    else:
        print(' No input dlatents found')
        exit()
    key_dlatents = key_dlatents[:, np.newaxis]  # [frm,1,18,512]
    print(' key dlatents', key_dlatents.shape)

    # replace higher layers with single (style) latent
    if a.style_npy_file is not None:
        print(' styling with latent', a.style_npy_file)
        style_dlatent = load_latents(a.style_npy_file)
        while len(style_dlatent.shape) < 4:
            style_dlatent = np.expand_dims(style_dlatent, 0)
        # try replacing 5 by other value, less than dl_dim
        key_dlatents[:, :,
                     range(5, dl_dim), :] = style_dlatent[:, :,
                                                          range(5, dl_dim), :]

    frames = key_dlatents.shape[0] * a.fstep

    dlatents = latent_anima(dlat_shape,
                            frames,
                            a.fstep,
                            key_latents=key_dlatents,
                            cubic=a.cubic,
                            verbose=True)  # [frm,1,512]
    print(' dlatents', dlatents.shape)
    frame_count = dlatents.shape[0]

    # truncation trick
    dlatent_avg = Gs.get_var('dlatent_avg')  # (512,)
    tr_range = range(0, 8)
    dlatents[:, :, tr_range, :] = dlatent_avg + (dlatents[:, :, tr_range, :] -
                                                 dlatent_avg) * a.trunc

    # distort image by tweaking initial const layer
    if a.digress > 0:
        try:
            latent_size = Gs.static_kwargs['latent_size']
        except:
            latent_size = 512  # default latent size
        try:
            init_res = Gs.static_kwargs['init_res']
        except:
            init_res = (4, 4)  # default initial layer size
        dconst = a.digress * latent_anima([1, latent_size, *init_res],
                                          frames,
                                          a.fstep,
                                          cubic=True,
                                          verbose=False)
    else:
        dconst = np.zeros([frame_count, 1, 1, 1, 1])

    # generate images from latent timeline
    pbar = ProgressBar(frame_count)
    for i in range(frame_count):

        if a.digress is True:
            tf.get_default_session().run(tf.assign(wvars[0], wts[i]))

        # generate multi-latent result
        if Gs.num_inputs == 2:
            output = Gs.components.synthesis.run(dlatents[i],
                                                 randomize_noise=False,
                                                 output_transform=fmt,
                                                 minibatch_size=1)
        else:
            output = Gs.components.synthesis.run(dlatents[i], [None],
                                                 dconst[i],
                                                 randomize_noise=False,
                                                 output_transform=fmt,
                                                 minibatch_size=1)

        ext = 'png' if output.shape[3] == 4 else 'jpg'
        filename = osp.join(a.out_dir, "%06d.%s" % (i, ext))
        imsave(filename, output[0])
        pbar.upd()
Пример #7
0
def main():
    os.makedirs(a.out_dir, exist_ok=True)
    np.random.seed(seed=696)

    # parse filename to model parameters
    mparams = basename(a.model).split('-')
    res = int(mparams[1])
    cfg = mparams[2]
    
    # setup generator
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.func_name = 'training.stylegan2_custom.G_main'
    Gs_kwargs.verbose = False
    Gs_kwargs.resolution = res
    Gs_kwargs.size = a.size
    Gs_kwargs.scale_type = a.scale_type
    Gs_kwargs.latent_size = a.latent_size
    Gs_kwargs.impl = a.ops
    
    if cfg.lower() == 'f':
        Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2'
    elif cfg.lower() == 'e':
        Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2'
        Gs_kwargs.fmap_base = 8 << 10
    else:
        print(' old modes [A-D] not implemented'); exit()
    
    # check initial model resolution
    if len(mparams) > 3: 
        if 'x' in mparams[3].lower():
            init_res = [int(x) for x in mparams[3].lower().split('x')]
            Gs_kwargs.init_res = list(reversed(init_res)) # [H,W] 

    # load model, check channels
    sess = tflib.init_tf({'allow_soft_placement':True})
    pkl_name = osp.splitext(a.model)[0]
    with open(pkl_name + '.pkl', 'rb') as file:
        network = pickle.load(file, encoding='latin1')
    try: _, _, Gs = network
    except:    Gs = network
    Gs_kwargs.num_channels = Gs.output_shape[1]

    # reload custom network, if needed
    if '.pkl' in a.model.lower(): 
        print(' .. Gs from pkl ..')
    else: 
        print(' .. Gs custom ..')
        Gs = tflib.Network('Gs', **Gs_kwargs)
        Gs.copy_vars_from(network)
    # Gs.print_layers()
    print(' out shape', Gs.output_shape[1:])
    if a.size is None: a.size = Gs.output_shape[2:]

    z_dim = Gs.input_shape[1]
    shape = (1, z_dim)
    
    print(' making timeline..')
    latents = latent_anima(shape, a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, verbose=True) # [frm,1,512]
    print(' latents', latents.shape)
    
    # generate images from latent timeline
    frame_count = latents.shape[0]
    pbar = ProgressBar(frame_count)
    for i in range(frame_count):
        output = Gs.run(latents[i], [None], truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt)
        ext = 'png' if output.shape[3]==4 else 'jpg'
        filename = osp.join(a.out_dir, "%05d.%s" % (i,ext))
        imsave(filename, output[0])
        pbar.upd()

    # convert latents to dlatents, save them
    latents = latents.squeeze(1) # [frm,512]
    dlatents = Gs.components.mapping.run(latents, None, latent_size=z_dim, dtype='float16') # [frm,18,512]
    filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0])
    filename = osp.join(osp.dirname(a.out_dir), filename)
    np.save(filename, dlatents)
    print('saved dlatents', dlatents.shape, 'to', filename)
Пример #8
0
def main():
    os.makedirs(a.out_dir, exist_ok=True)

    # parse filename to model parameters
    mparams = basename(a.model).split('-')
    res = int(mparams[1])
    cfg = mparams[2]

    # setup generator
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.func_name = 'training.stylegan2_custom.G_main'
    Gs_kwargs.verbose = False
    Gs_kwargs.resolution = res
    Gs_kwargs.size = a.size
    Gs_kwargs.scale_type = a.scale_type
    Gs_kwargs.latent_size = a.latent_size

    if cfg.lower() == 'f':
        Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2'
    elif cfg.lower() == 'e':
        Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2'
        Gs_kwargs.fmap_base = 8 << 10
    else:
        print(' old modes [A-D] not implemented')
        exit()

    # check initial model resolution
    if len(mparams) > 3:
        if 'x' in mparams[3].lower():
            init_res = [int(x) for x in mparams[3].lower().split('x')]
            Gs_kwargs.init_res = list(reversed(init_res))  # [H,W]

    # load model, check channels
    sess = tflib.init_tf({'allow_soft_placement': True})
    pkl_name = osp.splitext(a.model)[0]
    with open(pkl_name + '.pkl', 'rb') as file:
        network = pickle.load(file, encoding='latin1')
    try:
        _, _, Gs = network
    except:
        Gs = network
    Gs_kwargs.num_channels = Gs.output_shape[1]

    # reload custom network, if needed
    if '.pkl' in a.model.lower():
        print(' .. Gs from pkl ..')
    else:
        print(' .. Gs custom ..')
        Gs = tflib.Network('Gs', **Gs_kwargs)
        Gs.copy_vars_from(network)

    z_dim = Gs.input_shape[1]
    dz_dim = a.dlatent_size  # 512
    dl_dim = 2 * (int(np.floor(np.log2(res))) - 1)
    dlat_shape = (1, dl_dim, dz_dim)  # [1,18,512]

    # read saved latents
    if a.dlatents is not None and osp.isfile(a.dlatents):
        key_dlatents = load_latents(a.dlatents)
        if len(key_dlatents.shape) == 2:
            key_dlatents = np.expand_dims(key_dlatents, 0)
    elif a.dlatents is not None and osp.isdir(a.dlatents):
        # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1]
        key_dlatents = []
        npy_list = file_list(a.dlatents, 'npy')
        for npy in npy_list:
            key_dlatent = load_latents(npy)
            if len(key_dlatent.shape) == 2:
                key_dlatent = np.expand_dims(key_dlatent, 0)
            key_dlatents.append(key_dlatent)
        key_dlatents = np.concatenate(key_dlatents)  # [frm,18,512]
    else:
        print(' No input dlatents found')
        exit()
    key_dlatents = key_dlatents[:, np.newaxis]  # [frm,1,18,512]
    print(' key dlatents', key_dlatents.shape)

    # replace higher layers with single (style) latent
    if a.style_npy_file is not None:
        print(' styling with latent', a.style_npy_file)
        style_dlatent = load_latents(a.style_npy_file)
        while len(style_dlatent.shape) < 4:
            style_dlatent = np.expand_dims(style_dlatent, 0)
        # try other values < dl_dim besides 5
        key_dlatents[:, :,
                     range(5, dl_dim), :] = style_dlatent[:, :,
                                                          range(5, dl_dim), :]

    frames = key_dlatents.shape[0] * a.fstep

    dlatents = latent_anima(dlat_shape,
                            frames,
                            a.fstep,
                            key_latents=key_dlatents,
                            cubic=a.cubic,
                            verbose=True)  # [frm,1,512]
    print(' dlatents', dlatents.shape)

    # truncation trick
    dlatent_avg = Gs.get_var('dlatent_avg')  # (512,)
    tr_range = range(0, 8)
    dlatents[:, :, tr_range, :] = dlatent_avg + (dlatents[:, :, tr_range, :] -
                                                 dlatent_avg) * a.trunc

    # loop for graph frame by frame
    frame_count = dlatents.shape[0]
    pbar = ProgressBar(frame_count)
    for i in range(frame_count):

        dlatent = dlatents[i]

        output = Gs.components.synthesis.run(dlatent,
                                             randomize_noise=False,
                                             output_transform=fmt,
                                             minibatch_size=1)

        ext = 'png' if output.shape[3] == 4 else 'jpg'
        filename = osp.join(a.out_dir, "%05d.%s" % (i, ext))
        imsave(filename, output[0])
        pbar.upd()
Пример #9
0
def main():
    os.makedirs(a.out_dir, exist_ok=True)
    np.random.seed(seed=696)

    # setup generator
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.func_name = 'training.stylegan2_multi.G_main'
    Gs_kwargs.verbose = a.verbose
    Gs_kwargs.size = a.size
    Gs_kwargs.scale_type = a.scale_type
    Gs_kwargs.impl = a.ops

    # mask/blend latents with external latmask or by splitting the frame
    if a.latmask is None:
        if a.verbose is True:
            print(' Latent blending w/split frame %d x %d' % (a.nX, a.nY))
        n_mult = a.nX * a.nY
        lmask = np.tile(np.asarray([[[[None]]]]), (1, n_mult, 1, 1))
        Gs_kwargs.countW = a.nX
        Gs_kwargs.countH = a.nY
        Gs_kwargs.splitfine = a.splitfine
    else:
        if a.verbose is True: print(' Latent blending with mask', a.latmask)
        n_mult = 2
        if os.path.isfile(a.latmask):  # single file
            lmask = np.asarray([[imread(a.latmask)[:, :, 0] / 255.]])  # [h,w]
        elif os.path.isdir(a.latmask):  # directory with frame sequence
            lmask = np.asarray([[
                imread(f)[:, :, 0] / 255. for f in img_list(a.latmask)
            ]])  # [h,w]
        else:
            print(' !! Blending mask not found:', a.latmask)
            exit(1)
        lmask = np.concatenate((lmask, 1 - lmask), 1)  # [frm,2,h,w]
        Gs_kwargs.latmask_res = lmask.shape[2:]

    # load model with arguments
    sess = tflib.init_tf({'allow_soft_placement': True})
    pkl_name = osp.splitext(a.model)[0]
    with open(pkl_name + '.pkl', 'rb') as file:
        network = pickle.load(file, encoding='latin1')
    try:
        _, _, network = network
    except:
        pass
    for k in list(network.static_kwargs.keys()):
        Gs_kwargs[k] = network.static_kwargs[k]

    # reload custom network, if needed
    if '.pkl' in a.model.lower():
        print(' .. Gs from pkl ..', basename(a.model))
        Gs = network
    else:  # reconstruct network
        print(' .. Gs custom ..', basename(a.model))
        # print(Gs_kwargs)
        Gs = tflib.Network('Gs', **Gs_kwargs)
        Gs.copy_vars_from(network)
    if a.verbose is True:
        print('kwargs:', [
            '%s: %s' % (kv[0], kv[1])
            for kv in sorted(Gs.static_kwargs.items())
        ])

    if a.verbose is True: print(' out shape', Gs.output_shape[1:])
    if a.size is None: a.size = Gs.output_shape[2:]

    if a.verbose is True: print(' making timeline..')
    lats = []  # list of [frm,1,512]
    for i in range(n_mult):
        lat_tmp = latent_anima((1, Gs.input_shape[1]),
                               a.frames,
                               a.fstep,
                               cubic=a.cubic,
                               gauss=a.gauss,
                               verbose=False)  # [frm,1,512]
        lats.append(lat_tmp)  # list of [frm,1,512]
    latents = np.concatenate(lats, 1)  # [frm,X,512]
    print(' latents', latents.shape)
    frame_count = latents.shape[0]

    # distort image by tweaking initial const layer
    if a.digress > 0:
        try:
            latent_size = Gs.static_kwargs['latent_size']
        except:
            latent_size = 512  # default latent size
        try:
            init_res = Gs.static_kwargs['init_res']
        except:
            init_res = (4, 4)  # default initial layer size
        dconst = []
        for i in range(n_mult):
            dc_tmp = a.digress * latent_anima([1, latent_size, *init_res],
                                              a.frames,
                                              a.fstep,
                                              cubic=True,
                                              verbose=False)
            dconst.append(dc_tmp)
        dconst = np.concatenate(dconst, 1)
    else:
        dconst = np.zeros([frame_count, 1, 1, 1, 1])

    # generate images from latent timeline
    pbar = ProgressBar(frame_count)
    for i in range(frame_count):

        latent = latents[i]  # [X,512]
        latmask = lmask[i %
                        len(lmask)] if lmask is not None else [None]  # [X,h,w]
        dc = dconst[i % len(dconst)]  # [X,512,4,4]

        # generate multi-latent result
        if Gs.num_inputs == 2:
            output = Gs.run(latent, [None],
                            truncation_psi=a.trunc,
                            randomize_noise=False,
                            output_transform=fmt)
        else:
            output = Gs.run(latent, [None],
                            latmask,
                            dc,
                            truncation_psi=a.trunc,
                            randomize_noise=False,
                            output_transform=fmt)

        # save image
        ext = 'png' if output.shape[3] == 4 else 'jpg'
        filename = osp.join(a.out_dir, "%06d.%s" % (i, ext))
        imsave(filename, output[0])
        pbar.upd()

    # convert latents to dlatents, save them
    if a.save_lat is True:
        latents = latents.squeeze(1)  # [frm,512]
        dlatents = Gs.components.mapping.run(latents, None,
                                             dtype='float16')  # [frm,18,512]
        filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1],
                                         a.size[0])
        filename = osp.join(osp.dirname(a.out_dir), filename)
        np.save(filename, dlatents)
        print('saved dlatents', dlatents.shape, 'to', filename)