source: trunk/src/formats/assimp_loader.cc @ 292

// Copyright (C) 2014 ChaosForge Ltd
// http://chaosforge.org/
//
// This file is part of NV Libraries.
// For conditions of distribution and use, see copyright notice in nv.hh

#include "nv/formats/assimp_loader.hh"
#include <unordered_map>
#include <glm/gtx/transform.hpp>
#include "nv/io/std_stream.hh"
#include "nv/lib/assimp.hh"

using namespace nv;

const int MAX_BONES = 64;

struct assimp_key_p  { float time; vec3 position; };
struct assimp_key_r  { float time; quat rotation; };
struct assimp_key_s  { float time; vec3 scale; };
struct assimp_key_tr { transform tform; };


nv::assimp_loader::assimp_loader( const string& a_ext, const mat3& a_rotate_transform, float a_scale, bool pre_transform, uint32 a_assimp_flags /*= 0 */ ) 
        : m_scene( nullptr ), m_flat( pre_transform ), m_mesh_count(0)
{
        initialize( a_ext, a_rotate_transform, a_scale, a_assimp_flags );
}

nv::assimp_loader::assimp_loader( const string& a_ext, bool pre_transform, uint32 a_assimp_flags /*= 0 */ ) 
        : m_scene( nullptr ), m_flat( pre_transform ), m_mesh_count(0)
{
        initialize( a_ext, mat3(), 1.0f, a_assimp_flags );
}


void nv::assimp_loader::initialize( const string& a_ext, const mat3& a_rotate_transform, float a_scale, uint32 a_assimp_flags )
{
        m_ext   = a_ext;
        m_r33   = a_rotate_transform;
        m_ri33  = glm::transpose( m_r33 );
        m_scale = a_scale;
        m_assimp_flags = a_assimp_flags;
        if ( m_assimp_flags == 0 )
        {
                m_assimp_flags = ( 
                        aiProcess_CalcTangentSpace                              |  
                        aiProcess_GenSmoothNormals                              |  
                        aiProcess_JoinIdenticalVertices                 |   
                        aiProcess_ImproveCacheLocality                  |  
                        aiProcess_LimitBoneWeights                              |  
                        aiProcess_RemoveRedundantMaterials      |  
                        aiProcess_SplitLargeMeshes                              |  
                        aiProcess_Triangulate                                   |  
                        aiProcess_GenUVCoords                   |  
                        aiProcess_SortByPType                   |  
                        aiProcess_FindDegenerates               |  
                        aiProcess_FindInvalidData               |  
                        0 );
        }
}


bool nv::assimp_loader::load( stream& source )
{
        load_assimp_library();
        if ( m_scene != nullptr ) aiReleaseImport( (const aiScene*)m_scene );
        m_scene = nullptr;
        m_mesh_count = 0;
        NV_LOG( nv::LOG_NOTICE, "AssImp loading file..." );
        int size = (int)source.size();
        char* data  = new char[ size ];
        source.read( data, size, 1 );
        const aiScene* scene = aiImportFileFromMemory( data, size, m_assimp_flags, m_ext.c_str() );

        if( !scene)
        {
                NV_LOG( nv::LOG_ERROR, aiGetErrorString() );
                return false;
        }
        m_scene      = scene;
        m_mesh_count = scene->mNumMeshes;
        NV_LOG( nv::LOG_NOTICE, "Loading successfull" );
        return true;
}

