eladrich / pixel2style2pixel

Official Implementation for "Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation" (CVPR 2021) presenting the pixel2style2pixel (pSp) framework
https://eladrich.github.io/pixel2style2pixel/
MIT License
3.19k stars 570 forks source link

how to get each picture's latent code? #36

Closed alip7 closed 3 years ago

alip7 commented 3 years ago

I set the 'return_latents=True' in psp.py , but get TypeError :

【 Traceback (most recent call last): File "/home/zhz1/anaconda3/envs/env_36/lib/python3.7/site-packages/PIL/Image.py", line 2751, in fromarray mode, rawmode = _fromarray_typemap[typekey] KeyError: ((1, 1, 1, 256), '|u1')

The above exception was the direct cause of the following exception:

Traceback (most recent call last): File "scripts/inference.py", line 136, in run() File "scripts/inference.py", line 82, in run result = tensor2im(result_batch[i]) File "./utils/common.py", line 23, in tensor2im return Image.fromarray(var.astype('uint8')) File "/home/zhz1/anaconda3/envs/env_36/lib/python3.7/site-packages/PIL/Image.py", line 2753, in fromarray raise TypeError("Cannot handle this data type: %s, %s" % typekey) from e TypeError: Cannot handle this data type: (1, 1, 1, 256), |u1 】

could you please tell me how to output each picture's latent code?

yuval-alaluf commented 3 years ago

Did you only make a change in psp.py?
You most likely got the error because you did not make all the necessary changes needed.

You should make the changes in inference.py instead. You should change the call https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L111 in inference.py to return a tuple as such:

images, result_latents = net(...,return_latents=True) 

Then, in inference.py in line 129 you should change https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L129 to

return images, result_latents

Finally, in line 73, you should change the line https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L73 to something like

result_batch, result_latents = run_on_batch(input_cuda, net, opts)

Now you have the latents of the batch and can save them if you wish. In short, returning latents is very much possible with some minimal changes to the inference.py file. I may have missed a small step, but the above changes should assist you in getting started. If you still encounter issues feel free to reply.

alip7 commented 3 years ago

Did you only make a change in psp.py? You most likely got the error because you did not make all the necessary changes needed.

You should make the changes in inference.py instead. You should change the call https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L111

in inference.py to return a tuple as such:

images, result_latents = net(...,return_latents=True) 

Then, in inference.py in line 129 you should change https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L129

to

return images, result_latents

Finally, in line 73, you should change the line https://github.com/eladrich/pixel2style2pixel/blob/751fcbd9ffcb86487afdaa6a7e5953b94e28369f/scripts/inference.py#L73

to something like

result_batch, result_latents = run_on_batch(input_cuda, net, opts)

Now you have the latents of the batch and can save them if you wish. In short, returning latents is very much possible with some minimal changes to the inference.py file. I may have missed a small step, but the above changes should assist you in getting started. If you still encounter issues feel free to reply.

Thanks for your instruction and your great work! It helps a lot. now I get the (18,512) of latent code for each picture.

yuval-alaluf commented 3 years ago

Awesome!