From 6114370bc29c57fc78113d2635927785853f4441 Mon Sep 17 00:00:00 2001 From: John Haddon Date: Fri, 15 Nov 2024 13:22:15 +0000 Subject: [PATCH] GafferML : Add TensorToImage This allows tensors to be converted back to GafferImage images, after they have been processed by the Inference node. --- include/GafferML/TensorToImage.h | 91 +++++++ python/GafferMLTest/TensorToImageTest.py | 130 +++++++++ python/GafferMLTest/__init__.py | 1 + python/GafferMLUI/TensorToImageUI.py | 95 +++++++ python/GafferMLUI/__init__.py | 1 + src/GafferML/TensorToImage.cpp | 329 +++++++++++++++++++++++ src/GafferMLModule/GafferMLModule.cpp | 2 + 7 files changed, 649 insertions(+) create mode 100644 include/GafferML/TensorToImage.h create mode 100644 python/GafferMLTest/TensorToImageTest.py create mode 100644 python/GafferMLUI/TensorToImageUI.py create mode 100644 src/GafferML/TensorToImage.cpp diff --git a/include/GafferML/TensorToImage.h b/include/GafferML/TensorToImage.h new file mode 100644 index 00000000000..de0eafcc35e --- /dev/null +++ b/include/GafferML/TensorToImage.h @@ -0,0 +1,91 @@ +////////////////////////////////////////////////////////////////////////// +// +// Copyright (c) 2024, Cinesite VFX Ltd. All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above +// copyright notice, this list of conditions and the following +// disclaimer. +// +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following +// disclaimer in the documentation and/or other materials provided with +// the distribution. +// +// * Neither the name of John Haddon nor the names of +// any other contributors to this software may be used to endorse or +// promote products derived from this software without specific prior +// written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// +////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include "GafferML/Export.h" +#include "GafferML/TensorPlug.h" + +#include "GafferImage/FlatImageSource.h" + +namespace GafferML +{ + +class GAFFERML_API TensorToImage : public GafferImage::FlatImageSource +{ + + public : + + explicit TensorToImage( const std::string &name=defaultName() ); + ~TensorToImage() override; + + GAFFER_NODE_DECLARE_TYPE( GafferML::TensorToImage, TensorToImageTypeId, GafferImage::FlatImageSource ); + + TensorPlug *tensorPlug(); + const TensorPlug *tensorPlug() const; + + Gaffer::StringVectorDataPlug *channelsPlug(); + const Gaffer::StringVectorDataPlug *channelsPlug() const; + + Gaffer::BoolPlug *interleavedChannelsPlug(); + const Gaffer::BoolPlug *interleavedChannelsPlug() const; + + void affects( const Gaffer::Plug *input, AffectedPlugsContainer &outputs ) const override; + + protected : + + void hashMetadata( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const override; + IECore::ConstCompoundDataPtr computeMetadata( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const override; + + void hashFormat( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const override; + GafferImage::Format computeFormat( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const override; + + void hashDataWindow( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const override; + Imath::Box2i computeDataWindow( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const override; + + void hashChannelNames( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const override; + IECore::ConstStringVectorDataPtr computeChannelNames( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const override; + + void hashChannelData( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const override; + IECore::ConstFloatVectorDataPtr computeChannelData( const std::string &channelName, const Imath::V2i &tileOrigin, const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const override; + + static size_t g_firstPlugIndex; + +}; + +IE_CORE_DECLAREPTR( TensorToImage ) + +} // namespace GafferML diff --git a/python/GafferMLTest/TensorToImageTest.py b/python/GafferMLTest/TensorToImageTest.py new file mode 100644 index 00000000000..b7a2b05d7a5 --- /dev/null +++ b/python/GafferMLTest/TensorToImageTest.py @@ -0,0 +1,130 @@ +########################################################################## +# +# Copyright (c) 2024, Cinesite VFX Ltd. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: +# +# * Redistributions of source code must retain the above +# copyright notice, this list of conditions and the following +# disclaimer. +# +# * Redistributions in binary form must reproduce the above +# copyright notice, this list of conditions and the following +# disclaimer in the documentation and/or other materials provided with +# the distribution. +# +# * Neither the name of John Haddon nor the names of +# any other contributors to this software may be used to endorse or +# promote products derived from this software without specific prior +# written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +########################################################################## + +import unittest + +import imath + +import IECore + +import Gaffer +import GafferTest +import GafferImage +import GafferImageTest +import GafferML + +class TensorToImageTest( GafferImageTest.ImageTestCase ) : + + def testNoInput( self ) : + + node = GafferML.TensorToImage() + with self.assertRaisesRegex( Gaffer.ProcessException, "Empty tensor" ) : + node["out"].dataWindow() + + def testNonMatchingChannels( self ) : + + tensor = GafferML.Tensor( + IECore.Color3fVectorData( [ imath.Color3f( 1, 2, 3 ) ] ), + [ 1, 1, 3 ] + ) + + tensorToImage = GafferML.TensorToImage() + tensorToImage["tensor"].setValue( tensor ) + tensorToImage["interleavedChannels"].setValue( True ) + self.assertEqual( tensorToImage["out"].dataWindow(), imath.Box2i( imath.V2i( 0 ), imath.V2i( 1 ) ) ) + + # Only two channels specified. + + tensorToImage["channels"].setValue( IECore.StringVectorData( [ "R", "G" ] ) ) + self.assertEqual( tensorToImage["out"].channelNames(), IECore.StringVectorData( [ "R", "G" ] ) ) + self.assertEqual( tensorToImage["out"].channelData( "R", imath.V2i( 0 ) )[0], 1 ) + self.assertEqual( tensorToImage["out"].channelData( "G", imath.V2i( 0 ) )[0], 2 ) + + with self.assertRaisesRegex( RuntimeError, 'Invalid channel "B"' ) : + tensorToImage["out"].channelData( "B", imath.V2i( 0 ) ) + + # Duplicate channels specified. We just take the first. + + tensorToImage["channels"].setValue( IECore.StringVectorData( [ "R", "R", "B" ] ) ) + self.assertEqual( tensorToImage["out"].channelNames(), IECore.StringVectorData( [ "R", "B" ] ) ) + self.assertEqual( tensorToImage["out"].channelData( "R", imath.V2i( 0 ) )[0], 1 ) + self.assertEqual( tensorToImage["out"].channelData( "B", imath.V2i( 0 ) )[0], 3 ) + + with self.assertRaisesRegex( RuntimeError, 'Invalid channel "G' ) : + tensorToImage["out"].channelData( "G", imath.V2i( 0 ) ) + + # Too many channels specified. We error if the extra channel is accessed. + + tensorToImage["channels"].setValue( IECore.StringVectorData( [ "R", "G", "B", "A" ] ) ) + self.assertEqual( tensorToImage["out"].channelNames(), IECore.StringVectorData( [ "R", "G", "B", "A" ] ) ) + self.assertEqual( tensorToImage["out"].channelData( "R", imath.V2i( 0 ) )[0], 1 ) + self.assertEqual( tensorToImage["out"].channelData( "G", imath.V2i( 0 ) )[0], 2 ) + self.assertEqual( tensorToImage["out"].channelData( "B", imath.V2i( 0 ) )[0], 3 ) + + with self.assertRaisesRegex( RuntimeError, 'Channel "A" out of range' ) : + tensorToImage["out"].channelData( "A", imath.V2i( 0 ) ) + + # Channels skipped by entering empty strings. + + tensorToImage["channels"].setValue( IECore.StringVectorData( [ "R", "", "B" ] ) ) + self.assertEqual( tensorToImage["out"].channelNames(), IECore.StringVectorData( [ "R", "B" ] ) ) + self.assertEqual( tensorToImage["out"].channelData( "R", imath.V2i( 0 ) )[0], 1 ) + self.assertEqual( tensorToImage["out"].channelData( "B", imath.V2i( 0 ) )[0], 3 ) + + with self.assertRaisesRegex( RuntimeError, 'Invalid channel "G' ) : + tensorToImage["out"].channelData( "G", imath.V2i( 0 ) ) + + def testRoundTripWithImageToTensor( self ) : + + image = GafferImage.Checkerboard() + + imageToTensor = GafferML.ImageToTensor() + imageToTensor["image"].setInput( image["out"] ) + imageToTensor["channels"].setInput( image["out"]["channelNames"]) + + tensorToImage = GafferML.TensorToImage() + tensorToImage["tensor"].setInput( imageToTensor["tensor"] ) + tensorToImage["channels"].setInput( image["out"]["channelNames"]) + + self.assertImagesEqual( tensorToImage["out"], image["out"] ) + + imageToTensor["interleaveChannels"].setValue( True ) + tensorToImage["interleavedChannels"].setValue( True ) + + self.assertImagesEqual( tensorToImage["out"], image["out"] ) + +if __name__ == "__main__": + unittest.main() diff --git a/python/GafferMLTest/__init__.py b/python/GafferMLTest/__init__.py index 6c6c8d2202e..b926385c6b5 100644 --- a/python/GafferMLTest/__init__.py +++ b/python/GafferMLTest/__init__.py @@ -39,6 +39,7 @@ from .DataToTensorTest import DataToTensorTest from .InferenceTest import InferenceTest from .ImageToTensorTest import ImageToTensorTest +from .TensorToImageTest import TensorToImageTest if __name__ == "__main__": import unittest diff --git a/python/GafferMLUI/TensorToImageUI.py b/python/GafferMLUI/TensorToImageUI.py new file mode 100644 index 00000000000..d9e04ef5c77 --- /dev/null +++ b/python/GafferMLUI/TensorToImageUI.py @@ -0,0 +1,95 @@ +########################################################################## +# +# Copyright (c) 2024, Cinesite VFX Ltd. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: +# +# * Redistributions of source code must retain the above +# copyright notice, this list of conditions and the following +# disclaimer. +# +# * Redistributions in binary form must reproduce the above +# copyright notice, this list of conditions and the following +# disclaimer in the documentation and/or other materials provided with +# the distribution. +# +# * Neither the name of John Haddon nor the names of +# any other contributors to this software may be used to endorse or +# promote products derived from this software without specific prior +# written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +########################################################################## + +import Gaffer +import GafferML + +Gaffer.Metadata.registerNode( + + GafferML.TensorToImage, + + plugs = { + + "tensor" : [ + + "description", + """ + The input tensor to be turned into an image. Typically this would be connected + to the output of an Inference node that is doing image processing. + """, + + "plugValueWidget:type", "", + "nodule:type", "GafferUI::StandardNodule", + + ], + + "channels" : [ + + "description", + """ + The names to give to the channels in the output image. These + channels are unpacked from the tensor in the order in which they are + specified. For example, an order of `[ "B", "G", "R" ]` might be + needed for use with models trained on images using OpenCV + conventions. An empty channel name may be used to skip a channel + when unpacking. + """, + + ], + + "interleavedChannels" : [ + + "description", + """ + Indicates that the channels are interleaved in the input tensor, in + which case they will be deinterleaved when converting to the output + image. Whether or not channels are interleaved will depend on the + model from which the tensor is obtained. + """, + + ], + + "out" : [ + + "description", + """ + The output image. + """, + + ], + + } +) diff --git a/python/GafferMLUI/__init__.py b/python/GafferMLUI/__init__.py index 466a103c9e9..446a2d81fac 100644 --- a/python/GafferMLUI/__init__.py +++ b/python/GafferMLUI/__init__.py @@ -37,5 +37,6 @@ from . import DataToTensorUI from . import InferenceUI from . import ImageToTensorUI +from . import TensorToImageUI __import__( "IECore" ).loadConfig( "GAFFER_STARTUP_PATHS", subdirectory = "GafferMLUI" ) diff --git a/src/GafferML/TensorToImage.cpp b/src/GafferML/TensorToImage.cpp new file mode 100644 index 00000000000..3f2ec34716a --- /dev/null +++ b/src/GafferML/TensorToImage.cpp @@ -0,0 +1,329 @@ +////////////////////////////////////////////////////////////////////////// +// +// Copyright (c) 2024, Cinesite VFX Ltd. All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are +// met: +// +// * Redistributions of source code must retain the above +// copyright notice, this list of conditions and the following +// disclaimer. +// +// * Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following +// disclaimer in the documentation and/or other materials provided with +// the distribution. +// +// * Neither the name of John Haddon nor the names of +// any other contributors to this software may be used to endorse or +// promote products derived from this software without specific prior +// written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// +////////////////////////////////////////////////////////////////////////// + +#include "GafferML/TensorToImage.h" + +#include "GafferImage/BufferAlgo.h" +#include "GafferImage/ImageAlgo.h" +#include "GafferImage/Sampler.h" + +#include "Gaffer/Context.h" + +using namespace std; +using namespace Imath; +using namespace IECore; +using namespace Gaffer; +using namespace GafferImage; +using namespace GafferML; + +////////////////////////////////////////////////////////////////////////// +// Internal utilities +////////////////////////////////////////////////////////////////////////// + +namespace +{ + +struct ImageShape +{ + Box2i dataWindow; + size_t numChannels; +}; + +ImageShape imageShape( const Tensor *tensor, bool interleavedChannels ) +{ + if( !tensor->value() ) + { + throw IECore::Exception( "Empty tensor" ); + } + + const auto shape = tensor->value().GetTensorTypeAndShapeInfo().GetShape(); + if( shape.size() < 3 ) + { + throw IECore::Exception( "Expected tensor with at least 3 dimensions" ); + } + + size_t i = shape.size() - 3; + for( size_t d = 0; d < i; ++d ) + { + if( shape[d] != 1 ) + { + throw IECore::Exception( + fmt::format( + "Expected {} dimensional tensor to have size 1 in dimension {}", + shape.size(), d + ) + ); + } + } + + if( interleavedChannels ) + { + return { + Box2i( V2i( 0 ), V2i( (int)shape[i+1], (int)shape[i] ) ), + (size_t)shape[i+2] + }; + } + else + { + return { + Box2i( V2i( 0 ), V2i( (int)shape[i+2], (int)shape[i+1] ) ), + (size_t)shape[i] + }; + } +} + +} // namespace + +////////////////////////////////////////////////////////////////////////// +// TensorToImage +////////////////////////////////////////////////////////////////////////// + +GAFFER_NODE_DEFINE_TYPE( TensorToImage ); + +size_t TensorToImage::g_firstPlugIndex = 0; + +TensorToImage::TensorToImage( const std::string &name ) + : FlatImageSource( name ) +{ + storeIndexOfNextChild( g_firstPlugIndex ); + addChild( new TensorPlug( "tensor" ) ); + addChild( new StringVectorDataPlug( "channels", Plug::In, new StringVectorData( { "R", "G", "B" } ) ) ); + addChild( new BoolPlug( "interleavedChannels" ) ); +} + +TensorToImage::~TensorToImage() +{ +} + +TensorPlug *TensorToImage::tensorPlug() +{ + return getChild( g_firstPlugIndex ); +} + +const TensorPlug *TensorToImage::tensorPlug() const +{ + return getChild( g_firstPlugIndex ); +} + +Gaffer::StringVectorDataPlug *TensorToImage::channelsPlug() +{ + return getChild( g_firstPlugIndex + 1 ); +} + +const Gaffer::StringVectorDataPlug *TensorToImage::channelsPlug() const +{ + return getChild( g_firstPlugIndex + 1 ); +} + +Gaffer::BoolPlug *TensorToImage::interleavedChannelsPlug() +{ + return getChild( g_firstPlugIndex + 2 ); +} + +const Gaffer::BoolPlug *TensorToImage::interleavedChannelsPlug() const +{ + return getChild( g_firstPlugIndex + 2 ); +} + +void TensorToImage::affects( const Gaffer::Plug *input, AffectedPlugsContainer &outputs ) const +{ + FlatImageSource::affects( input, outputs ); + + if( input == channelsPlug() ) + { + outputs.push_back( outPlug()->channelNamesPlug() ); + } + + if( input == tensorPlug() || input == interleavedChannelsPlug() ) + { + outputs.push_back( outPlug()->dataWindowPlug() ); + } + + if( input == tensorPlug() || input == channelsPlug() || input == interleavedChannelsPlug() ) + { + outputs.push_back( outPlug()->channelDataPlug() ); + } + + if( input == outPlug()->dataWindowPlug() ) + { + outputs.push_back( outPlug()->formatPlug() ); + } +} + +void TensorToImage::hashMetadata( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const +{ + h = outPlug()->metadataPlug()->defaultHash(); +} + +IECore::ConstCompoundDataPtr TensorToImage::computeMetadata( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const +{ + return outPlug()->metadataPlug()->defaultValue(); +} + +void TensorToImage::hashFormat( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const +{ + FlatImageSource::hashFormat( parent, context, h ); + outPlug()->dataWindowPlug()->hash( h ); +} + +GafferImage::Format TensorToImage::computeFormat( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const +{ + return Format( outPlug()->dataWindowPlug()->getValue() ); +} + +void TensorToImage::hashDataWindow( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const +{ + FlatImageSource::hashDataWindow( parent, context, h ); + tensorPlug()->hash( h ); + interleavedChannelsPlug()->hash( h ); +} + +Imath::Box2i TensorToImage::computeDataWindow( const Gaffer::Context *context, const ImagePlug *parent ) const +{ + ConstTensorPtr tensor = tensorPlug()->getValue(); + return imageShape( tensor.get(), interleavedChannelsPlug()->getValue() ).dataWindow; +} + +void TensorToImage::hashChannelNames( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const +{ + FlatImageSource::hashChannelNames( parent, context, h ); + channelsPlug()->hash( h ); +} + +IECore::ConstStringVectorDataPtr TensorToImage::computeChannelNames( const Gaffer::Context *context, const GafferImage::ImagePlug *parent ) const +{ + ConstStringVectorDataPtr channels = channelsPlug()->getValue()->copy(); + // `channels` might be in a non-standard order, to facilitate unpacking from + // a shuffled buffer, and could contain duplicates since it's easy to create + // them while shuffling the list in the UI. Sort into a more natural order + // and remove duplicates. + StringVectorDataPtr result = new StringVectorData( ImageAlgo::sortedChannelNames( channels->readable() ) ); + result->writable().erase( + std::unique( + result->writable().begin(), + result->writable().end() + ), + result->writable().end() + ); + // Remove empty channel names since they wouldn't be valid in the output. + result->writable().erase( + std::remove_if( + result->writable().begin(), + result->writable().end(), + [] ( const string &channelName ) { + return channelName.empty(); + } + ), + result->writable().end() + ); + return result; +} + +void TensorToImage::hashChannelData( const GafferImage::ImagePlug *parent, const Gaffer::Context *context, IECore::MurmurHash &h ) const +{ + FlatImageSource::hashChannelData( parent, context, h ); + { + ImagePlug::GlobalScope globalScope( context ); + tensorPlug()->hash( h ); + channelsPlug()->hash( h ); + interleavedChannelsPlug()->hash( h ); + } + + h.append( context->get( ImagePlug::channelNameContextName ) ); + h.append( context->get( ImagePlug::tileOriginContextName ) ); +} + +IECore::ConstFloatVectorDataPtr TensorToImage::computeChannelData( const std::string &channelName, const Imath::V2i &tileOrigin, const Gaffer::Context *context, const ImagePlug *parent ) const +{ + ConstTensorPtr tensorData; + ConstStringVectorDataPtr channelsData; + bool interleavedChannels; + { + ImagePlug::GlobalScope globalScope( context ); + tensorData = tensorPlug()->getValue(); + channelsData = channelsPlug()->getValue(); + interleavedChannels = interleavedChannelsPlug()->getValue(); + } + + const auto channelIt = std::find( channelsData->readable().begin(), channelsData->readable().end(), channelName ); + if( channelIt == channelsData->readable().end() ) + { + throw IECore::Exception( fmt::format( "Invalid channel \"{}\"", channelName ) ); + } + const size_t channelIndex = channelIt - channelsData->readable().begin(); + + const ImageShape imageShape = ::imageShape( tensorData.get(), interleavedChannels ); + if( channelIndex >= imageShape.numChannels ) + { + throw IECore::Exception( fmt::format( "Channel \"{}\" out of range", channelName ) ); + } + + FloatVectorDataPtr outData = new FloatVectorData; + vector &out = outData->writable(); + + const Box2i dataWindow = imageShape.dataWindow; + const Box2i tileBound( tileOrigin, tileOrigin + V2i( ImagePlug::tileSize() ) ); + const Box2i validTileBound = BufferAlgo::intersection( dataWindow, tileBound ); + out.resize( ImagePlug::tileSize() * ImagePlug::tileSize() ); + + const float *sourceData = tensorData->value().GetTensorData(); + size_t sourceStride; + if( interleavedChannels ) + { + sourceData += channelIndex; + sourceStride = imageShape.numChannels; + } + else + { + sourceData += dataWindow.size().x * dataWindow.size().y * channelIndex; + sourceStride = 1; + } + float *dstData = out.data(); + + for( V2i p = validTileBound.min; p.y < validTileBound.max.y ; ++p.y ) + { + size_t srcIndex = BufferAlgo::index( V2i( p.x, dataWindow.max.y - p.y - 1 ), dataWindow ) * sourceStride; + size_t dstIndex = BufferAlgo::index( p, tileBound ); + + for( int x = validTileBound.min.x; x < validTileBound.max.x; ++x ) + { + dstData[dstIndex++] = sourceData[srcIndex]; + srcIndex += sourceStride; + } + } + + return outData; +} diff --git a/src/GafferMLModule/GafferMLModule.cpp b/src/GafferMLModule/GafferMLModule.cpp index 8f1418ad918..dc0e32b16cf 100644 --- a/src/GafferMLModule/GafferMLModule.cpp +++ b/src/GafferMLModule/GafferMLModule.cpp @@ -45,6 +45,7 @@ #include "GafferML/Inference.h" #include "GafferML/Tensor.h" #include "GafferML/TensorPlug.h" +#include "GafferML/TensorToImage.h" #include "IECorePython/RunTimeTypedBinding.h" @@ -238,5 +239,6 @@ BOOST_PYTHON_MODULE( _GafferML ) ; GafferBindings::DependencyNodeClass(); + GafferBindings::DependencyNodeClass(); }