mesh_data* nv::assimp_loader::release_mesh_data( size_t index /*= 0 */ )
{
        if ( index >= m_mesh_count ) return nullptr;
        mesh_data* result = new mesh_data;
        load_mesh_data( result, index );
        return result;
}
void nv::assimp_loader::load_mesh_data( mesh_data* data, size_t index )
{
        const aiScene* scene = (const aiScene*)m_scene;
        const aiMesh*  mesh  = scene->mMeshes[ index ];
        data->set_name( mesh->mName.data );

        vec3 vertex_offset     = glm::vec3(); 
        mat3 vertex_transform  = m_scale * m_r33;
        mat3 normal_transform  = m_r33;

        bool skinned = mesh->mNumBones > 0;
        mesh_raw_channel* channel = nullptr;
        if ( skinned )
                channel = mesh_raw_channel::create< assimp_skinned_vtx >( mesh->mNumVertices );
        else
                channel = mesh_raw_channel::create< assimp_plain_vtx >( mesh->mNumVertices );

        data->add_channel( channel );
        for (unsigned int i=0; i<mesh->mNumVertices; i++)
        {
                vec3 v = vertex_transform * assimp_vec3_cast( mesh->mVertices[ i ] ) + vertex_offset;
                vec3 n = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mNormals[ i ] ) );
                vec3 t = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mTangents[ i ] ) );
                vec3 b = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mBitangents[ i ] ) );
                vec2 s = assimp_st_cast( mesh->mTextureCoords[ 0 ][ i ] );

                glm::vec3 t_i = glm::normalize (t - n * glm::dot (n, t));
                float det = (glm::dot (glm::cross (n, t), b));
                det = (det < 0.0f ? -1.0f : 1.0f );
                nv::vec4 vt ( t_i[0], t_i[1], t_i[2], det );
                if ( skinned )
                        ((assimp_skinned_vtx*)channel->data)[i] = assimp_skinned_vtx( v, s, n, vt );
                else
                        ((assimp_plain_vtx*)channel->data)[i] = assimp_plain_vtx( v, s, n, vt );
        }

        if ( skinned )
        {
                assimp_skinned_vtx* vtx = (assimp_skinned_vtx*)channel->data;
                for (unsigned int m=0; m<mesh->mNumBones; m++)
                {
                        aiBone* bone  = mesh->mBones[m];
                        for (unsigned int w=0; w<bone->mNumWeights; w++)
                        {
                                assimp_skinned_vtx& v = vtx[ bone->mWeights[w].mVertexId ];
                                bool found = false;
                                for (nv::uint32 i = 0 ; i < 4; ++i)
                                {
                                        if ( v.boneweight[i] <= 0.0f ) 
                                        {
                                                v.boneindex[i] = m;
                                                v.boneweight[i] = bone->mWeights[w].mWeight;
                                                found = true;
                                                break;
                                        }
                                }
                                NV_ASSERT( found, "Too many weights!" );
                        }
                }
        }

        mesh_raw_channel* ichannel = mesh_raw_channel::create_index( USHORT, mesh->mNumFaces * 3 );
        data->add_channel( ichannel );
        uint16* indices = (uint16*)ichannel->data;
        for (unsigned int i=0; i<mesh->mNumFaces; i++)
        {
                const aiFace* face = &mesh->mFaces[i];
                for (unsigned int j=0; j<face->mNumIndices; j++)
                {
                        indices[ i*3 + j ] = (uint16)face->mIndices[j];
                }
        }
}

nv::assimp_loader::~assimp_loader()
{
        if ( m_scene != nullptr ) aiReleaseImport( (const aiScene*)m_scene );
}

bool nv::assimp_loader::load_bones( size_t index, std::vector< mesh_node_data >& bones )
{
        if ( m_scene == nullptr ) return false;
        const aiScene* scene = (const aiScene*)m_scene;
        const aiMesh*  mesh  = scene->mMeshes[ index ];

        mat4 bone_transform     = mat4( ( 1.f/m_scale * m_ri33 ) ); 
        mat4 bone_pre_transform = mat4( m_scale * m_r33 );

        for (unsigned int m=0; m<mesh->mNumBones; m++)
        {
                aiBone* bone   = mesh->mBones[m];
                mat4    offset = bone_pre_transform * assimp_mat4_cast( bone->mOffsetMatrix ) * bone_transform;
                bones[m].name = bone->mName.data;
                bones[m].data = nullptr;
                bones[m].parent_id = -1;
                bones[m].target_id = -1;
                bones[m].transform = offset;
        }
        return true;
}

