Open antongonz opened 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_face
s 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
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?
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?
Because it is weird that the output is blank white. At least, it should be a noise image.
Getting closer, I think. It looks like this after the swap:
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!
Well, a result like this may happen sometimes because of the bad latent codes predicted by HieRFE. You may fix it in two ways:
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.
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:
After the swap however, the face is not swapped:
Did I do something incorrect with my input src?