Added ShadowMask as buffer visualization output

Config/DefaultEngine.ini

@@ -49,6 +49,7 @@ gc.PendingKillEnabled=False
 ;LogEOSSdk=VeryVerbose
 ;LogHandshake=VeryVerbose
 LogHotfixManager=Log
+LogNeuralNetworkInference=VeryVerbose
 
 [/Script/Engine.Player]
 ; These numbers should match TotalNetBandwidth
@@ -313,3 +314,9 @@ ManualIPAddress=
 [Internationalization]
 +LocalizationPaths=%GAMEDIR%Content/Localization/EngineOverrides
 
+
+
+[Engine.BufferVisualizationMaterials]
+ViewTangent=(Material="/StyleTransfer/ViewTangent.ViewTangent", Name=LOCTEXT("ViewTangentMat", "View Tangent"))
+ViewNormal=(Material="/StyleTransfer/ViewNormal.ViewNormal", Name=LOCTEXT("ViewNormalMat", "View Normal"))
+ShadowMask=(Material="/StyleTransfer/ShadowMask.ShadowMask", Name=LOCTEXT("ShadowMaskMat", "Shadow Mask"))

Lyra.uproject

@@ -149,7 +149,8 @@
 			"SupportedTargetPlatforms": [
 				"Win64",
 				"Linux",
-				"Android"
+				"Android",
+				"HoloLens"
 			]
 		},
 		{

Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf

@@ -1,11 +1,11 @@
 // Copyright 2022 Manuel Wagner - All rights reserved
 
-// this assumes that the OutputTexture has
-// the exact same dimensions as InputTensor!
-uint TensorVolume;
-uint2 TextureSize;
 RWTexture2D<float4> OutputTexture;
 Buffer<float> InputTensor;
+uint TensorVolume;
+// this assumes that the OutputTexture has
+// the exact same dimensions as InputTensor!
+uint2 TextureSize;
 
 // DispatchThreadID corresponds to InputTensor shape dimensions not texture XY -> DispatchThreadID.X = Texture.Y
 [numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)]

Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf

@@ -15,7 +15,7 @@ void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThre
 		return;
 	}
 
-	const uint GlobalIndex = ((OutputUAVTexelCoordinate.x * OutputDimensions.y + OutputUAVTexelCoordinate.y)*3);
+	const uint GlobalIndex = ((OutputUAVTexelCoordinate.x * OutputDimensions.y  + OutputUAVTexelCoordinate.y)*3);
 
 	// note that the OutputUAV has shape (1, Y, X, C)
 	// which is why we need to flip the indexing

Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf

@@ -15,7 +15,7 @@ void ShadowMaskToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThre
 		return;
 	}
 
-	const uint GlobalIndex = ((OutputUAVTexelCoordinate.x * OutputDimensions.y + OutputUAVTexelCoordinate.y));
+	const uint GlobalIndex = OutputUAVTexelCoordinate.x * OutputDimensions.y + OutputUAVTexelCoordinate.y;
 
 	// note that the OutputUAV has shape (1, Y, X, C)
 	// which is why we need to flip the indexing

Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp

