singinwhale
/
RealtimeStyleTransferRuntime
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Added ShadowMask as buffer visualization output

This commit is contained in:
Manuel Wagner 2022-10-13 15:32:30 +02:00
parent 7e227f24f2
commit c3318e806a
13 changed files with 90 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
Config/DefaultEngine.ini
View File

@ -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"))

3
Lyra.uproject
View File

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

BIN
Plugins/StyleTransfer/Content/NN_PredictorVGG.uasset (Stored with Git LFS)
View File

Binary file not shown.

BIN
Plugins/StyleTransfer/Content/NN_TransferVGG.uasset (Stored with Git LFS)
View File

Binary file not shown.

BIN
Plugins/StyleTransfer/Content/ShadowMask.uasset (Stored with Git LFS) Normal file
View File

Binary file not shown.

BIN
Plugins/StyleTransfer/Content/ViewNormal.uasset (Stored with Git LFS) Normal file
View File

Binary file not shown.

BIN
Plugins/StyleTransfer/Content/ViewTangent.uasset (Stored with Git LFS) Normal file
View File

Binary file not shown.

8
Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf
View File

@ -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)]

2
Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf
View File

@ -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

2
Plugins/StyleTransfer/Shaders/Private/ShadowMaskToInputTensor.usf
View File

@ -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

94
Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp
View File

@ -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 TextureToTensor(FRDGBuilder& GraphBuilder, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor)
void TextureToTensorRGB(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());
typename ShaderType::FParameters* SceneColorToInputTensorParameters = GraphBuilder.AllocParameters<typename ShaderType::FParameters>();
SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
SceneColorToInputTensorParameters->InputTexture = SourceTexture;
SceneColorToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
SceneColorToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(StyleTransferContentInputBuffer);
SceneColorToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y};
SceneColorToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / SceneColorRenderTargetDimensions.X, 0.5 / SceneColorRenderTargetDimensions.Y);
FIntVector SceneColorToInputTensorGroupCount = FComputeShaderUtils::GetGroupCount(
const FRDGBufferRef DestinationTensorBuffer = GraphBuilder.RegisterExternalBuffer(DestinationTensor.GetPooledBuffer());
FSceneColorToInputTensorCS::FParameters* RgbToInputTensorParameters = GraphBuilder.AllocParameters<FSceneColorToInputTensorCS::FParameters>();
RgbToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(DestinationTensor.Num());
RgbToInputTensorParameters->InputTexture = SourceTexture;
RgbToInputTensorParameters->InputTextureSampler = TStaticSamplerState<SF_Bilinear>::GetRHI();
RgbToInputTensorParameters->OutputUAV = GraphBuilder.CreateUAV(DestinationTensorBuffer);
RgbToInputTensorParameters->OutputDimensions = {InputTensorDimensions.X, InputTensorDimensions.Y};
RgbToInputTensorParameters->HalfPixelUV = FVector2f(0.5f / RgbRenderTargetDimensions.X, 0.5 / RgbRenderTargetDimensions.Y);
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()),
SceneColorToInputTensorParameters,
RDG_EVENT_NAME("TextureToTensorRGB(%s)", FSceneColorToInputTensorCS::StaticType.GetName()),
RgbToInputTensorParameters,
ERDGPassFlags::Compute,
[SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList)
[RgbToInputTensorCS, RgbToInputTensorParameters, ComputeGroupCount](FRHICommandList& RHICommandList)
{
FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS,
*SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount);
FComputeShaderUtils::Dispatch(RHICommandList, RgbToInputTensorCS,
*RgbToInputTensorParameters, ComputeGroupCount);
}
);
}
template <uint8 Channels>
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)
void TextureToTensorGrayscale(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);
GScreenShadowMaskTexture = nullptr;
::TextureToTensorGrayscale(GraphBuilder, GScreenShadowMaskTexture, StyleTransferStyleWeightsInputTensor);
::TextureToTensor<3>(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
::TextureToTensorRGB(GraphBuilder, SceneColor.Texture, StyleTransferContentInputTensor);
StyleTransferNetwork->Run(GraphBuilder, *InferenceContext);

2
Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp
View File

@ -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);

3
Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h
View File

@ -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: