TextureBaker.h

Go to the documentation of this file.

//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//
 
#ifndef MATERIALX_TEXTUREBAKER
#define MATERIALX_TEXTUREBAKER

 
#include <iostream>
 
#include <MaterialXCore/Unit.h>
 
#include <MaterialXRender/Export.h>
#include <MaterialXFormat/File.h>
#include <MaterialXRender/ImageHandler.h>
#include <MaterialXGenShader/GenContext.h>
 
MATERIALX_NAMESPACE_BEGIN

using BakedDocumentVec = std::vector<std::pair<std::string, DocumentPtr>>;

template <typename Renderer, typename ShaderGen>
class TextureBaker : public Renderer
{
  public:
    void setExtension(const string& extension)
    {
        _extension = extension;
    }

    const string& getExtension() const
    {
        return _extension;
    }

    void setColorSpace(const string& colorSpace)
    {
        _colorSpace = colorSpace;
    }

    const string& getColorSpace() const
    {
        return _colorSpace;
    }

    void setDistanceUnit(const string& unitSpace)
    {
        _distanceUnit = unitSpace;
    }

    const string& getDistanceUnit() const
    {
        return _distanceUnit;
    }

    void setAverageImages(bool enable)
    {
        _averageImages = enable;
    }

    bool getAverageImages() const
    {
        return _averageImages;
    }

    void setOptimizeConstants(bool enable)
    {
        _optimizeConstants = enable;
    }

    bool getOptimizeConstants() const
    {
        return _optimizeConstants;
    }

    void setOutputImagePath(const FilePath& outputImagePath)
    {
        _outputImagePath = outputImagePath;
    }

    const FilePath& getOutputImagePath()
    {
        return _outputImagePath;
    }

    void setBakedGraphName(const string& name)
    {
        _bakedGraphName = name;
    }

    const string& getBakedGraphName() const
    {
        return _bakedGraphName;
    }

    void setBakedGeomInfoName(const string& name)
    {
        _bakedGeomInfoName = name;
    }

    const string& getBakedGeomInfoName() const
    {
        return _bakedGeomInfoName;
    }

    const string& getTextureFilenameTemplate() const
    {
        return _textureFilenameTemplate;
    }

    void setTextureFilenameTemplate(const string& filenameTemplate)
    {
        _textureFilenameTemplate = (filenameTemplate.find("$EXTENSION") == string::npos) ?
            filenameTemplate + ".$EXTENSION" : filenameTemplate;
    }

    void setFilenameTemplateVarOverride(const string& key, const string& value)
    {
        if (_permittedOverrides.count(key))
        {
            _texTemplateOverrides[key] = value;
        }
    }

    void setOutputStream(std::ostream* outputStream)
    {
        _outputStream = outputStream;
    }

    std::ostream* getOutputStream() const
    {
        return _outputStream;
    }

    void setHashImageNames(bool enable)
    {
        _hashImageNames = enable;
    }

    bool getHashImageNames() const
    {
        return _hashImageNames;
    }

    void setTextureSpaceMin(const Vector2& min)
    {
        _textureSpaceMin = min;
    }

    Vector2 getTextureSpaceMin() const
    {
        return _textureSpaceMin;
    }

    void setTextureSpaceMax(const Vector2& max)
    {
        _textureSpaceMax = max;
    }

    Vector2 getTextureSpaceMax() const
    {
        return _textureSpaceMax;
    }

    void setupUnitSystem(DocumentPtr unitDefinitions);

    void bakeShaderInputs(NodePtr material, NodePtr shader, GenContext& context, const string& udim = EMPTY_STRING);

    void bakeGraphOutput(OutputPtr output, GenContext& context, const StringMap& filenameTemplateMap);

    void optimizeBakedTextures(NodePtr shader);

    DocumentPtr bakeMaterialToDoc(DocumentPtr doc, const FileSearchPath& searchPath, const string& materialPath,
                                  const StringVec& udimSet, std::string& documentName);

    void bakeAllMaterials(DocumentPtr doc, const FileSearchPath& searchPath, const FilePath& outputFileName);

    void writeDocumentPerMaterial(bool value)
    {
        _writeDocumentPerMaterial = value;
    }
 
    string getValueStringFromColor(const Color4& color, const string& type);
 
  protected:
    class BakedImage
    {
      public:
        FilePath filename;
        Color4 uniformColor;
        bool isUniform = false;
    };
    class BakedConstant
    {
      public:
        Color4 color;
        bool isDefault = false;
    };
    using BakedImageVec = vector<BakedImage>;
    using BakedImageMap = std::unordered_map<OutputPtr, BakedImageVec>;
    using BakedConstantMap = std::unordered_map<OutputPtr, BakedConstant>;
 
  protected:
    TextureBaker(unsigned int width, unsigned int height, Image::BaseType baseType, bool flipSavedImage);
 
    // Populate file template variable naming map
    StringMap initializeFileTemplateMap(InputPtr input, NodePtr shader, const string& udim = EMPTY_STRING);
 
    // Find first occurrence of variable in filename from start index onwards
    size_t findVarInTemplate(const string& filename, const string& var, size_t start = 0);
 
    // Generate a texture filename for the given graph output.
    FilePath generateTextureFilename(const StringMap& fileTemplateMap);
 
    // Create document that links shader outputs to a material.
    DocumentPtr generateNewDocumentFromShader(NodePtr shader, const StringVec& udimSet);
 
    // Write a baked image to disk, returning true if the write was successful.
    bool writeBakedImage(const BakedImage& baked, ImagePtr image);
 
  protected:
    string _extension;
    string _colorSpace;
    string _distanceUnit;
    bool _averageImages;
    bool _optimizeConstants;
    FilePath _outputImagePath;
    string _bakedGraphName;
    string _bakedGeomInfoName;
    string _textureFilenameTemplate;
    std::ostream* _outputStream;
    bool _hashImageNames;
    Vector2 _textureSpaceMin;
    Vector2 _textureSpaceMax;
 
    ShaderGeneratorPtr _generator;
    ConstNodePtr _material;
    ImagePtr _frameCaptureImage;
    BakedImageMap _bakedImageMap;
    BakedConstantMap _bakedConstantMap;
    StringSet _permittedOverrides;
    StringMap _texTemplateOverrides;
    StringMap _bakedInputMap;
 
    std::unordered_map<string, NodePtr> _worldSpaceNodes;
 
    bool _flipSavedImage;
 
    bool _writeDocumentPerMaterial;
    DocumentPtr _bakedTextureDoc;
};
 
MATERIALX_NAMESPACE_END
 
#include <MaterialXRender/TextureBaker.inl>
 
#endif