@@ -38,11 +38,11 @@ TAutoConsoleVariable<bool> CVarAutoCaptureStyleTransfer(
 template <class OutType, class InType>
 OutType CastNarrowingSafe(InType InValue)
 {
-	if (InValue > TNumericLimits<OutType>::Max())
+	if (!ensure(InValue <= TNumericLimits<OutType>::Max()))
 	{
 		return TNumericLimits<OutType>::Max();
 	}
-	if (InValue < TNumericLimits<OutType>::Min())
+	if (!ensure(InValue >= TNumericLimits<OutType>::Min()))
 	{
 		return TNumericLimits<OutType>::Min();
 	}
@@ -247,54 +247,79 @@ FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& B
 	return OutputTexture;
 }
 
-template <class ShaderType = FSceneColorToInputTensorCS>
+void TextureToTensorRGB(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
-void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
 {
 	const FIntVector InputTensorDimensions = {
 		CastNarrowingSafe<int32>(DestinationTensor.GetSize(1)),
 		CastNarrowingSafe<int32>(DestinationTensor.GetSize(2)),
 		CastNarrowingSafe<int32>(DestinationTensor.GetSize(3)),
 	};
-	const FIntPoint SceneColorRenderTargetDimensions = SourceTexture->Desc.Extent;
+	const FIntPoint RgbRenderTargetDimensions = SourceTexture->Desc.Extent;
 
-	FRDGBufferRef StyleTransferContentInputBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer());
+	const FRDGBufferRef DestinationTensorBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer());
-	typename ShaderType::FParameters* SceneColorToInputTensorParameters = GraphBuilder.AllocParameters<typename ShaderType::FParameters>();
+	FSceneColorToInputTensorCS::FParameters* RgbToInputTensorParameters = GraphBuilder.AllocParameters<FSceneColorToInputTensorCS::FParameters>();
-	SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
+	RgbToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
-	SceneColorToInputTensorParameters->InputTexture = SourceTexture;
+	RgbToInputTensorParameters->InputTexture = SourceTexture;
-	SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
+	RgbToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
-	SceneColorToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(StyleTransferContentInputBuffer);
+	RgbToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(DestinationTensorBuffer);
-	SceneColorToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y};
+	RgbToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y};
-	SceneColorToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / SceneColorRenderTargetDimensions.X, 0.5 / SceneColorRenderTargetDimensions.Y);
+	RgbToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / RgbRenderTargetDimensions.X, 0.5 / RgbRenderTargetDimensions.Y);
-	FIntVector SceneColorToInputTensorGroupCount = FComputeShaderUtils::GetGroupCount(
+	FIntVector ComputeGroupCount = FComputeShaderUtils::GetGroupCount(
 		{InputTensorDimensions.X, InputTensorDimensions.Y, 1},
-		ShaderType::ThreadGroupSize
+		FSceneColorToInputTensorCS::ThreadGroupSize
 	);
 
-	TShaderMapRef<ShaderType> SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
+	TShaderMapRef<FSceneColorToInputTensorCS> RgbToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
-
 	GraphBuilder.AddPass(
-		RDG_EVENT_NAME("TextureToTensor(%s)", ShaderType::StaticType.GetName()),
+		RDG_EVENT_NAME("TextureToTensorRGB(%s)", FSceneColorToInputTensorCS::StaticType.GetName()),
-		SceneColorToInputTensorParameters,
+		RgbToInputTensorParameters,
 		ERDGPassFlags::Compute,
-		[SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList)
+		[RgbToInputTensorCS, RgbToInputTensorParameters, ComputeGroupCount](FRHICommandList& RHICommandList)
 		{
-			FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS,
+			FComputeShaderUtils::Dispatch(RHICommandList, RgbToInputTensorCS,
-			                              *SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount);
+			                              *RgbToInputTensorParameters, ComputeGroupCount);
 		}
 	);
 }
 
