zyainfal / One-Shot-Face-Swapping-on-Megapixels

One Shot Face Swapping on Megapixels.
Other
314 stars 40 forks source link

Swap with custom input src not working #29

Open antongonz opened 2 years ago

antongonz commented 2 years ago

I tried to follow your instructions about face detection and alignment. I ended up with a face image that looks like this after the preprocess resize:

image

After the swap however, the face is not swapped:

image

Did I do something incorrect with my input src?

zyainfal commented 2 years ago

It is correct that the input image size is 256.

You may find the mistake by looking into the run function:

    def run(self, src_idx, tgt_idx, refine=True):
        src_face_rgb, tgt_face_rgb, tgt_mask = self.read_pair(src_idx, tgt_idx)
        source, target = self.preprocess(src_face_rgb, tgt_face_rgb)
        swapped_face = self.swap(source, target)
        swapped_face = self.postprocess(swapped_face, tgt_face_rgb, tgt_mask)

        result = np.hstack((src_face_rgb[:,:,::-1], tgt_face_rgb[:,:,::-1], swapped_face))

In this snippet, you may check the source, target, and two swapped_faces before and after postprocess to see if they are blank. If the swapped_face before postprocess does give some results, the mistake should locate in postprocess .

In addition, if the swapped_face before postprocess has nothing, you may check if the model loaded correctly by reconstructing the input images. The function refine is a sample for reconstruction for you:

    def refine(self, swapped_tensor): # swapped_tensor is of size (1, 3, 256, 256), as same as an input image tensor
        with torch.no_grad():
            lats, struct = self.encoder(swapped_tensor.cuda())

            fake_refine, _ = self.generator(struct, [lats, None], randomize_noise=False)

            fake_refine_max = torch.max(fake_refine)
            fake_refine_min = torch.min(fake_refine)
            denormed_fake_refine = (fake_refine[0] - fake_refine_min) / (fake_refine_max - fake_refine_min) * 255.0
            fake_refine_numpy = denormed_fake_refine.permute((1, 2, 0)).cpu().numpy()
        return fake_refine_numpy
antongonz commented 2 years ago

The images after preprocess looks fine. After the swap I try:

        cv2.imshow('swapped_face ', swapped_face )
        cv2.waitKey(0)

but this does not show an image at all, just blank white. Do I need to do something different to see the image after swap?

zyainfal commented 2 years ago

Have you denormalized the tensor back to [0, 255]? If you did, then I think the problem is at the generation process. Can you generate random faces by StyleGAN2?

zyainfal commented 2 years ago

Because it is weird that the output is blank white. At least, it should be a noise image.

antongonz commented 2 years ago

Getting closer, I think. It looks like this after the swap:

image

If I swap using CelebA pictures, everything works normal. Essentially my code looks like this for custom images:

face = handler.get_face(cv2.imread('test.jpg')) #detect face and align
face = cv2.resize(faces,(1024,1024),interpolation=cv2.INTER_LINEAR) #resize

and then I send this to preprocess. Any idea what might be wrong?

Thanks!

zyainfal commented 2 years ago

Well, a result like this may happen sometimes because of the bad latent codes predicted by HieRFE. You may fix it in two ways:

  1. fine-tune HieRFE on your data, which is slow
  2. use projector.py in stylegan2-pytorch repo to reconstruct your image in stylegan fashion, then used this reconstructed image as the input of our model.

Anyway, the keypoints are: 1. use more data to train a HieRFE with better generalization, or 2. directly find good latent codes for swapping.