script to generate gltf from pre-trained banmo models

gistfile1.txt

"""
# script to generate gltf from pre-trained banmo models
python scripts/visualize/obj2glb.py --flagfile logdir/hmnerf-cate-pretrain-shiba-haru-1-e120-b72-init-dog80-e120-b512-ft2/opts.log
"""
import os
import sys
import cv2
import trimesh
import glob
import numpy as np
import pdb
from scipy.spatial.transform import Rotation as R
import configparser
from utils.gltf_lib import load_mesh, load_bg, load_rtks, load_bone_rtks, create_gltf
from absl import flags, app
from nnutils.train_utils import v2s_trainer
opts = flags.FLAGS
  
def main(_):
    # get urdf file
    opts.model_path = ''
    trainer = v2s_trainer(opts, is_eval=True)
    data_info = trainer.init_dataset()
    trainer.define_model(data_info)
    urdf = trainer.model.nerf_body_rts.urdf
    
    model_dir=sys.argv[2].rsplit('/', 1)[0]
    seqname=opts.seqname
    vidid_all=range(len(data_info['offset'])-1)
    
    # input output 
    #model_dir=sys.argv[1]
    #seqname=sys.argv[2]
    #vidid=int(sys.argv[3])

    for vidid in vidid_all:
        outpath='%s/mesh-%s.gltf'%(model_dir, vidid)
        
        config = configparser.RawConfigParser()
        config.read('configs/%s.config'%seqname)
        vidname=config['data_%d'%vidid]['datapath'].split('/')[4]
        
        # load meshes
        mesh, weights,rest_bone = load_mesh(model_dir, vidid)
        mesh_bg = load_bg(model_dir, vidname)
        # load transforms
        root_rtks = load_rtks('%s/*%s*-cam-*.txt'%(model_dir, vidname))
        bone_rtks = load_bone_rtks('%s/*%s*se3*.txt'%(model_dir, vidname)) # x = T x* => x = j2x* j*2j x*2j* x*
        # transform camera space motion to world space
        #try:
        #    bg_rtks = load_rtks('%s/*%s*-bgcam-*.txt'%(model_dir, vidname))
        #    root_rtks = np.linalg.inv(bg_rtks)@root_rtks # transform to bg canonical space
        #except: pass

        # set root transforms to identity
        root_rtks[:,:3,:3] = np.eye(3)[None]
        root_rtks[:,:3,3] = 0
        
        gltf = create_gltf(mesh, weights, rest_bone, root_rtks, bone_rtks, urdf=urdf)
        #gltf = create_gltf(mesh, weights, rest_bone, root_rtks, bone_rtks, mesh_bg=mesh_bg)
        gltf.save(outpath)

if __name__ == '__main__':
    app.run(main)