Skip to content

Commit

Permalink
change 2
Browse files Browse the repository at this point in the history
  • Loading branch information
@os-netizen
os-netizen committed Apr 20, 2023
1 parent 502a742 commit 5b30d4b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 17 deletions.
1 change: 1 addition & 0 deletions generate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 98,7 @@ def generate_images(
img = G.synthesis(w.unsqueeze(0), noise_mode=noise_mode)
img = (img.permute(0, 2, 3, 1) * 127.5 128).clamp(0, 255).to(torch.uint8)
imgArr=img[0].cpu().numpy()
imgArr=imgArr.reshape((1, 3, 256, 256))
img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/proj{idx:02d}.png')
return imgArr

Expand Down
38 changes: 21 additions & 17 deletions projector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 37,7 @@ def project(
verbose = False,
device: torch.device
):
print(f"target shape is - {target.shape}")
assert target.shape == (G.img_channels, G.img_resolution, G.img_resolution)

def logprint(*args):
Expand Down Expand Up @@ -132,13 133,13 @@ def logprint(*args):

#----------------------------------------------------------------------------

@click.command()
@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
@click.option('--target', 'target_fname', help='Target image file to project to', required=True, metavar='FILE')
@click.option('--num-steps', help='Number of optimization steps', type=int, default=1000, show_default=True)
@click.option('--seed', help='Random seed', type=int, default=303, show_default=True)
@click.option('--save-video', help='Save an mp4 video of optimization progress', type=bool, default=True, show_default=True)
@click.option('--outdir', help='Where to save the output images', required=True, metavar='DIR')
# @click.command()
# @click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
# @click.option('--target', 'target_fname', help='Target image file to project to', required=True, metavar='FILE')
# @click.option('--num-steps', help='Number of optimization steps', type=int, default=1000, show_default=True)
# @click.option('--seed', help='Random seed', type=int, default=303, show_default=True)
# @click.option('--save-video', help='Save an mp4 video of optimization progress', type=bool, default=True, show_default=True)
# @click.option('--outdir', help='Where to save the output images', required=True, metavar='DIR')
def run_projection(
network_pkl: str,
target_fname: str,
Expand All @@ -165,16 166,19 @@ def run_projection(
G = legacy.load_network_pkl(fp)['G_ema'].requires_grad_(False).to(device) # type: ignore

# Load target image.
target_pil = PIL.Image.open(target_fname).convert('RGB')
# print(target_fname.shape)
target_fname=target_fname.reshape(256, 256, 3)
target_pil = PIL.Image.fromarray(target_fname, "RGB")
w, h = target_pil.size
s = min(w, h)
# target_pil = target_pil.crop(((w - s) // 2, (h - s) // 2, (w s) // 2, (h s) // 2))
# target_pil = target_pil.resize((G.img_resolution, G.img_resolution), PIL.Image.LANCZOS)
# target_uint8 = np.array(target_pil, dtype=np.uint8)

targ=target_fname.crop(((w - s) // 2, (h - s) // 2, (w s) // 2, (h s) // 2))
# targ=target_fname.crop(((w - s) // 2, (h - s) // 2, (w s) // 2, (h s) // 2))
# targ=targ.resize(targ, (G.img_resolution, G.img_resolution))
target_uint8 = np.array(targ, dtype=np.uint8)
target_uint8 = np.array(target_pil, dtype=np.uint8)
print(f"targ_uint8 shape is - {target_uint8.shape}")
# Optimize projection.
start_time = perf_counter()
projected_w_steps = project(
Expand All @@ -187,7 191,7 @@ def run_projection(
print (f'Elapsed: {(perf_counter()-start_time):.1f} s')

# Render debug output: optional video and projected image and W vector.
os.makedirs(outdir, exist_ok=True)
# os.makedirs(outdir, exist_ok=True)
# if save_video:
# video = imageio.get_writer(f'{outdir}/proj.mp4', mode='I', fps=10, codec='libx264', bitrate='16M')
# print (f'Saving optimization progress video "{outdir}/proj.mp4"')
Expand All @@ -199,13 203,13 @@ def run_projection(
# video.close()

# Save final projected frame and W vector.
target_pil.save(f'{outdir}/target.png')
# target_pil.save(f'{outdir}/target.png')
projected_w = projected_w_steps[-1]
synth_image = G.synthesis(projected_w.unsqueeze(0), noise_mode='const')
synth_image = (synth_image 1) * (255/2)
synth_image = synth_image.permute(0, 2, 3, 1).clamp(0, 255).to(torch.uint8)[0].cpu().numpy()
PIL.Image.fromarray(synth_image, 'RGB').save(f'{outdir}/proj.png')
np.savez(f'{outdir}/projected_w.npz', w=projected_w.unsqueeze(0).cpu().numpy())
# synth_image = G.synthesis(projected_w.unsqueeze(0), noise_mode='const')
# synth_image = (synth_image 1) * (255/2)
# synth_image = synth_image.permute(0, 2, 3, 1).clamp(0, 255).to(torch.uint8)[0].cpu().numpy()
# PIL.Image.fromarray(synth_image, 'RGB').save(f'{outdir}/proj.png')
# np.savez(f'{outdir}/projected_w.npz', w=projected_w.unsqueeze(0).cpu().numpy())

return projected_w
#----------------------------------------------------------------------------
Expand Down

0 comments on commit 5b30d4b

Please sign in to comment.