Files
UnrealEngine/Engine/Source/ThirdParty/MaterialX/MaterialX-1.38.10/source/MaterialXRender/TextureBaker.inl
2025-05-18 13:04:45 +08:00

652 lines
23 KiB
C++

//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//
#include <MaterialXRender/OiioImageLoader.h>
#include <MaterialXRender/StbImageLoader.h>
#include <MaterialXRender/Util.h>
#include <MaterialXGenShader/DefaultColorManagementSystem.h>
#include <MaterialXFormat/XmlIo.h>
MATERIALX_NAMESPACE_BEGIN
namespace
{
const string SRGB_TEXTURE = "srgb_texture";
const string LIN_REC709 = "lin_rec709";
const string SHADER_PREFIX = "SR_";
const string DEFAULT_UDIM_PREFIX = "_";
} // anonymous namespace
template <typename Renderer, typename ShaderGen>
string TextureBaker<Renderer, ShaderGen>::getValueStringFromColor(const Color4& color, const string& type)
{
if (type == "color4" || type == "vector4")
{
return toValueString(color);
}
if (type == "color3" || type == "vector3")
{
return toValueString(Vector3(color[0], color[1], color[2]));
}
if (type == "vector2")
{
return toValueString(Vector2(color[0], color[1]));
}
if (type == "float")
{
return toValueString(color[0]);
}
return EMPTY_STRING;
}
template <typename Renderer, typename ShaderGen>
TextureBaker<Renderer, ShaderGen>::TextureBaker(unsigned int width, unsigned int height, Image::BaseType baseType, bool flipSavedImage) :
Renderer(width, height, baseType),
_distanceUnit("meter"),
_averageImages(false),
_optimizeConstants(true),
_bakedGraphName("NG_baked"),
_bakedGeomInfoName("GI_baked"),
_textureFilenameTemplate("$MATERIAL_$SHADINGMODEL_$INPUT$UDIMPREFIX$UDIM.$EXTENSION"),
_outputStream(&std::cout),
_hashImageNames(false),
_textureSpaceMin(0.0f),
_textureSpaceMax(1.0f),
_generator(ShaderGen::create()),
_permittedOverrides({ "$ASSET", "$MATERIAL", "$UDIMPREFIX" }),
_flipSavedImage(flipSavedImage),
_writeDocumentPerMaterial(true),
_bakedTextureDoc(nullptr)
{
if (baseType == Image::BaseType::UINT8)
{
#if MATERIALX_BUILD_OIIO
_extension = ImageLoader::TIFF_EXTENSION;
#else
_extension = ImageLoader::PNG_EXTENSION;
#endif
_colorSpace = SRGB_TEXTURE;
}
else
{
#if MATERIALX_BUILD_OIIO
_extension = ImageLoader::EXR_EXTENSION;
#else
_extension = ImageLoader::HDR_EXTENSION;
#endif
_colorSpace = LIN_REC709;
}
// Initialize our base renderer.
Renderer::initialize();
// Initialize our image handler.
Renderer::_imageHandler = Renderer::createImageHandler(StbImageLoader::create());
#if MATERIALX_BUILD_OIIO
Renderer::_imageHandler->addLoader(OiioImageLoader::create());
#endif
// Create our dedicated frame capture image.
_frameCaptureImage = Image::create(width, height, 4, baseType);
_frameCaptureImage->createResourceBuffer();
}
template <typename Renderer, typename ShaderGen>
size_t TextureBaker<Renderer, ShaderGen>::findVarInTemplate(const string& filename, const string& var, size_t start)
{
size_t i = filename.find(var, start);
if (var == "$UDIM" && i != string::npos)
{
size_t udimPrefix = filename.find("$UDIMPREFIX", start);
if (i == udimPrefix)
{
i = filename.find(var, i + 1);
}
}
return i;
}
template <typename Renderer, typename ShaderGen>
FilePath TextureBaker<Renderer, ShaderGen>::generateTextureFilename(const StringMap& filenameTemplateMap)
{
string bakedImageName = _textureFilenameTemplate;
for (auto& pair : filenameTemplateMap)
{
string replacement = (_texTemplateOverrides.count(pair.first)) ? _texTemplateOverrides[pair.first] : pair.second;
replacement = (filenameTemplateMap.at("$UDIM").empty() && pair.first == "$UDIMPREFIX") ? EMPTY_STRING : replacement;
for (size_t i = 0; (i = findVarInTemplate(bakedImageName, pair.first, i)) != string::npos; i++)
{
bakedImageName.replace(i, pair.first.length(), replacement);
}
}
if (_hashImageNames)
{
std::stringstream hashStream;
hashStream << std::hash<std::string>{}(bakedImageName);
hashStream << "." + getExtension();
bakedImageName = hashStream.str();
}
return _outputImagePath / bakedImageName;
}
template <typename Renderer, typename ShaderGen>
StringMap TextureBaker<Renderer, ShaderGen>::initializeFileTemplateMap(InputPtr input, NodePtr shader, const string& udim)
{
FilePath assetPath = FilePath(shader->getActiveSourceUri());
assetPath.removeExtension();
StringMap filenameTemplateMap;
filenameTemplateMap["$ASSET"] = assetPath.getBaseName();
filenameTemplateMap["$INPUT"] = _bakedInputMap[input->getName()];
filenameTemplateMap["$EXTENSION"] = _extension;
filenameTemplateMap["$MATERIAL"] = _material->getName();
filenameTemplateMap["$SHADINGMODEL"] = shader->getCategory();
filenameTemplateMap["$UDIM"] = udim;
filenameTemplateMap["$UDIMPREFIX"] = DEFAULT_UDIM_PREFIX;
return filenameTemplateMap;
}
template <typename Renderer, typename ShaderGen>
bool TextureBaker<Renderer, ShaderGen>::writeBakedImage(const BakedImage& baked, ImagePtr image)
{
if (!Renderer::_imageHandler->saveImage(baked.filename, image, _flipSavedImage))
{
if (_outputStream)
{
*_outputStream << "Failed to write baked image: " << baked.filename.asString() << std::endl;
}
return false;
}
if (_outputStream)
{
*_outputStream << "Wrote baked image: " << baked.filename.asString() << std::endl;
}
return true;
}
template <typename Renderer, typename ShaderGen>
void TextureBaker<Renderer, ShaderGen>::bakeShaderInputs(NodePtr material, NodePtr shader, GenContext& context, const string& udim)
{
_material = material;
if (!shader)
{
return;
}
std::unordered_map<OutputPtr, InputPtr> bakedOutputMap;
for (InputPtr input : shader->getInputs())
{
OutputPtr output = input->getConnectedOutput();
if (output && !bakedOutputMap.count(output))
{
bakedOutputMap[output] = input;
_bakedInputMap[input->getName()] = input->getName();
// When possible, nodes with world-space outputs are applied outside of the baking process.
NodePtr worldSpaceNode = connectsToWorldSpaceNode(output);
if (worldSpaceNode)
{
output->setConnectedNode(worldSpaceNode->getConnectedNode("in"));
_worldSpaceNodes[input->getName()] = worldSpaceNode;
}
StringMap filenameTemplateMap = initializeFileTemplateMap(input, shader, udim);
bakeGraphOutput(output, context, filenameTemplateMap);
}
else if (bakedOutputMap.count(output))
{
// When the input shares the same output as a previously baked input, we use the already baked input.
_bakedInputMap[input->getName()] = bakedOutputMap[output]->getName();
}
}
// Release all images used to generate this set of shader inputs.
Renderer::_imageHandler->clearImageCache();
}
template <typename Renderer, typename ShaderGen>
void TextureBaker<Renderer, ShaderGen>::bakeGraphOutput(OutputPtr output, GenContext& context, const StringMap& filenameTemplateMap)
{
if (!output)
{
return;
}
bool encodeSrgb = _colorSpace == SRGB_TEXTURE && output->isColorType();
Renderer::getFramebuffer()->setEncodeSrgb(encodeSrgb);
ShaderPtr shader = _generator->generate("BakingShader", output, context);
Renderer::createProgram(shader);
// Render and capture the requested image.
Renderer::renderTextureSpace(getTextureSpaceMin(), getTextureSpaceMax());
string texturefilepath = generateTextureFilename(filenameTemplateMap);
Renderer::captureImage(_frameCaptureImage);
// Construct a baked image record.
BakedImage baked;
baked.filename = texturefilepath;
if (_averageImages)
{
baked.uniformColor = _frameCaptureImage->getAverageColor();
baked.isUniform = true;
}
else if (_frameCaptureImage->isUniformColor(&baked.uniformColor))
{
baked.isUniform = true;
}
_bakedImageMap[output].push_back(baked);
// TODO: Write images to memory rather than to disk.
// Write non-uniform images to disk.
if (!baked.isUniform)
{
writeBakedImage(baked, _frameCaptureImage);
}
}
template <typename Renderer, typename ShaderGen>
void TextureBaker<Renderer, ShaderGen>::optimizeBakedTextures(NodePtr shader)
{
if (!shader)
{
return;
}
// Check for fully uniform outputs.
for (auto& pair : _bakedImageMap)
{
bool outputIsUniform = true;
for (BakedImage& baked : pair.second)
{
if (!baked.isUniform || baked.uniformColor != pair.second[0].uniformColor)
{
outputIsUniform = false;
continue;
}
}
if (outputIsUniform)
{
BakedConstant bakedConstant;
bakedConstant.color = pair.second[0].uniformColor;
_bakedConstantMap[pair.first] = bakedConstant;
}
}
// Check for uniform outputs at their default values.
NodeDefPtr shaderNodeDef = shader->getNodeDef();
if (shaderNodeDef)
{
for (InputPtr shaderInput : shader->getInputs())
{
OutputPtr output = shaderInput->getConnectedOutput();
if (output && _bakedConstantMap.count(output))
{
InputPtr input = shaderNodeDef->getInput(shaderInput->getName());
if (input)
{
Color4 uniformColor = _bakedConstantMap[output].color;
string uniformColorString = getValueStringFromColor(uniformColor, input->getType());
string defaultValueString = input->hasValue() ? input->getValue()->getValueString() : EMPTY_STRING;
if (uniformColorString == defaultValueString)
{
_bakedConstantMap[output].isDefault = true;
}
}
}
}
}
// Remove baked images that have been replaced by constant values.
for (auto& pair : _bakedConstantMap)
{
if (pair.second.isDefault || _optimizeConstants || _averageImages)
{
_bakedImageMap.erase(pair.first);
}
}
}
template <typename Renderer, typename ShaderGen>
DocumentPtr TextureBaker<Renderer, ShaderGen>::generateNewDocumentFromShader(NodePtr shader, const StringVec& udimSet)
{
if (!shader)
{
return nullptr;
}
// Create document.
if (!_bakedTextureDoc || _writeDocumentPerMaterial)
{
_bakedTextureDoc = createDocument();
}
if (shader->getDocument()->hasColorSpace())
{
_bakedTextureDoc->setColorSpace(shader->getDocument()->getColorSpace());
}
// Create node graph and geometry info.
NodeGraphPtr bakedNodeGraph;
if (!_bakedImageMap.empty())
{
_bakedGraphName = _bakedTextureDoc->createValidChildName(_bakedGraphName);
bakedNodeGraph = _bakedTextureDoc->addNodeGraph(_bakedGraphName);
bakedNodeGraph->setColorSpace(_colorSpace);
}
_bakedGeomInfoName = _bakedTextureDoc->createValidChildName(_bakedGeomInfoName);
GeomInfoPtr bakedGeom = !udimSet.empty() ? _bakedTextureDoc->addGeomInfo(_bakedGeomInfoName) : nullptr;
if (bakedGeom)
{
bakedGeom->setGeomPropValue(UDIM_SET_PROPERTY, udimSet, "stringarray");
}
// Create a shader node.
NodePtr bakedShader = _bakedTextureDoc->addNode(shader->getCategory(), shader->getName(), shader->getType());
// Optionally create a material node, connecting it to the new shader node.
if (_material)
{
string materialName = (_texTemplateOverrides.count("$MATERIAL")) ? _texTemplateOverrides["$MATERIAL"] : _material->getName();
NodePtr bakedMaterial = _bakedTextureDoc->addNode(_material->getCategory(), materialName, _material->getType());
for (auto sourceMaterialInput : _material->getInputs())
{
const string& sourceMaterialInputName = sourceMaterialInput->getName();
NodePtr upstreamShader = sourceMaterialInput->getConnectedNode();
if (upstreamShader && (upstreamShader->getNamePath() == shader->getNamePath()))
{
InputPtr bakedMaterialInput = bakedMaterial->getInput(sourceMaterialInputName);
if (!bakedMaterialInput)
{
bakedMaterialInput = bakedMaterial->addInput(sourceMaterialInputName, sourceMaterialInput->getType());
}
bakedMaterialInput->setNodeName(bakedShader->getName());
}
}
}
// Create and connect inputs on the new shader node.
for (ValueElementPtr valueElem : shader->getChildrenOfType<ValueElement>())
{
// Get the source input and its connected output.
InputPtr sourceInput = valueElem->asA<Input>();
if (!sourceInput)
{
continue;
}
OutputPtr output = sourceInput->getConnectedOutput();
// Skip uniform outputs at their default values.
if (output && _bakedConstantMap.count(output) && _bakedConstantMap[output].isDefault)
{
continue;
}
// Find or create the baked input.
const string& sourceName = sourceInput->getName();
const string& sourceType = sourceInput->getType();
InputPtr bakedInput = bakedShader->getInput(sourceName);
if (!bakedInput)
{
bakedInput = bakedShader->addInput(sourceName, sourceType);
}
// Assign image or constant data to the baked input.
if (output)
{
// Store a constant value for uniform outputs.
if (_optimizeConstants && _bakedConstantMap.count(output))
{
Color4 uniformColor = _bakedConstantMap[output].color;
string uniformColorString = getValueStringFromColor(uniformColor, bakedInput->getType());
bakedInput->setValueString(uniformColorString);
if (bakedInput->isColorType())
{
bakedInput->setColorSpace(_colorSpace);
}
continue;
}
if (!_bakedImageMap.empty())
{
// Add the image node.
NodePtr bakedImage = bakedNodeGraph->addNode("image", sourceName, sourceType);
InputPtr input = bakedImage->addInput("file", "filename");
StringMap filenameTemplateMap = initializeFileTemplateMap(bakedInput, shader, udimSet.empty() ? EMPTY_STRING : UDIM_TOKEN);
input->setValueString(generateTextureFilename(filenameTemplateMap));
// Reconstruct any world-space nodes that were excluded from the baking process.
auto worldSpacePair = _worldSpaceNodes.find(sourceInput->getName());
if (worldSpacePair != _worldSpaceNodes.end())
{
NodePtr origWorldSpaceNode = worldSpacePair->second;
if (origWorldSpaceNode)
{
NodePtr newWorldSpaceNode = bakedNodeGraph->addNode(origWorldSpaceNode->getCategory(), sourceName + "_map", sourceType);
newWorldSpaceNode->copyContentFrom(origWorldSpaceNode);
InputPtr mapInput = newWorldSpaceNode->getInput("in");
if (mapInput)
{
mapInput->setNodeName(bakedImage->getName());
}
bakedImage = newWorldSpaceNode;
}
}
// Add the graph output.
OutputPtr bakedOutput = bakedNodeGraph->addOutput(sourceName + "_output", sourceType);
bakedOutput->setConnectedNode(bakedImage);
bakedInput->setConnectedOutput(bakedOutput);
}
}
else
{
bakedInput->copyContentFrom(sourceInput);
}
}
// Generate uniform images and write to disk.
ImagePtr uniformImage = createUniformImage(4, 4, 4, Renderer::_baseType, Color4());
for (const auto& pair : _bakedImageMap)
{
for (const BakedImage& baked : pair.second)
{
if (baked.isUniform)
{
uniformImage->setUniformColor(baked.uniformColor);
writeBakedImage(baked, uniformImage);
}
}
}
// Clear cached information after each material bake
_bakedImageMap.clear();
_bakedConstantMap.clear();
_worldSpaceNodes.clear();
_bakedInputMap.clear();
_material = nullptr;
// Return the baked document on success.
return _bakedTextureDoc;
}
template <typename Renderer, typename ShaderGen>
DocumentPtr TextureBaker<Renderer, ShaderGen>::bakeMaterialToDoc(DocumentPtr doc, const FileSearchPath& searchPath, const string& materialPath,
const StringVec& udimSet, string& documentName)
{
if (_outputStream)
{
*_outputStream << "Processing material: " << materialPath << std::endl;
}
// Set up generator context for material
GenContext genContext(_generator);
genContext.getOptions().targetColorSpaceOverride = LIN_REC709;
genContext.getOptions().fileTextureVerticalFlip = true;
genContext.getOptions().targetDistanceUnit = _distanceUnit;
DefaultColorManagementSystemPtr cms = DefaultColorManagementSystem::create(genContext.getShaderGenerator().getTarget());
cms->loadLibrary(doc);
genContext.registerSourceCodeSearchPath(searchPath);
genContext.getShaderGenerator().setColorManagementSystem(cms);
// Compute the material tag set.
StringVec materialTags = udimSet;
if (materialTags.empty())
{
materialTags.push_back(EMPTY_STRING);
}
ElementPtr elem = doc->getDescendant(materialPath);
if (!elem || !elem->isA<Node>())
{
return nullptr;
}
NodePtr materialNode = elem->asA<Node>();
vector<NodePtr> shaderNodes = getShaderNodes(materialNode);
NodePtr shaderNode = shaderNodes.empty() ? nullptr : shaderNodes[0];
if (!shaderNode)
{
return nullptr;
}
StringResolverPtr resolver = StringResolver::create();
// Iterate over material tags.
for (const string& tag : materialTags)
{
// Always clear any cached implementations before generation.
genContext.clearNodeImplementations();
ShaderPtr hwShader = createShader("Shader", genContext, shaderNode);
if (!hwShader)
{
continue;
}
Renderer::_imageHandler->setSearchPath(searchPath);
resolver->setUdimString(tag);
Renderer::_imageHandler->setFilenameResolver(resolver);
bakeShaderInputs(materialNode, shaderNode, genContext, tag);
// Optimize baked textures.
optimizeBakedTextures(shaderNode);
}
// Link the baked material and textures in a MaterialX document.
documentName = shaderNode->getName();
return generateNewDocumentFromShader(shaderNode, udimSet);
}
template <typename Renderer, typename ShaderGen>
void TextureBaker<Renderer, ShaderGen>::bakeAllMaterials(DocumentPtr doc, const FileSearchPath& searchPath, const FilePath& outputFilename)
{
if (_outputImagePath.isEmpty())
{
_outputImagePath = outputFilename.getParentPath();
if (!_outputImagePath.exists())
{
_outputImagePath.createDirectory();
}
}
std::vector<TypedElementPtr> renderableMaterials = findRenderableElements(doc);
// Compute the UDIM set.
ValuePtr udimSetValue = doc->getGeomPropValue(UDIM_SET_PROPERTY);
StringVec udimSet;
if (udimSetValue && udimSetValue->isA<StringVec>())
{
udimSet = udimSetValue->asA<StringVec>();
}
// Bake all materials in documents to memory.
BakedDocumentVec bakedDocuments;
for (size_t i = 0; i < renderableMaterials.size(); i++)
{
if (_outputStream && i > 0)
{
*_outputStream << std::endl;
}
const TypedElementPtr& element = renderableMaterials[i];
string documentName;
DocumentPtr bakedMaterialDoc = bakeMaterialToDoc(doc, searchPath, element->getNamePath(), udimSet, documentName);
if (_writeDocumentPerMaterial && bakedMaterialDoc)
{
bakedDocuments.push_back(make_pair(documentName, bakedMaterialDoc));
}
}
if (_writeDocumentPerMaterial)
{
// Write documents in memory to disk.
size_t bakeCount = bakedDocuments.size();
for (size_t i = 0; i < bakeCount; i++)
{
if (bakedDocuments[i].second)
{
FilePath writeFilename = outputFilename;
// Add additional filename decorations if there are multiple documents.
if (bakedDocuments.size() > 1)
{
const string extension = writeFilename.getExtension();
writeFilename.removeExtension();
string filenameSeparator = writeFilename.isDirectory() ? EMPTY_STRING : "_";
writeFilename = FilePath(writeFilename.asString() + filenameSeparator + bakedDocuments[i].first + "." + extension);
}
writeToXmlFile(bakedDocuments[i].second, writeFilename);
if (_outputStream)
{
*_outputStream << "Wrote baked document: " << writeFilename.asString() << std::endl;
}
}
}
}
else if (_bakedTextureDoc)
{
writeToXmlFile(_bakedTextureDoc, outputFilename);
if (_outputStream)
{
*_outputStream << "Wrote baked document: " << outputFilename.asString() << std::endl;
}
}
}
template <typename Renderer, typename ShaderGen>
void TextureBaker<Renderer, ShaderGen>::setupUnitSystem(DocumentPtr unitDefinitions)
{
UnitTypeDefPtr distanceTypeDef = unitDefinitions ? unitDefinitions->getUnitTypeDef("distance") : nullptr;
UnitTypeDefPtr angleTypeDef = unitDefinitions ? unitDefinitions->getUnitTypeDef("angle") : nullptr;
if (!distanceTypeDef && !angleTypeDef)
{
return;
}
UnitSystemPtr unitSystem = UnitSystem::create(_generator->getTarget());
if (!unitSystem)
{
return;
}
_generator->setUnitSystem(unitSystem);
UnitConverterRegistryPtr registry = UnitConverterRegistry::create();
registry->addUnitConverter(distanceTypeDef, LinearUnitConverter::create(distanceTypeDef));
registry->addUnitConverter(angleTypeDef, LinearUnitConverter::create(angleTypeDef));
_generator->getUnitSystem()->loadLibrary(unitDefinitions);
_generator->getUnitSystem()->setUnitConverterRegistry(registry);
}
MATERIALX_NAMESPACE_END