yuval-alaluf / stylegan3-editing

Official Implementation of "Third Time's the Charm? Image and Video Editing with StyleGAN3" (AIM ECCVW 2022) https://arxiv.org/abs/2201.13433
https://yuval-alaluf.github.io/stylegan3-editing/
MIT License
654 stars 73 forks source link

Generating image using Labels #35

Closed rut00 closed 2 years ago

rut00 commented 2 years ago

I have the StyleGAN3 network trained on the labeled dataset. I want to know how and where to pass a label to generate images while editing? My labels are in numeric form like 1 and 2.

yuval-alaluf commented 2 years ago

Good question, I actually haven't tried this myself. Typically, when we want to generate an image we perform the following:

        w = generator.mapping(z, None, truncation_psi=truncation_psi)
        img = generator.synthesis(w, noise_mode='const')

It seems like None is value of the label argument. Looking at the official StyleGAN3 code, they generate images using:

img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)

where label is defined here. I would try to follow their format and try to pass the label when generating an image. That is, if you look at the code, calling G(.) is the same as:

        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas)
        img = self.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)

So you probably want to change our code to the following:

        w = generator.mapping(z, c, truncation_psi=truncation_psi)
        w_edit = [edit you latent code w]
        img = generator.synthesis(w, noise_mode='const')

where c is our label as defined in the official SG3 repo. Hope this helps.

rut00 commented 2 years ago

For generating images with label, we added below lines in edit_synthetic.py file under the function get_random_image:

label = torch.zeros([1, G.c_dim], device=device)
# Label 1
label[:, 1] = 1

label = torch.zeros([1, G.c_dim], device=device)
# Label 2
label[:, 2] = 1

Thank you for your help.