From 20144b8b59eea3eec7c6794aa676f14ba111f6c5 Mon Sep 17 00:00:00 2001 From: Manuel Wagner Date: Wed, 31 Aug 2022 18:27:27 +0200 Subject: [PATCH] style transfer seems to be active now. Produces garbage images though. --- .../StyleTransfer/Content/T_StyleImage.uasset | 3 ++ .../Private/OutputTensorToSceneColor.usf | 7 ++--- .../Private/SceneColorToInputTensor.usf | 6 ++-- .../StyleTransferSceneViewExtension.cpp | 29 ++++++++++++++++++- .../Private/StyleTransferSubsystem.cpp | 28 +++++++++++++----- .../Public/StyleTransferSceneViewExtension.h | 11 +++++-- .../Public/StyleTransferSubsystem.h | 3 +- .../StyleTransfer/StyleTransfer.Build.cs | 4 ++- Plugins/StyleTransfer/StyleTransfer.uplugin | 4 +++ 9 files changed, 74 insertions(+), 21 deletions(-) create mode 100644 Plugins/StyleTransfer/Content/T_StyleImage.uasset diff --git a/Plugins/StyleTransfer/Content/T_StyleImage.uasset b/Plugins/StyleTransfer/Content/T_StyleImage.uasset new file mode 100644 index 00000000..33212ee9 --- /dev/null +++ b/Plugins/StyleTransfer/Content/T_StyleImage.uasset @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:936df5dec6fff94f1da76ce02bac39fcaa46848fe7a760c7f9c2fea6b4cdcfb8 +size 2714591 diff --git a/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf index a41de6f2..e5d8f5cf 100644 --- a/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf +++ b/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf @@ -1,7 +1,4 @@ -// Copyright Epic Games, Inc. All Rights Reserved. - -#include "/Engine/Public/Platform.ush" - +// Copyright 2022 Manuel Wagner - All rights reserved // this assumes that the OutputTexture has // the exact same dimensions as InputTensor! @@ -33,3 +30,5 @@ void OutputTensorToSceneColorCS(in const uint3 DispatchThreadID : SV_DispatchThr ); OutputTexture[TextureCoords] = RGBAColor; } + +#include "/Engine/Public/Platform.ush" \ No newline at end of file diff --git a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf index f762454a..207d322d 100644 --- a/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf +++ b/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf @@ -1,6 +1,4 @@ -// Copyright Epic Games, Inc. All Rights Reserved. - -#include "/Engine/Public/Platform.ush" +// Copyright 2022 Manuel Wagner - All rights reserved Texture2D InputTexture; SamplerState InputTextureSampler; @@ -29,3 +27,5 @@ void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThre OutputUAV[GlobalIndex + 1] = TextureValue.g; OutputUAV[GlobalIndex + 2] = TextureValue.b; } + +#include "/Engine/Public/Platform.ush" \ No newline at end of file diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp index 24685448..07024ee7 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp @@ -12,6 +12,9 @@ #include "SceneView.h" #include "ScreenPass.h" #include "CommonRenderResources.h" +#include "IPixWinPlugin.h" +#include "IPixWinPlugin.h" +#include "IRenderCaptureProvider.h" #include "NeuralNetwork.h" #include "RenderGraphEvent.h" #include "PostProcess/PostProcessing.h" @@ -108,6 +111,30 @@ void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProces } } +void FStyleTransferSceneViewExtension::PreRenderViewFamily_RenderThread(FRDGBuilder& GraphBuilder, FSceneViewFamily& InViewFamily) +{ + const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName(); + if(!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType)) + return; + + IRenderCaptureProvider& RenderCaptureProvider = IModularFeatures::Get().GetModularFeature(RenderCaptureProviderType); + if(bIsEnabled && NumFramesCaptured == -1) + { + RenderCaptureProvider.BeginCapture(&GRHICommandList.GetImmediateCommandList()); + NumFramesCaptured = 0; + } + + if(NumFramesCaptured >= 0) + { + ++NumFramesCaptured; + } + + if(NumFramesCaptured == 10) + { + RenderCaptureProvider.EndCapture(&GRHICommandList.GetImmediateCommandList()); + } +} + FRDGTexture* FStyleTransferSceneViewExtension::TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor) { FIntVector SourceTensorDimensions = { @@ -216,7 +243,7 @@ FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap StyleTransferNetwork->Run(GraphBuilder, *InferenceContext); - const FNeuralTensor& StyleTransferContentOutputTensor = StyleTransferNetwork->GetInputTensorForContext(*InferenceContext, 0); + const FNeuralTensor& StyleTransferContentOutputTensor = StyleTransferNetwork->GetOutputTensorForContext(*InferenceContext, 0); FRDGTexture* StyleTransferRenderTargetTexture = TensorToTexture(GraphBuilder, SceneColor.Texture->Desc, StyleTransferContentOutputTensor); TSharedPtr StyleTransferOutputTarget = MakeShared(StyleTransferRenderTargetTexture, SceneColor.ViewRect, diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp index cf85574e..a21aea1e 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp @@ -5,8 +5,10 @@ #include "NeuralNetwork.h" #include "RenderGraphUtils.h" +#include "ScreenPass.h" #include "StyleTransferModule.h" #include "StyleTransferSceneViewExtension.h" +#include "Rendering/Texture2DResource.h" TAutoConsoleVariable CVarStyleTransferEnabled( TEXT("r.StyleTransfer.Enabled"), @@ -44,8 +46,10 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl StyleTransferInferenceContext = MakeShared(StyleTransferNetwork->CreateInferenceContext()); checkf(*StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork")); - FlushRenderingCommands(); + UTexture2D* StyleTexture = LoadObject(this, TEXT("/Script/Engine.Texture2D'/StyleTransfer/T_StyleImage.T_StyleImage'")); + UpdateStyle(StyleTexture); StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef()); + } StyleTransferSceneViewExtension->SetEnabled(true); } @@ -66,17 +70,27 @@ void UStyleTransferSubsystem::StopStylizingViewport() } } -void UStyleTransferSubsystem::UpdateStyle(const FNeuralTensor& StyleImage) +void UStyleTransferSubsystem::UpdateStyle(UTexture2D* StyleTexture) { - checkf(StyleTransferSceneViewExtension.IsValid(), TEXT("Can not update style while not stylizing")); - checkf(StyleTransferInferenceContext.IsValid(), TEXT("Can not update style without inference context")); + checkf(StyleTransferInferenceContext.IsValid() && (*StyleTransferInferenceContext) != INDEX_NONE, TEXT("Can not infer style without inference context")); + checkf(StylePredictionInferenceContext != INDEX_NONE, TEXT("Can not update style without inference context")); - StylePredictionNetwork->SetInputFromArrayCopy(StyleImage.GetArrayCopy()); - - ENQUEUE_RENDER_COMMAND(StylePrediction)([this](FRHICommandListImmediate& RHICommandList) + ENQUEUE_RENDER_COMMAND(StylePrediction)([this, StyleTexture](FRHICommandListImmediate& RHICommandList) { FRDGBuilder GraphBuilder(RHICommandList); + + const FNeuralTensor& InputStyleImageTensor = StylePredictionNetwork->GetInputTensorForContext(StylePredictionInferenceContext, 0); + + FTextureResource* StyleTextureResource = StyleTexture->GetResource(); + if(!StyleTextureResource->IsInitialized()) + { + StyleTextureResource->UpdateRHI(); + } + + FRDGTextureRef RDGStyleTexture = GraphBuilder.RegisterExternalTexture(CreateRenderTarget(StyleTextureResource->TextureRHI, TEXT("StyleInputTexture"))); + FStyleTransferSceneViewExtension::TextureToTensor(GraphBuilder, FScreenPassTexture(RDGStyleTexture), InputStyleImageTensor); + StylePredictionNetwork->Run(GraphBuilder, StylePredictionInferenceContext); const FNeuralTensor& OutputStyleParams = StylePredictionNetwork->GetOutputTensorForContext(StylePredictionInferenceContext, 0); diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h index 1eca292c..7f161c8e 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h @@ -15,6 +15,7 @@ public: // - ISceneViewExtension virtual void SubscribeToPostProcessingPass(EPostProcessingPass Pass, FAfterPassCallbackDelegateArray& InOutPassCallbacks, bool bIsPassEnabled) override; + virtual void PreRenderViewFamily_RenderThread(FRDGBuilder& GraphBuilder, FSceneViewFamily& InViewFamily) override; FScreenPassTexture PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs); virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override; @@ -33,6 +34,12 @@ public: void SetEnabled(bool bInIsEnabled) { bIsEnabled = bInIsEnabled; } bool IsEnabled() const { return bIsEnabled; } + + + 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, const FScreenPassTexture& SourceTexture, const FNeuralTensor& DestinationTensor); + private: /** The actual Network pointer is not tracked so we need a WeakPtr too so we can check its validity on the game thread. */ TWeakObjectPtr StyleTransferNetworkWeakPtr; @@ -44,7 +51,5 @@ private: bool bIsEnabled = true; - void AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget); - FRDGTexture* TensorToTexture(FRDGBuilder& GraphBuilder, const FRDGTextureDesc& BaseDestinationDesc, const FNeuralTensor& SourceTensor); - void TextureToTensor(FRDGBuilder& GraphBuilder, const FScreenPassTexture& SourceTexture, const FNeuralTensor& DestinationTensor); + int32 NumFramesCaptured = -1; }; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h index 023cdd83..bf419bd5 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSubsystem.h @@ -3,7 +3,6 @@ #pragma once #include "CoreMinimal.h" -#include "NeuralTensor.h" #include "StyleTransferSceneViewExtension.h" #include "Subsystems/GameInstanceSubsystem.h" #include "UObject/Object.h" @@ -26,7 +25,7 @@ public: void StartStylizingViewport(FViewportClient* ViewportClient); void StopStylizingViewport(); - void UpdateStyle(const FNeuralTensor& StyleImage); + void UpdateStyle(UTexture2D* StyleTexture); private: FStyleTransferSceneViewExtension::Ptr StyleTransferSceneViewExtension; diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs b/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs index e67641e4..1593a211 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs +++ b/Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs @@ -42,7 +42,9 @@ public class StyleTransfer : ModuleRules "RenderCore", "Renderer", "Projects", - "StyleTransferShaders" + "StyleTransferShaders", + "PixWinPlugin", + "InputDevice", } ); diff --git a/Plugins/StyleTransfer/StyleTransfer.uplugin b/Plugins/StyleTransfer/StyleTransfer.uplugin index 67ed84fa..ae4d515c 100644 --- a/Plugins/StyleTransfer/StyleTransfer.uplugin +++ b/Plugins/StyleTransfer/StyleTransfer.uplugin @@ -32,6 +32,10 @@ { "Name": "NeuralNetworkInference", "Enabled": true + }, + { + "Name": "PixWinPlugin", + "Enabled": true } ] } \ No newline at end of file