fantasyfish/dreamgaussian

1import os2import glob3import argparse4 5parser = argparse.ArgumentParser()6parser.add_argument('--dir', default='data', type=str, help='Directory where processed images are stored')7parser.add_argument('--out', default='logs', type=str, help='Directory where obj files will be saved')8parser.add_argument('--video-out', default='videos', type=str, help='Directory where videos will be saved')9parser.add_argument('--gpu', default=0, type=int, help='ID of GPU to use')10parser.add_argument('--elevation', default=0, type=int, help='Elevation angle of view in degrees')11parser.add_argument('--config', default='configs', type=str, help='Path to config directory, which contains image.yaml')12args = parser.parse_args()13 14files = glob.glob(f'{args.dir}/*_rgba.png')15configs_dir = args.config16 17# check if image.yaml exists18if not os.path.exists(os.path.join(configs_dir, 'image.yaml')):19    raise FileNotFoundError(20        f'image.yaml not found in {configs_dir} directory. Please check if the directory is correct.'21    )22 23# create output directories if not exists24out_dir = args.out25os.makedirs(out_dir, exist_ok=True)26video_dir = args.video_out27os.makedirs(video_dir, exist_ok=True)28 29 30for file in files:31    name = os.path.basename(file).replace("_rgba.png", "")32    print(f'======== processing {name} ========')33    # first stage34    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main.py '35              f'--config {configs_dir}/image.yaml '36              f'input={file} '37              f'save_path={name} elevation={args.elevation}')38    # second stage39    os.system(f'CUDA_VISIBLE_DEVICES={args.gpu} python main2.py '40              f'--config {configs_dir}/image.yaml '41              f'input={file} '42              f'save_path={name} elevation={args.elevation}')43    # export video44    mesh_path = os.path.join(out_dir, f'{name}.obj')45    os.system(f'python -m kiui.render {mesh_path} '46              f'--save_video {video_dir}/{name}.mp4 '47              f'--wogui '48              f'--elevation {args.elevation}')49