void nv::assimp_loader::scene_report() const
{
        const aiScene* scene = (const aiScene*)m_scene;
        if ( scene == nullptr ) return;

        NV_LOG( nv::LOG_NOTICE, "------------------------" );
        NV_LOG( nv::LOG_NOTICE, "Texture   count - " << scene->mNumTextures );
        NV_LOG( nv::LOG_NOTICE, "Animation count - " << scene->mNumAnimations );
        NV_LOG( nv::LOG_NOTICE, "Material  count - " << scene->mNumMaterials );
        NV_LOG( nv::LOG_NOTICE, "Meshes    count - " << scene->mNumMeshes );
        NV_LOG( nv::LOG_NOTICE, "------------------------" );

        aiNode* root = scene->mRootNode;
        if (root)
        {
                NV_LOG( nv::LOG_NOTICE, "Root node  - " << root->mName.data );
                NV_LOG( nv::LOG_NOTICE, "  meshes   - " << root->mNumMeshes );
                NV_LOG( nv::LOG_NOTICE, "  children - " << root->mNumChildren );
        }
        else
        {
                NV_LOG( nv::LOG_NOTICE, "No root node!" );
        }
        NV_LOG( nv::LOG_NOTICE, "------------------------" );

        if ( scene->mNumMeshes > 0 )
        {
                for ( nv::uint32 mc = 0; mc < scene->mNumMeshes; mc++ )
                {
                        aiMesh* mesh = scene->mMeshes[mc];

                        NV_LOG( nv::LOG_NOTICE, "Mesh #"<<mc<<"   - " << std::string( mesh->mName.data ) );
                        NV_LOG( nv::LOG_NOTICE, "  bones   - " << mesh->mNumBones );
                        NV_LOG( nv::LOG_NOTICE, "  uvs     - " << mesh->mNumUVComponents[0] );
                        NV_LOG( nv::LOG_NOTICE, "  verts   - " << mesh->mNumVertices );
                        NV_LOG( nv::LOG_NOTICE, "  faces   - " << mesh->mNumFaces );

                        //                      NV_LOG( nv::LOG_NOTICE, "Bones:" );
                        //                      for (unsigned int m=0; m<mesh->mNumBones; m++)
                        //                      {
                        //                              aiBone* bone  = mesh->mBones[m];
                        //                              NV_LOG( nv::LOG_DEBUG, bone->mName.C_Str() );
                        //                      }
                }
        }
        else
        {
                NV_LOG( nv::LOG_NOTICE, "No meshes!" );
        }
        NV_LOG( nv::LOG_NOTICE, "------------------------" );


        //      if ( scene->mNumMaterials > 0 )
        //      {
        //              for (unsigned int m=0; m < scene->mNumMaterials; m++)
        //              {
        //                      int texIndex = 0;
        //                      aiReturn texFound = aiReturn_SUCCESS;
        //                      aiString path;  // filename
        //                      while (texFound == aiReturn_SUCCESS)
        //                      {
        //                              texFound = scene->mMaterials[m]->GetTexture(aiTextureType_DIFFUSE, texIndex, &path);
        //                              NV_LOG( nv::LOG_NOTICE, "  material - " << path.data );
        //                              texIndex++;
        //                      }
        //              }
        //      }
        //      else
        //      {
        //              NV_LOG( nv::LOG_NOTICE, "No materials" );
        //      }
        //      NV_LOG( nv::LOG_NOTICE, "------------------------" );

}

mesh_nodes_data* nv::assimp_loader::release_merged_bones( mesh_data* meshes )
{
        const aiScene* scene = (const aiScene*)m_scene;
        std::vector< mesh_node_data > final_bones;
        std::unordered_map< std::string, uint16 > names;
        for ( unsigned int m = 0; m < m_mesh_count; ++m )
        {
                sint16 translate[MAX_BONES];
                std::vector< mesh_node_data > bones;
                const aiMesh*  mesh  = scene->mMeshes[ m ];
                if ( mesh->mNumBones != 0 )
                {
                        bones.resize( mesh->mNumBones );
                        load_bones( m, bones );
                        for ( unsigned int b = 0; b < mesh->mNumBones; ++b )
                        {

                                mesh_node_data& bone = bones[b];
                                auto iname = names.find( bone.name );
                                if ( iname == names.end() )
                                {
                                        NV_ASSERT( final_bones.size() < MAX_BONES, "Too many bones to merge!" );
                                        sint16 index = (sint16)final_bones.size();
                                        final_bones.push_back( bone );
                                        names[ bone.name ] = index;
                                        translate[b] = index;
                                }
                                else
                                {
                                        translate[b] = (sint16)iname->second;
                                }
                        }
                        if ( m > 0 && bones.size() > 0 )
                        {
                                mesh_raw_channel* channel = meshes[m].get_raw_channels()[0];
                                NV_ASSERT( !channel->is_index(), "index channel in release_merged!" );
                                assimp_skinned_vtx* va = (assimp_skinned_vtx*)channel->data;
                                for ( unsigned v = 0; v < channel->count; ++v )
                                {
                                        assimp_skinned_vtx& vertex = va[v];

                                        for (uint32 i = 0 ; i < 4; ++i)
                                        {
                                                if ( vertex.boneweight[i] > 0.0f ) 
                                                {
                                                        vertex.boneindex[i] = translate[vertex.boneindex[i]];
                                                }
                                        }
                                }
                        }
                }       
        }
        mesh_node_data* bones = new mesh_node_data[ final_bones.size() ];
        std::copy( final_bones.begin(), final_bones.end(), bones );
        return new mesh_nodes_data( "bones", final_bones.size(), bones );
}