-template <uint8 Channels>
+void TextureToTensorGrayscale(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
-void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);
-template <>
-void TextureToTensor<1>(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor){return TextureToTensor<FShadowMaskToInputTensorCS>(GraphBuilder, SourceTexture, DestinationTensor);}
-template <>
-void TextureToTensor<3>(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor){return TextureToTensor<FSceneColorToInputTensorCS>(GraphBuilder, SourceTexture, DestinationTensor);}
-
-
-void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
 {
-	::TextureToTensor<3>(GraphBuilder, SourceTexture, DestinationTensor);
+	const FIntVector InputTensorDimensions = {
+		CastNarrowingSafe<int32>(DestinationTensor.GetSize(1)),
+		CastNarrowingSafe<int32>(DestinationTensor.GetSize(2)),
+		CastNarrowingSafe<int32>(DestinationTensor.GetSize(3)),
+	};
+	const FIntPoint GrayscaleRenderTargetDimensions = SourceTexture->Desc.Extent;
+
+	const FRDGBufferRef DestinationTensorBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer());
+	FShadowMaskToInputTensorCS::FParameters* GrayscaleToInputTensorParameters = GraphBuilder.AllocParameters<FShadowMaskToInputTensorCS::FParameters>();
+	GrayscaleToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
+	GrayscaleToInputTensorParameters->InputTexture = SourceTexture;
+	GrayscaleToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
+	GrayscaleToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(DestinationTensorBuffer);
+	GrayscaleToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y};
+	GrayscaleToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / GrayscaleRenderTargetDimensions.X, 0.5 / GrayscaleRenderTargetDimensions.Y);
+	FIntVector ComputeGroupCount = FComputeShaderUtils::GetGroupCount(
+		{InputTensorDimensions.X, InputTensorDimensions.Y, 1},
+		FShadowMaskToInputTensorCS::ThreadGroupSize
+	);
+
+	TShaderMapRef<FShadowMaskToInputTensorCS> GrayscaleToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
+	GraphBuilder.AddPass(
+		RDG_EVENT_NAME("TextureToTensorGrayscale(%s)", FShadowMaskToInputTensorCS::StaticType.GetName()),
+		GrayscaleToInputTensorParameters,
+		ERDGPassFlags::Compute,
+		[GrayscaleToInputTensorCS, GrayscaleToInputTensorParameters, ComputeGroupCount](FRHICommandList& RHICommandList)
+		{
+			FComputeShaderUtils::Dispatch(RHICommandList, GrayscaleToInputTensorCS,
+										  *GrayscaleToInputTensorParameters, ComputeGroupCount);
+		}
+	);
+}
+
+void FStyleTransferSceneViewExtension::TextureToTensorRGB(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
+{
+	::TextureToTensorRGB(GraphBuilder, SourceTexture, DestinationTensor);
 }
 
 void FStyleTransferSceneViewExtension::InterpolateTensors(FRDGBuilder& GraphBuilder, const FNeuralTensor& DestinationTensor, const FNeuralTensor& InputTensorA, const FNeuralTensor& InputTensorB, float Alpha)
@@ -359,10 +384,9 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap
 	const FNeuralTensor& StyleTransferStyleWeightsInputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, StyleWeightsInputTensorIndex);
 
 	check(GScreenShadowMaskTexture);
-	::TextureToTensor<1>(GraphBuilder, GScreenShadowMaskTexture, StyleTransferStyleWeightsInputTensor);
+	::TextureToTensorGrayscale(GraphBuilder, GScreenShadowMaskTexture, StyleTransferStyleWeightsInputTensor);
-	GScreenShadowMaskTexture = nullptr;
 
-	::TextureToTensor<3>(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
+	::TextureToTensorRGB(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
 
 	StyleTransferNetwork->Run(GraphBuilder, *InferenceContext);
 

Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp

@@ -144,7 +144,7 @@ void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture, uint32 Style
 			const FNeuralTensor& InputStyleImageTensor = StylePredictionNetwork->GetInputTensorForContext(StylePredictionInferenceContext, 0);
 			FTextureResource* StyleTextureResource = StyleTexture->GetResource();
 			FRDGTextureRef RDGStyleTexture = GraphBuilder.RegisterExternalTexture(CreateRenderTarget(StyleTextureResource->TextureRHI, TEXT("StyleInputTexture")));
-			FStyleTransferSceneViewExtension::TextureToTensor(GraphBuilder, RDGStyleTexture, InputStyleImageTensor);
+			FStyleTransferSceneViewExtension::TextureToTensorRGB(GraphBuilder, RDGStyleTexture, InputStyleImageTensor);
 
 			StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext);
 

Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h

@@ -42,7 +42,8 @@ public:
 
 	static void AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget);
 	static FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor);
-	static void TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);
+	static void TextureToTensorRGB(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);
+	static void TextureToTensorGrayscale(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor);
 	static void InterpolateTensors(FRDGBuilder& GraphBuilder, const FNeuralTensor& DestinationTensor, const FNeuralTensor& InputTensorA, const FNeuralTensor& InputTensorB, float Alpha);
 
 private: