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;
}