mesh_nodes_data* nv::assimp_loader::release_animation( size_t index )
{
        if ( m_scene == nullptr ) return nullptr;
        const aiScene* scene = (const aiScene*)m_scene;
        if ( scene->mRootNode == nullptr || scene->mAnimations == nullptr || scene->mAnimations[index] == nullptr) return nullptr;

        const aiNode*      root = scene->mRootNode;
        const aiAnimation* anim = scene->mAnimations[index];

        uint32 count = count_nodes( scene->mRootNode );
        mesh_node_data* data    = new mesh_node_data[count];

        uint16 frame_rate     = (uint16)anim->mTicksPerSecond;
        float  duration       = (float)anim->mDuration;
        bool   flat           = m_flat;
        uint32 max_frames     = 0;

        load_node( index, data, root, 0, -1, max_frames );

        // DAE pre_transform hack
        if ( m_flat && frame_rate == 1 )
        {
                frame_rate = 32;
                duration   = (float)max_frames;
        }

        return new mesh_nodes_data( anim->mName.data, count, data, frame_rate, duration, flat ) ;
}

nv::uint32 nv::assimp_loader::count_nodes( const void* node ) const
{
        const aiNode* ainode = (const aiNode*)node;
        nv::uint32 count = 1;
        for ( unsigned i = 0; i < ainode->mNumChildren; ++i )
        {
                count += count_nodes( ainode->mChildren[i] );
        }
        return count;
}

nv::sint16 nv::assimp_loader::load_node( uint32 anim_id, mesh_node_data* nodes, const void* vnode, sint16 this_id, sint16 parent_id, uint32& max_frames )
{
        const aiScene* scene = (const aiScene*)m_scene;
        const aiNode*  node  = (const aiNode*)vnode;
        string name( node->mName.data );
        const aiAnimation* anim  = scene->mAnimations[anim_id];
        const aiNodeAnim*  anode = nullptr;

        for ( unsigned i = 0 ; i < anim->mNumChannels ; i++ )
        {
                anode = anim->mChannels[i];
                if ( std::string( anode->mNodeName.data ) == name ) break;
                anode = nullptr;
        }

        mesh_node_data& a_data = nodes[ this_id ];

        a_data.name      = name;
        a_data.target_id = -1;
        a_data.parent_id = parent_id;
        // This value is ignored by the create_transformed_keys, but needed by create_direct_keys!
        // TODO: find a common solution!
        //       This is bad because create_transformed_keys never uses node-transformations for
        //       node's without keys
        a_data.transform = nv::assimp_mat4_cast( node->mTransformation );
        if (this_id == 0)
                a_data.transform = mat4();
        a_data.data = nullptr;

        if (anode) 
        {
                if ( m_flat )
                {
                        create_transformed_keys( &a_data, anode, parent_id >= 0 ? &(nodes[ parent_id ]) : nullptr );
                }
                else
                {
                        create_direct_keys( &a_data, anode );
                }
                max_frames = glm::max<uint32>( a_data.data->get_channel(0)->count, max_frames );
        }

        nv::sint16 next = this_id + 1;
        for ( unsigned i = 0; i < node->mNumChildren; ++i )
        {
                next = load_node( anim_id, nodes, node->mChildren[i], next, this_id, max_frames );
        }

        return next;
}

void nv::assimp_loader::create_transformed_keys( mesh_node_data* data, const void* vnode, const mesh_node_data* parent )
{
        const aiNodeAnim* node = (const aiNodeAnim*)vnode;
        size_t max_keys = glm::max( node->mNumPositionKeys, node->mNumRotationKeys );

        key_raw_channel* raw_channel = key_raw_channel::create<assimp_key_tr>( max_keys );
        data->data = new key_data;
        data->data->add_channel( raw_channel );
        assimp_key_tr* channel = ((assimp_key_tr*)(raw_channel->data));

        for ( unsigned n = 0; n < max_keys; ++n )
        {
                size_t pn = glm::min( node->mNumPositionKeys - 1, n );
                size_t rn = glm::min( node->mNumRotationKeys - 1, n );
                nv::vec3 pos = m_r33 * nv::assimp_vec3_cast(node->mPositionKeys[pn].mValue) * m_scale;
                nv::quat rot = glm::quat_cast( m_r33 * glm::mat3_cast( assimp_quat_cast(node->mRotationKeys[rn].mValue ) ) * m_ri33 );
                // TODO: only do the calculation when a rotate transform is present!
                nv::transform ptr;
                if ( parent && parent->data )
                {
                        const key_raw_channel* pchannel = parent->data->get_channel(0);
                        if ( pchannel && pchannel->count > 0 )
                        {
                                ptr = ((assimp_key_tr*)pchannel->data)[ glm::min( n, pchannel->count-1 ) ].tform;
                        }
                }

                nv::transform key( ptr * nv::transform( pos, rot ) );
                channel[n].tform = key;
        }
}

