export_annotation_animal.py

import argparse
import os

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--output_dir', type=str, default='../results/animal')
    parser.add_argument('--partnet_root', type=str, metavar='PATH', default='./data/partnet/')
    parser.add_argument('--gso_root', type=str, metavar='PATH', default='./data/GSO/')
    parser.add_argument('--background_hdr_path', type=str, default='./data/hdri/')
    parser.add_argument('--animal_root', type=str, default='./data/deformingthings4d')
    parser.add_argument('--camera_root', type=str, metavar='PATH',
                        default='./data/camera_trajectory/MannequinChallenge')
    parser.add_argument('--num_assets', type=int, default=5)
    parser.add_argument('--use_gpu', default=False, action='store_true')
    parser.add_argument('--render_engine', type=str, default='CYCLES')
    parser.add_argument('--force_num', type=int, default=5)
    parser.add_argument('--add_force', default=False, action='store_true')
    parser.add_argument('--force_step', type=int, default=3)
    parser.add_argument('--force_interval', type=int, default=120)
    parser.add_argument('--material_path', type=str, default='./data/blender_assets/animal_material.blend')
    parser.add_argument('--add_smoke', default=False, action='store_true')
    parser.add_argument('--animal_name', type=str, metavar='PATH', default=None)

    # rendering settings
    parser.add_argument('--rendering',  default=False, action='store_true')


    # exr settings
    parser.add_argument('--exr',  default=False, action='store_true')
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--frame_idx', type=int, default=1)

    # export obj settings
    parser.add_argument('--export_obj',  default=False, action='store_true')

    # export tracking settings
    parser.add_argument('--export_tracking',  default=False, action='store_true')
    parser.add_argument('--sampling_points', type=int, default=5000)
    parser.add_argument('--sampling_scene_points', type=int, default=1000)
    args = parser.parse_args()


    current_path = os.path.dirname(os.path.realpath(__file__))

    if args.rendering:
        rendering_script = 'blender --background --python {}/render_animal.py \
        -- --output_dir {} --partnet_root {} --gso_root {} --background_hdr_path {} --animal_root {} --camera_root {} ' \
                           '--num_assets {} --render_engine {} --force_num {} ' \
                           '--force_step {} --force_interval {} --material_path {}'.format(
            current_path, args.output_dir, args.partnet_root, args.gso_root,
            args.background_hdr_path, args.animal_root, args.camera_root, args.num_assets,
            args.render_engine, args.force_num, args.force_step, args.force_interval, args.material_path)
        if args.use_gpu:
            rendering_script += ' --use_gpu'
        if args.add_force:
            rendering_script += ' --add_force'
        if args.add_smoke:
            rendering_script += ' --add_smoke'
        if args.animal_name is not None:
            rendering_script += ' --animal_name {}'.format(args.animal_name)
        os.system(rendering_script)
    if args.export_obj:
        obj_script = 'blender --background --python {}/utils/export_obj.py \
         -- --scene_root {} --output_dir {}'.format(
            current_path, os.path.join(args.output_dir, 'scene.blend'), args.output_dir)
        os.system(obj_script)
    if args.exr:
        exr_script = 'python -m utils.openexr_utils --data_dir {} --output_dir {} --batch_size {} --frame_idx {}'.format(
            args.output_dir, args.output_dir + '/exr_img', args.batch_size, args.frame_idx)
        os.system(exr_script)
    if args.export_tracking:
        tracking_script = 'python -m utils.gen_tracking --data_root {} --cp_root {} --sampling_points {} --sampling_scene_points {}'.format(
            args.output_dir, args.output_dir, args.sampling_points, args.sampling_scene_points)
        os.system(tracking_script)

    # rm redundant files
    current_path = os.path.dirname(os.path.realpath(__file__))
    output_dir_abs = os.path.abspath(args.output_dir)
    os.system('cd {} && rm -rf ./*.blend1 && rm -r exr && rm -r tmp'.format(output_dir_abs))
    os.system('cd {}'.format(current_path))