From 919004644b78a7e5e4c32fbb2d2e546e0097c10b Mon Sep 17 00:00:00 2001 From: Manuel Wagner Date: Tue, 13 Sep 2022 09:08:25 +0200 Subject: [PATCH] made style transfer world specific --- .../StyleTransferSceneViewExtension.cpp | 22 ++++++++----------- .../Private/StyleTransferSubsystem.cpp | 5 ++++- .../Public/StyleTransferSceneViewExtension.h | 12 ++++++---- .../Private/OutputTensorToSceneColorCS.cpp | 1 - .../Private/SceneColorToInputTensorCS.cpp | 1 - 5 files changed, 21 insertions(+), 20 deletions(-) diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp index 9e4033d0..d3c13d56 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSceneViewExtension.cpp @@ -40,8 +40,8 @@ OutType CastNarrowingSafe(InType InValue) } -FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext) - : FSceneViewExtensionBase(AutoRegister) +FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext) + : FWorldSceneViewExtension(AutoRegister, World) , StyleTransferNetworkWeakPtr(InStyleTransferNetwork) , StyleTransferNetwork(InStyleTransferNetwork) , LinkedViewportClient(AssociatedViewportClient) @@ -50,14 +50,12 @@ FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRe ensure(InStyleTransferNetwork->GetDeviceType() == ENeuralDeviceType::GPU); } -void FStyleTransferSceneViewExtension::SetupViewFamily(FSceneViewFamily& InViewFamily) -{ -} - bool FStyleTransferSceneViewExtension::IsActiveThisFrame_Internal(const FSceneViewExtensionContext& Context) const { check(IsInGameThread()); - return bIsEnabled && *InferenceContext != -1 && StyleTransferNetworkWeakPtr.IsValid(); + return FWorldSceneViewExtension::IsActiveThisFrame_Internal(Context) + && bIsEnabled + && *InferenceContext != -1 && StyleTransferNetworkWeakPtr.IsValid(); } void FStyleTransferSceneViewExtension::AddRescalingTextureCopy(FRDGBuilder& GraphBuilder, FRDGTexture& RDGSourceTexture, FScreenPassRenderTarget& DestinationRenderTarget) @@ -115,22 +113,22 @@ void FStyleTransferSceneViewExtension::PreRenderViewFamily_RenderThread(FRDGBuil { return; const FName RenderCaptureProviderType = IRenderCaptureProvider::GetModularFeatureName(); - if(!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType)) + if (!IModularFeatures::Get().IsModularFeatureAvailable(RenderCaptureProviderType)) return; IRenderCaptureProvider& RenderCaptureProvider = IModularFeatures::Get().GetModularFeature(RenderCaptureProviderType); - if(bIsEnabled && NumFramesCaptured == -1) + if (bIsEnabled && NumFramesCaptured == -1) { RenderCaptureProvider.BeginCapture(&GRHICommandList.GetImmediateCommandList()); NumFramesCaptured = 0; } - if(NumFramesCaptured >= 0) + if (NumFramesCaptured >= 0) { ++NumFramesCaptured; } - if(NumFramesCaptured == 10) + if (NumFramesCaptured == 10) { RenderCaptureProvider.EndCapture(&GRHICommandList.GetImmediateCommandList()); } @@ -219,8 +217,6 @@ void FStyleTransferSceneViewExtension::TextureToTensor(FRDGBuilder& GraphBuilder FScreenPassTexture FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View, const FPostProcessMaterialInputs& InOutInputs) { - const FSceneViewFamily& ViewFamily = *View.Family; - const FScreenPassTexture& SceneColor = InOutInputs.Textures[(uint32)EPostProcessMaterialInput::SceneColor]; if (!EnumHasAnyFlags(SceneColor.Texture->Desc.Flags, TexCreate_ShaderResource)) diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp index dd2d234e..5cdc2000 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Private/StyleTransferSubsystem.cpp @@ -52,10 +52,12 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl UTexture2D* StyleTexture = StyleTransferSettings->StyleTexture.LoadSynchronous(); //UTexture2D* StyleTexture = LoadObject(this, TEXT("/Script/Engine.Texture2D'/StyleTransfer/T_StyleImage.T_StyleImage'")); +#if WITH_EDITOR FTextureCompilingManager::Get().FinishCompilation({StyleTexture}); +#endif UpdateStyle(StyleTexture); //UpdateStyle(FPaths::GetPath("C:\\projects\\realtime-style-transfer\\temp\\style_params_tensor.bin")); - StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef()); + StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension(ViewportClient->GetWorld(), ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef()); } StyleTransferSceneViewExtension->SetEnabled(true); @@ -68,6 +70,7 @@ void UStyleTransferSubsystem::StopStylizingViewport() if (StylePredictionInferenceContext != INDEX_NONE) { StylePredictionNetwork->DestroyInferenceContext(StylePredictionInferenceContext); + StylePredictionInferenceContext = INDEX_NONE; } if (StyleTransferInferenceContext && *StyleTransferInferenceContext != INDEX_NONE) { diff --git a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h index 5f05b458..d37d0dd2 100644 --- a/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h +++ b/Plugins/StyleTransfer/Source/StyleTransfer/Public/StyleTransferSceneViewExtension.h @@ -5,20 +5,25 @@ struct FNeuralTensor; struct FScreenPassRenderTarget; class UNeuralNetwork; -class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase +class FStyleTransferSceneViewExtension : public FWorldSceneViewExtension { public: using Ptr = TSharedPtr; using Ref = TSharedRef; - FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext); + FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef InInferenceContext); // - 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; + + virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override + { + } virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override { @@ -35,7 +40,6 @@ public: 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, FRDGTextureRef SourceTexture, const FNeuralTensor& DestinationTensor); diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp index b2dbacf3..e54bb785 100644 --- a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/OutputTensorToSceneColorCS.cpp @@ -1,7 +1,6 @@ // Copyright Epic Games, Inc. All Rights Reserved. #include "OutputTensorToSceneColorCS.h" -#include "Utils.h" const FIntVector FOutputTensorToSceneColorCS::ThreadGroupSize{8, 8, 1}; diff --git a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp index 1cec857c..c2d7069d 100644 --- a/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp +++ b/Plugins/StyleTransfer/Source/StyleTransferShaders/Private/SceneColorToInputTensorCS.cpp @@ -1,7 +1,6 @@ // Copyright Epic Games, Inc. All Rights Reserved. #include "SceneColorToInputTensorCS.h" -#include "Utils.h" const FIntVector FSceneColorToInputTensorCS::ThreadGroupSize{8, 8, 1};