- Main Page
- Related Pages
- Modules
- Classes
- Examples

# cvColorShader.cpp

    //-
    // ==========================================================================
    // Copyright 1995,2006,2008 Autodesk, Inc. All rights reserved.
    //
    // Use of this software is subject to the terms of the Autodesk
    // license agreement provided at the time of installation or download,
    // or which otherwise accompanies this software in either electronic
    // or hard copy form.
    // ==========================================================================
    //+

    #include <math.h>

    #include <maya/MPxNode.h>
    #include <maya/MIOStream.h>
    #include <maya/MString.h>
    #include <maya/MTypeId.h>
    #include <maya/MPlug.h>
    #include <maya/MDataBlock.h>
    #include <maya/MDataHandle.h>
    #include <maya/MFnNumericAttribute.h>
    #include <maya/MFloatVector.h>
    #include <maya/MPointArray.h>
    #include <maya/MFnPlugin.h>
    #include <maya/MItMeshPolygon.h>
    #include <maya/MFnMesh.h>
    #include <maya/MItMeshVertex.h>
    #include <maya/MDagPath.h>
    #include <maya/MFnSingleIndexedComponent.h>
    #include <maya/MMutexLock.h>
    #include <maya/MRenderUtil.h>
    #include <maya/MSelectionList.h>

    class cvColorShader : public MPxNode
    {
        public:
                        cvColorShader();
        virtual         ~cvColorShader();

        virtual MStatus compute( const MPlug&, MDataBlock& );
        virtual void    postConstructor();

        static  void *  creator();
        static  MStatus initialize();

        //  Id tag for use with binary file format
        static  MTypeId id;

        private:

        static inline float dotProd(const MFloatVector &, const MFloatVector &);

        static MStatus      getTriangleInfo(
                                const MDagPath& meshPath,
                                long            triangleId,
                                MPointArray&    vertPositions,
                                MColorArray&    vertColours
                            );

        static MMutexLock   fCriticalSection;

        // Input attributes

        static MObject aReverseAlpha;

        static MObject aPointObj;               // Implicit attribute
        static MObject aPrimitiveId;            // Implicit attribute
        static MObject aObjectId;           // Implicit attribute

        // Output attributes
        static MObject aOutColor;
        static MObject aOutAlpha;
    };

    // Static data
    MTypeId cvColorShader::id( 0x8000f );
    MMutexLock cvColorShader::fCriticalSection;

    // Attributes
    MObject cvColorShader::aReverseAlpha;
    MObject cvColorShader::aPointObj;
    MObject cvColorShader::aPrimitiveId;
    MObject cvColorShader::aObjectId;
    MObject cvColorShader::aOutColor;
    MObject cvColorShader::aOutAlpha;

    void cvColorShader::postConstructor( )
    {
        setMPSafe(true);
    }

    cvColorShader::cvColorShader()
    {
    }

    cvColorShader::~cvColorShader()
    {
    }

    void * cvColorShader::creator()
    {
        return new cvColorShader();
    }

    MStatus cvColorShader::initialize()
    {
        MFnNumericAttribute nAttr;

        aReverseAlpha = nAttr.create( "reverseAlpha", "ra",
                                      MFnNumericData::kBoolean);
        CHECK_MSTATUS ( nAttr.setDefault( true ) );

        aPointObj  = nAttr.createPoint( "pointObj", "po" );
        CHECK_MSTATUS ( nAttr.setStorable(false) );
        CHECK_MSTATUS ( nAttr.setHidden(true) );

        aPrimitiveId = nAttr.create( "primitiveId", "pi", MFnNumericData::kLong);
        CHECK_MSTATUS ( nAttr.setHidden(true) );

        aObjectId = nAttr.createAddr("objectId", "oi");
        CHECK_MSTATUS ( nAttr.setHidden(true) );

        aOutColor = nAttr.createColor( "outColor", "oc" );
        CHECK_MSTATUS ( nAttr.setStorable(false) );
        CHECK_MSTATUS ( nAttr.setReadable(true) );
        CHECK_MSTATUS ( nAttr.setWritable(false) );

        aOutAlpha = nAttr.create( "outAlpha", "oa", MFnNumericData::kFloat);
        CHECK_MSTATUS (  nAttr.setDisconnectBehavior(MFnAttribute::kReset) );
        CHECK_MSTATUS ( nAttr.setStorable(false) );
        CHECK_MSTATUS ( nAttr.setReadable(true) );
        CHECK_MSTATUS ( nAttr.setWritable(false) );

        CHECK_MSTATUS ( addAttribute(aPointObj) );
        CHECK_MSTATUS ( addAttribute(aOutColor) );
        CHECK_MSTATUS ( addAttribute(aOutAlpha) );
        CHECK_MSTATUS ( addAttribute(aReverseAlpha) );
        CHECK_MSTATUS ( addAttribute(aPrimitiveId) );
        CHECK_MSTATUS ( addAttribute(aObjectId) );

        CHECK_MSTATUS ( attributeAffects(aPointObj,     aOutColor) );
        CHECK_MSTATUS ( attributeAffects(aPrimitiveId,  aOutColor) );
        CHECK_MSTATUS ( attributeAffects(aObjectId,  aOutColor) );

        CHECK_MSTATUS ( attributeAffects(aReverseAlpha, aOutAlpha) );
        CHECK_MSTATUS ( attributeAffects(aPointObj,     aOutAlpha) );
        CHECK_MSTATUS ( attributeAffects(aPrimitiveId,  aOutAlpha) );
        CHECK_MSTATUS ( attributeAffects(aObjectId,  aOutAlpha) );

        return MS::kSuccess;
    }

    // dot product on vectors
    inline float cvColorShader::dotProd(
        const MFloatVector & v1,
        const MFloatVector & v2)
    {
        return  v1.x*v2.x +  v1.y*v2.y + v1.z*v2.z;
    }

    //
    //
    MStatus cvColorShader::compute( const MPlug& plug, MDataBlock& block )
    {
        if ((plug != aOutColor) && (plug.parent() != aOutColor) &&
            (plug != aOutAlpha))
            return MS::kUnknownParameter;

        MStatus status;
        MObject thisNode = thisMObject();

        bool rev_flag = block.inputValue(aReverseAlpha).asBool();
        long triangleId = block.inputValue(aPrimitiveId).asLong();
        void* objectId = block.inputValue(aObjectId).asAddr();

        // Location of the point we are shading
        MFloatVector& pointObj = block.inputValue( aPointObj ).asFloatVector();

        MColor resultColor;

        // It's only worth continuing if the renderer was able to supply us
        // with a surface.
        if (objectId != NULL) {
            // Get the mesh that we are shading.
            MDagPath        meshPath;
            MSelectionList  list;
            status = MRenderUtil::renderObjectItem(objectId, list);
            list.getDagPath(0, meshPath);

            // Get the positions and colours of the triangle's vertices.
            //
            // Note that we could get the triangle's vertices
            MPointArray pos;
            MColorArray colours;
            status = getTriangleInfo(meshPath, triangleId, pos, colours);

            MFloatVector pos1((float)pos[0].x, (float)pos[0].y, (float)pos[0].z);
            MFloatVector pos2((float)pos[1].x, (float)pos[1].y, (float)pos[1].z);
            MFloatVector pos3((float)pos[2].x, (float)pos[2].y, (float)pos[2].z);

            // Compute the barycentric coordinates of the sample.

            pointObj = pointObj - pos3;             // Translate pos3 to origin
            pos1 = pos1 - pos3;
            pos2 = pos2 - pos3;

            MFloatVector norm = pos1 ^ pos2;        // Triangle normal
            float len = dotProd(norm, norm);
            len = dotProd(norm, pointObj)/len;

            pointObj = pointObj - (len * norm);     // Make sure the point is
                                                    // in the triangle

            float aa = dotProd(pos1, pos1);
            float bb = dotProd(pos2, pos2);
            float ab = dotProd(pos1, pos2);
            float am = dotProd(pos1, pointObj);
            float bm = dotProd(pos2, pointObj);
            float det = aa*bb - ab*ab;

            // a, b, c are the barycentric coordinates (assuming pnt
            // is in the triangle plane, best least square fit
            // otherwise.
            //
            float a = (am*bb - bm*ab) / det;
            float b = (bm*aa - am*ab) / det;
            float c = 1-a-b;

            resultColor = (a*colours[0]) + (b*colours[1]) + (c*colours[2]);

            if( rev_flag == true )
                resultColor.a = 1.0f - resultColor.a;
        }

        MDataHandle outColorHandle = block.outputValue( aOutColor );
        MFloatVector& outColor = outColorHandle.asFloatVector();
        outColor.x = resultColor.r;
        outColor.y = resultColor.g;
        outColor.z = resultColor.b;
        outColorHandle.setClean();

        MDataHandle outAlphaHandle = block.outputValue( aOutAlpha );
        float& outAlpha = outAlphaHandle.asFloat();
        outAlpha = resultColor.a;
        outAlphaHandle.setClean();

        return MS::kSuccess;
    }

    MStatus cvColorShader::getTriangleInfo(
        const MDagPath& meshPath,
        long            triangleId,
        MPointArray&    vertPositions,
        MColorArray&    vertColours
    ) {
        MStatus st;

        //  'triangleId' refers to the triangle currently being shaded. We need
        //  to find the positions and colours of the triangle's three vertices.
        //
        //  We could use the 'vertexCamera*' render attributes to determine the
        //  positions of the triangle's vertices, but to determine the color at
        //  a vertex we need to know the vertex's index within the mesh and
        //  there is no render attribute which will give us that. We will have
        //  to determine it ourselves by finding the face to which the triangle
        //  belongs and then using MItMeshPolygon::getTriangle() to get the
        //  indices of the triangle's vertices. The call to getTriangle() also
        //  happens to return the positions of the vertices, so we won't bother
        //  with the 'vertexCamera*' render attributes.
        //
        //  The way we find the face to which the triangle belongs is by
        //  running through all of the mesh's faces and counting the number of
        //  triangles in each one. When the count exceeds the value of
        //  'triangleId', the face which put us over is the face containing the
        //  triangle.
        //
        //  The renderer does not assign triangle ids to the mesh all at once,
        //  but separately for each shading group. For example, let's say that
        //  a mesh has 20 faces comprised of a total of 30 triangles, assigned
        //  to two different shaders, as follows:
        //
        //  - 12 faces, having a total of 19 triangles, are assigned to
        //    the first shader
        //  - 5 faces, having a total of 7 triangles, are assigned to
        //    the second shader
        //  - the remaining 3 faces, having a total of 4 triangles,
        //    are not assigned to any shader.
        //
        //  The 19 triangles in the first shader will be given primitiveIds 0 to 18.
        //  The 7 trianges in the second shader will be given primitiveIds 19 to 25.
        //  The 4 triangles which are not assigned to any shader won't have any
        //  primitiveIds assigned to them.
        //
        //  So when we're counting triangles, we must do it in shader order.

        //  The first step is to get all of the shaders used by this mesh and the
        //  faces to which they are assigned.
        MFnMesh         meshFn(meshPath);
        MObjectArray    shaders;
        MObjectArray    components;

        meshFn.getConnectedSetsAndMembers(
            meshPath.instanceNumber(), shaders, components, true
        );

        int         polygonId = -1;
        int         triangleCount = 0;
        MIntArray   vertIndices;

        //  Step through each shader.
        for (unsigned int s = 0; (polygonId < 0) && (s < shaders.length()); ++s) {
            //  Iterate over the faces assigned to this shader.
            //
            //  The constructor for MItMeshPolygon is not threadsafe as it may
            //  initiate a recalculation of the mesh's normals. So we must lock
            //  the thread while making the call.
            //
            //  Similarly, MItMeshMeshPolygon::hasValidTriangulation() may
            //  trigger triangulation of the mesh, which is also not thread
            //  safe. So we want to keep our lock until after the first call to
            //  it.
            bool    isLocked = true;
            fCriticalSection.lock();
            MItMeshPolygon faceIter(meshPath, components[s]);

            for (; !faceIter.isDone(); faceIter.next()) {
                if (faceIter.hasValidTriangulation()) {
                    //  Get the number of triangles in the current face.
                    int nTri;
                    faceIter.numTriangles(nTri);

                    //  If this face will put the count over 'triangleId' then
                    //  the triangle must belong to this face.
                    if (triangleId < triangleCount + nTri) {
                        //  Get the positions and indices of the triangle's
                        //  vertices, then break out of the loop. We subtract
                        //  'triangleCount' from 'triangleId' to get the index
                        //  of the triangle within the face.
                        polygonId = faceIter.index();
                        st = faceIter.getTriangle(
                                triangleId - triangleCount,
                                vertPositions,
                                vertIndices,
                                MSpace::kObject
                            );

                        break;
                    }

                    //  We haven't found the right face yet. Add the number of
                    //  triangles in this face to the count and keep going.
                    triangleCount += nTri;
                }

                //  If hasValidTriangulation() was going to triangulate the
                //  mesh it will have done so by now. Subsequent calls will
                //  use the existing triangulation so it's safe to remove the
                //  lock now.
                if (isLocked) {
                    fCriticalSection.unlock();
                    isLocked = false;
                }
            }

            //  If the shader has no face components assigned to it then the
            //  'for' loop above will not have run and fCriticalSection will
            //  still be locked, in which case we must unlock it now.
            if (isLocked) {
                fCriticalSection.unlock();
            }
        }

        if ((polygonId == -1) || !st) {
            return MS::kFailure;
        }

        //  Now that we know the indices of the triangle's vertices, get their
        //  colours.
        MItMeshVertex vertIter(meshPath);
        int preIndex = 0;
        vertColours.setLength(3);

        CHECK_MSTATUS ( vertIter.setIndex( vertIndices[0], preIndex) );
        CHECK_MSTATUS ( vertIter.getColor( vertColours[0], polygonId ) );

        CHECK_MSTATUS ( vertIter.setIndex( vertIndices[1], preIndex) );
        CHECK_MSTATUS ( vertIter.getColor( vertColours[1], polygonId ) );

        CHECK_MSTATUS ( vertIter.setIndex( vertIndices[2], preIndex) );
        CHECK_MSTATUS ( vertIter.getColor( vertColours[2], polygonId ) );

        return MS::kSuccess;
    }

    MStatus initializePlugin( MObject obj )
    {
        const MString UserClassify( "utility/color" );

        MFnPlugin plugin( obj, PLUGIN_COMPANY, "5.0", "Any");
        CHECK_MSTATUS ( plugin.registerNode( "cvColorShader", cvColorShader::id,
                             cvColorShader::creator,
                             cvColorShader::initialize,
                             MPxNode::kDependNode, &UserClassify ) );

        return MS::kSuccess;
    }

    MStatus uninitializePlugin( MObject obj )
    {
        MFnPlugin plugin( obj );
        CHECK_MSTATUS ( plugin.deregisterNode( cvColorShader::id ) );

        return MS::kSuccess;
    }