made style transfer world specific

This commit is contained in:
Manuel Wagner 2022-09-13 09:08:25 +02:00
parent c194cfa1ea
commit 919004644b
5 changed files with 21 additions and 20 deletions

View File

@ -40,8 +40,8 @@ OutType CastNarrowingSafe(InType InValue)
}
FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext)
: FSceneViewExtensionBase(AutoRegister)
FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> 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)
@ -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))

View File

@ -52,10 +52,12 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl
UTexture2D* StyleTexture = StyleTransferSettings->StyleTexture.LoadSynchronous();
//UTexture2D* StyleTexture = LoadObject<UTexture2D>(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<FStyleTransferSceneViewExtension>(ViewportClient, StyleTransferNetwork, StyleTransferInferenceContext.ToSharedRef());
StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(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)
{

View File

@ -5,20 +5,25 @@ struct FNeuralTensor;
struct FScreenPassRenderTarget;
class UNeuralNetwork;
class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase
class FStyleTransferSceneViewExtension : public FWorldSceneViewExtension
{
public:
using Ptr = TSharedPtr<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
using Ref = TSharedRef<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> InInferenceContext);
FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, UWorld* World, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork, TSharedRef<int32> 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);

View File

@ -1,7 +1,6 @@
// Copyright Epic Games, Inc. All Rights Reserved.
#include "OutputTensorToSceneColorCS.h"
#include "Utils.h"
const FIntVector FOutputTensorToSceneColorCS::ThreadGroupSize{8, 8, 1};

View File

@ -1,7 +1,6 @@
// Copyright Epic Games, Inc. All Rights Reserved.
#include "SceneColorToInputTensorCS.h"
#include "Utils.h"
const FIntVector FSceneColorToInputTensorCS::ThreadGroupSize{8, 8, 1};