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