From 2d1d7a9e5b3ca066f96a2ec9984344bd7ae3eb8b Mon Sep 17 00:00:00 2001 From: thf5000y2k Date: Sat, 26 Nov 2022 18:47:57 +0800 Subject: [PATCH] Fix bugs for the grayscale image saving --- gen_images.py | 8 +++++--- projector.py | 10 ++++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/gen_images.py b/gen_images.py index 66fdef44..1c0b2238 100644 --- a/gen_images.py +++ b/gen_images.py @@ -133,9 +133,11 @@ def generate_images( G.synthesis.input.transform.copy_(torch.from_numpy(m)) img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode) - img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8) - PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png') - + img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)[0].cpu().numpy() + if img.shape[2]== 1: + PIL.Image.fromarray(np.squeeze(img, axis=(2,)), 'L').save(f'{outdir}/seed{seed:04d}.png') + else: + PIL.Image.fromarray(img, 'RGB').save(f'{outdir}/seed{seed:04d}.png') #---------------------------------------------------------------------------- diff --git a/projector.py b/projector.py index 043d64ac..32797025 100755 --- a/projector.py +++ b/projector.py @@ -520,14 +520,20 @@ def run_projection( for step in tqdm(range(num_steps), desc='Saving projection results', unit='steps'): w = projected_w_steps[step] synth_image = gen_utils.w_to_img(G, dlatents=w, noise_mode='const')[0] - PIL.Image.fromarray(synth_image, 'RGB').save(f'{result_name}_step{step:0{n_digits}d}.jpg') + if synth_image.shape[2] == 1: # [H, W, C == 1] + PIL.Image.fromarray(np.squeeze(synth_image, axis=(2,)), 'L').save(f'{result_name}_final.png') + else: + PIL.Image.fromarray(synth_image, 'RGB').save(f'{result_name}_step{step:0{n_digits}d}.jpg') np.save(f'{npy_name}_step{step:0{n_digits}d}.npy', w.unsqueeze(0).cpu().numpy()) else: # Save only the final projected frame and W vector. print('Saving projection results...') projected_w = projected_w_steps[-1] synth_image = gen_utils.w_to_img(G, dlatents=projected_w, noise_mode='const')[0] - PIL.Image.fromarray(synth_image, 'RGB').save(f'{result_name}_final.jpg') + if synth_image.shape[2] == 1: # [H, W, C == 1] + PIL.Image.fromarray(np.squeeze(synth_image, axis=(2,)), 'L').save(f'{result_name}_final.png') + else: + PIL.Image.fromarray(synth_image, 'RGB').save(f'{result_name}_step{step:0{n_digits}d}.jpg') np.save(f'{npy_name}_final.npy', projected_w.unsqueeze(0).cpu().numpy()) # Save the optimization video and compress it if so desired