RameenAbdal / StyleFlow

StyleFlow: Attribute-conditioned Exploration of StyleGAN-generated Images using Conditional Continuous Normalizing Flows (ACM TOG 2021)
http://rameenabdal.github.io/StyleFlow
2.42k stars 342 forks source link

[Q]. How to only change face pose? #120

Closed univ-esuty closed 2 years ago

univ-esuty commented 2 years ago

I have a single font face picture for input image. How can I change pose continually by using run_generator.py? which parts of the code should I change?

FeiiYin commented 2 years ago

@univ-esuty Hi, I face the problems either, can you provide your solution or share your code? That will be a great help for me! Thanks in advance!

univ-esuty commented 2 years ago

I'm sorry to reply late. I found a simple solution.

I'll share you my code. Select a seed value first and call change function. The input value attr_index corresponds to attr_order .

from utils import Build_model
from options.test_options import TestOptions
from module.flow import cnf

from ui.ui import  attr_degree_list

import numpy as np
import PIL
import pickle
import torch
import random
import os
from tqdm import tqdm

keep_indexes = [i for i in range(0,100)]
keep_indexes = np.array(keep_indexes).astype(np.int)

raw_w = pickle.load(open("data/sg2latents.pickle", "rb"))
raw_TSNE = np.load('data/TSNE.npy')
raw_attr = np.load('data/attributes.npy')
raw_lights2 = np.load('data/light.npy')

raw_lights = raw_lights2
all_w = np.array(raw_w['Latent'])[keep_indexes]
all_attr = raw_attr[keep_indexes]
all_lights = raw_lights[keep_indexes]

IDX = 6 ## SEED value
os.makedirs(f'outputs/{IDX}')

w_current = all_w[IDX].copy()
attr_current = all_attr[IDX].copy()
light_current = all_lights[IDX].copy()

## load pre-trained network.
print('loading ...')
opt = TestOptions().parse()
model = Build_model(opt)
w_avg = model.Gs.get_var('dlatent_avg')
prior = cnf(512, '512-512-512-512-512', 17, 1)
prior.load_state_dict(torch.load('flow_weight/modellarge10k.pt'))
prior.eval()

## editor setting
def change(w_before, attr_index, real_value):
    zero_padding = torch.zeros(1, 18, 1).cuda()
    q_array = torch.from_numpy(w_before).cuda().clone().detach()
    array_source = torch.from_numpy(attr_current).type(torch.FloatTensor).cuda()
    array_light = torch.from_numpy(light_current).type(torch.FloatTensor).cuda()
    final_array_source = torch.cat([array_light, array_source.unsqueeze(0).unsqueeze(-1)], dim=1)
    final_array_target = torch.cat([array_light, array_source.unsqueeze(0).unsqueeze(-1)], dim=1)
    fws = prior(q_array, final_array_source, zero_padding)
    attr_order = ['Gender', 'Glasses', 'Yaw', 'Pitch', 'Baldness', 'Beard', 'Age', 'Expression']
    attr_current_list = [attr_current[i][0] for i in range(len(attr_order))]
    attr_change = real_value - attr_current_list[attr_index]
    attr_final = attr_degree_list[attr_index] * attr_change + attr_current_list[attr_index]
    final_array_target[0, attr_index + 9, 0, 0] = attr_final
    rev = prior(fws[0], final_array_target, zero_padding, True)
    return rev[0].detach().cpu().numpy()

## generate different pose and expression face image
print("generating...")
for i in tqdm(range(10000)):
    c_order = random.sample([2,3,7], 3)
    new_style = w_current.copy()

    for key in c_order:
        if key == 2 or key == 3:
            value = (random.random()-0.5)*40.0
        else:
            value = (random.random()-0.5)*2.0

        new_style = change(new_style, key, value)

    output_image = model.generate_im_from_w_space(new_style)[0]
    PIL.Image.fromarray(output_image, 'RGB').save(f'outputs/{IDX}/{i:05d}.png')
FeiiYin commented 2 years ago

Many Thanks!

LonglongaaaGo commented 1 year ago

@univ-esuty Thank you so much for your codes. I tried this but found that is can only turn the face from left to right. I would like it also to support turning right to left. Can you give some insights>?

Thank you