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

Add alignment to inference #125

Closed Nerdyvedi closed 3 years ago

Nerdyvedi commented 3 years ago

Hi, I wanted to add alignment step to inference. I would like to use a single script to give the end result.

What I did was added alignment step to inference script, Saving the aligned image , then reading it for inference.

Could you suggest a better method, I want to avoid saving the intermediate aligned image and want to reduce the total inference time as much as possible

yuval-alaluf commented 3 years ago

Hi @Nerdyvedi , I see a couple of options, but it really depends on if you are planning on running inference on the same images many times. If so, it doesn't seem efficient to run the alignment each time you run inference. I would therefore recommend running the alignment script before running inference and saving the images in a new directory. Then, when running inference, pass this new directory as your data path.
You said that you wanted a single script though, and this isn't exactly what you wanted, but this would be the most efficient solution. If you insist on using a single script, I would recommend adding the alignment to the dataloader. For example, in inference_dataset.py you can do:

def __getitem__(self, index):
   from_path = self.paths[index]
   try:
        from_im = align_face(filepath=from_path, predictor=predictor)
   except:
        raise Exception(f"Invalid image! Could not align: {from_path}"
    from_im = from_im.convert('RGB') if self.opts.label_nc == 0 else from_im.convert('L')
    if self.transform:
        from_im = self.transform(from_im)
    return from_im

The align_face() function is taken from the align_all_parallel script and predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat") I didn't test this, but maybe something like this? Notice that you don't need to save the aligned image and then read it before. What do you think?

Nerdyvedi commented 3 years ago

Thanks a lot, it works well. One more question, I want the output image to be the same size as input image(I am using PsP for toonification). What I have done is created a separate function that reads the image and returns the size.

It increases the total inference time, Is there a way I can get the input size from the loader?

Thank you

yuval-alaluf commented 3 years ago

Are all your images of different sizes?
The alignment step saves the images to be 256x256 by default and the generator generates images of size 1024x1024. You can either down-sample the output images or resize the input images to the size of the output. However, I wouldn't recommend down-sampling the output images. Therefore, you can either resize the inputs to 1024x1024 or alter the alignment script to output images of size 1024x1024 by changing the following lines from 256 to 1024. https://github.com/eladrich/pixel2style2pixel/blob/72f0c04a890adb9c8a406f879980572ac600cf29/scripts/align_all_parallel.py#L93-L94 Does this answer your question?

Nerdyvedi commented 3 years ago

Yes, Got it. Thanks