void nv::assimp_loader::create_direct_keys( mesh_node_data* data, const void* vnode )
{
        const aiNodeAnim* node = (const aiNodeAnim*)vnode;
        if ( node->mNumPositionKeys == 0 && node->mNumRotationKeys == 0 && node->mNumScalingKeys == 0 )
        {
                return;
        }

        data->data = new key_data;
        key_raw_channel* raw_pchannel = key_raw_channel::create<assimp_key_p>( node->mNumPositionKeys );
        key_raw_channel* raw_rchannel = key_raw_channel::create<assimp_key_r>( node->mNumRotationKeys );
        //key_raw_channel* raw_schannel = key_raw_channel::create<assimp_key_s>( node->mNumScalingKeys );
        data->data->add_channel( raw_pchannel );
        data->data->add_channel( raw_rchannel );
        //data->data->add_channel( raw_schannel );
        assimp_key_p* pchannel = ((assimp_key_p*)(raw_pchannel->data));
        assimp_key_r* rchannel = ((assimp_key_r*)(raw_rchannel->data));
        //assimp_key_s* schannel = ((assimp_key_s*)(raw_schannel->data));

        for ( unsigned np = 0; np < node->mNumPositionKeys; ++np )
        {
                pchannel[np].time     = (float)node->mPositionKeys[np].mTime;
                pchannel[np].position = m_r33 * assimp_vec3_cast(node->mPositionKeys[np].mValue) * m_scale;
        }
        for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
        {
                rchannel[np].time     = (float)node->mRotationKeys[np].mTime;
                rchannel[np].rotation = glm::quat_cast( m_r33 * glm::mat3_cast( assimp_quat_cast(node->mRotationKeys[np].mValue ) ) * m_ri33 );
        }
//      if ( node->mNumScalingKeys > 0 )
//      {
//              nv::vec3 scale_vec0 = assimp_vec3_cast( node->mScalingKeys[0].mValue );
//              float scale_value   = glm::length( glm::abs( scale_vec0 - nv::vec3(1,1,1) ) );
//              if ( node->mNumScalingKeys > 1 || scale_value > 0.001 ) 
//              {
//                      NV_LOG( nv::LOG_WARNING, "scale key significant!" );
//                      for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
//                      {
//                              schannel[np].time  = (float)node->mScalingKeys[np].mTime;
//                              schannel[np].scale = assimp_vec3_cast(node->mScalingKeys[np].mValue);
//                      }
//              }
//              else 
//              {
//                      schannel[0].time  = (float)node->mScalingKeys[0].mTime;
//                      schannel[0].scale = assimp_vec3_cast(node->mScalingKeys[0].mValue);
//              }
//      }

}

mesh_data_pack* nv::assimp_loader::release_mesh_data_pack()
{
        if ( m_scene == nullptr || m_mesh_count == 0 ) return nullptr;
        const aiScene* scene = (const aiScene*)m_scene;
        bool has_bones = false;
        mesh_data* meshes = new mesh_data[ m_mesh_count ];
        for ( size_t m = 0; m < m_mesh_count; ++m )
        {
                const aiMesh* mesh = scene->mMeshes[ m ];
                meshes[m].set_name( mesh->mName.data );
                if ( mesh->mNumBones > 0 ) has_bones = true;
                load_mesh_data(&meshes[m],m);
        }

        mesh_nodes_data* nodes = ( has_bones ? release_merged_bones( meshes ) : release_animation(0) );
        return new mesh_data_pack( m_mesh_count, meshes, nodes );
}

size_t nv::assimp_loader::get_nodes_data_count() const
{
        if ( m_scene == nullptr ) return 0;
        const aiScene* scene = (const aiScene*)m_scene;
        return scene->mNumAnimations;   
}

mesh_nodes_data* nv::assimp_loader::release_mesh_nodes_data( size_t index /*= 0 */ )
{
        return release_animation( index );
}