|
import os |
|
from os.path import join as opj |
|
import argparse |
|
from glob import glob |
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
|
|
|
|
parser.add_argument('--mode', type=str, required=True, choices=['image', 'video', 'manip', 'manip_from_inv'], |
|
help="image: Sample images and shapes, " |
|
"video: Sample pose-controlled videos, " |
|
"manip: Manipulated 3D reconstruction from images, " |
|
"manip_from_inv: Manipulated 3D reconstruction from inverted latent") |
|
parser.add_argument('--network', type=str, nargs='+', required=True) |
|
parser.add_argument('--generator_type', default='ffhq', type=str, choices=['ffhq', 'cat']) |
|
parser.add_argument('--outdir', type=str, default='test_runs') |
|
parser.add_argument('--trunc', type=float, default=0.7) |
|
parser.add_argument('--seeds', type=str, default='100-200') |
|
parser.add_argument('--down_src_eg3d_from_nvidia', default=True) |
|
parser.add_argument('--num_inv_steps', default=300, type=int) |
|
|
|
parser.add_argument('--indir', type=str, default='input_imgs') |
|
parser.add_argument('--name_tag', type=str, default='') |
|
|
|
parser.add_argument('--shape', default=True) |
|
parser.add_argument('--shape_format', type=str, choices=['.mrc', '.ply'], default='.mrc') |
|
parser.add_argument('--shape_only_first', type=bool, default=False) |
|
|
|
parser.add_argument('--grid', default='1x1') |
|
parser.add_argument('--w_frames', type=int, default=120) |
|
|
|
|
|
|
|
args = parser.parse_args() |
|
os.makedirs(args.outdir, exist_ok=True) |
|
print() |
|
|
|
|
|
network_command = '' |
|
for network_path in args.network: |
|
network_command += f"--network {opj('..', network_path)} " |
|
|
|
|
|
|
|
|
|
if args.mode == 'image': |
|
image_path = opj(args.outdir, f'image{args.name_tag}') |
|
os.makedirs(image_path, exist_ok=True) |
|
|
|
os.chdir('eg3d') |
|
command = f"""python gen_samples.py \ |
|
{network_command} \ |
|
--seeds={args.seeds} \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', image_path)} \ |
|
--shapes={args.shape} \ |
|
--shape_format={args.shape_format} \ |
|
--shape_only_first={args.shape_only_first} \ |
|
--trunc={args.trunc} \ |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
|
|
|
|
|
|
|
|
if args.mode == 'video': |
|
video_path = opj(args.outdir, f'video{args.name_tag}') |
|
os.makedirs(video_path, exist_ok=True) |
|
|
|
os.chdir('eg3d') |
|
command = f"""python gen_videos.py \ |
|
{network_command} \ |
|
--seeds={args.seeds} \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', video_path)} \ |
|
--shapes=False \ |
|
--trunc={args.trunc} \ |
|
--grid={args.grid} \ |
|
--w-frames={args.w_frames} |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
|
|
if args.mode == 'manip': |
|
input_path = opj(args.indir) |
|
align_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '1_align_result') |
|
pose_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '2_pose_result') |
|
inversion_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '3_inversion_result') |
|
manip_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '4_manip_result') |
|
|
|
os.makedirs(opj(args.outdir, f'manip_3D_recon{args.name_tag}'), exist_ok=True) |
|
os.makedirs(align_path, exist_ok=True) |
|
os.makedirs(pose_path, exist_ok=True) |
|
os.makedirs(inversion_path, exist_ok=True) |
|
os.makedirs(manip_path, exist_ok=True) |
|
|
|
os.chdir('eg3d') |
|
if args.generator_type == 'cat': |
|
generator_id = 'afhqcats512-128.pkl' |
|
else: |
|
generator_id = 'ffhqrebalanced512-128.pkl' |
|
generator_path = f'pretrained/{generator_id}' |
|
if not os.path.exists(generator_path): |
|
os.makedirs(f'pretrained', exist_ok=True) |
|
print("Pretrained EG3D model cannot be found. Downloading the pretrained EG3D models.") |
|
if args.down_src_eg3d_from_nvidia == True: |
|
os.system(f'wget -c https://api.ngc.nvidia.com/v2/models/nvidia/research/eg3d/versions/1/files/{generator_id} -O {generator_path}') |
|
else: |
|
os.system(f'wget https://huggingface.co/gwang-kim/datid3d-finetuned-eg3d-models/resolve/main/finetuned_models/nvidia_{generator_id} -O {generator_path}') |
|
os.chdir('..') |
|
|
|
|
|
os.chdir('pose_estimation') |
|
if not os.path.exists('checkpoints/pretrained/epoch_20.pth') or not os.path.exists('BFM'): |
|
print(f"BFM and pretrained DeepFaceRecon3D model cannot be found. Downloading the pretrained pose estimation model and BFM files, put epoch_20.pth in ./pose_estimation/checkpoints/pretrained/ and put unzip BFM.zip in ./pose_estimation/.") |
|
|
|
try: |
|
from gdown import download as drive_download |
|
drive_download(f'https://drive.google.com/uc?id=1mdqkEUepHZROeOj99pXogAPJPqzBDN2G', './BFM.zip', quiet=False) |
|
os.system('unzip BFM.zip') |
|
drive_download(f'https://drive.google.com/uc?id=1zawY7jYDJlUGnSAXn1pgIHgIvJpiSmj5', './checkpoints/pretrained/epoch_20.pth', quiet=False) |
|
except: |
|
os.system("pip install -U --no-cache-dir gdown --pre") |
|
from gdown import download as drive_download |
|
drive_download(f'https://drive.google.com/uc?id=1mdqkEUepHZROeOj99pXogAPJPqzBDN2G', './BFM.zip', quiet=False) |
|
os.system('unzip BFM.zip') |
|
drive_download(f'https://drive.google.com/uc?id=1zawY7jYDJlUGnSAXn1pgIHgIvJpiSmj5', './checkpoints/pretrained/epoch_20.pth', quiet=False) |
|
|
|
print() |
|
command = f"""python extract_pose.py 0 \ |
|
{opj('..', input_path)} {opj('..', align_path)} {opj('..', pose_path)} |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
os.chdir('eg3d') |
|
command = f"""python run_inversion.py \ |
|
--outdir={opj('..', inversion_path)} \ |
|
--latent_space_type=w_plus \ |
|
--network={generator_path} \ |
|
--image_path={opj('..', pose_path)} \ |
|
--num_steps={args.num_inv_steps} |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
os.chdir('eg3d') |
|
w_pths = sorted(glob(opj('..', inversion_path, '*.pt'))) |
|
if len(w_pths) == 0: |
|
print("No inverted latent") |
|
exit() |
|
for w_pth in w_pths: |
|
print(f"{w_pth} \n") |
|
|
|
command = f"""python gen_samples.py \ |
|
{network_command} \ |
|
--w_pth={w_pth} \ |
|
--seeds='100-200' \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', manip_path)} \ |
|
--shapes={args.shape} \ |
|
--shape_format={args.shape_format} \ |
|
--shape_only_first={args.shape_only_first} \ |
|
--trunc={args.trunc} \ |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
|
|
command = f"""python gen_videos.py \ |
|
{network_command} \ |
|
--w_pth={w_pth} \ |
|
--seeds='100-200' \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', manip_path)} \ |
|
--shapes=False \ |
|
--trunc={args.trunc} \ |
|
--grid=1x1 \ |
|
--w-frames={args.w_frames} |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
|
|
|
|
|
|
|
|
if args.mode == 'manip_from_inv': |
|
input_path = opj(args.indir) |
|
align_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '1_align_result') |
|
pose_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '2_pose_result') |
|
inversion_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '3_inversion_result') |
|
manip_path = opj(args.outdir, f'manip_3D_recon{args.name_tag}', '4_manip_result') |
|
|
|
os.makedirs(opj(args.outdir, f'manip_3D_recon{args.name_tag}'), exist_ok=True) |
|
os.makedirs(align_path, exist_ok=True) |
|
os.makedirs(pose_path, exist_ok=True) |
|
os.makedirs(inversion_path, exist_ok=True) |
|
os.makedirs(manip_path, exist_ok=True) |
|
|
|
|
|
os.chdir('eg3d') |
|
w_pths = sorted(glob(opj('..', inversion_path, '*.pt'))) |
|
if len(w_pths) == 0: |
|
print("No inverted latent") |
|
exit() |
|
for w_pth in w_pths: |
|
print(f"{w_pth} \n") |
|
|
|
command = f"""python gen_samples.py \ |
|
{network_command} \ |
|
--w_pth={w_pth} \ |
|
--seeds='100-200' \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', manip_path)} \ |
|
--shapes={args.shape} \ |
|
--shape_format={args.shape_format} \ |
|
--shape_only_first={args.shape_only_first} \ |
|
--trunc={args.trunc} \ |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
|
|
command = f"""python gen_videos.py \ |
|
{network_command} \ |
|
--w_pth={w_pth} \ |
|
--seeds='100-200' \ |
|
--generator_type={args.generator_type} \ |
|
--outdir={opj('..', manip_path)} \ |
|
--shapes=False \ |
|
--trunc={args.trunc} \ |
|
--grid=1x1 |
|
""" |
|
print(f"{command} \n") |
|
os.system(command) |
|
os.chdir('..') |
|
|
|
|
|
|
