style transfer seems to be active now. Produces garbage images though.

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"

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<float4> 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"

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<IRenderCaptureProvider>(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<FScreenPassRenderTarget> StyleTransferOutputTarget = MakeShared<FScreenPassRenderTarget>(StyleTransferRenderTargetTexture, SceneColor.ViewRect,

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<bool> CVarStyleTransferEnabled(
 	TEXT("r.StyleTransfer.Enabled"),
@@ -44,8 +46,10 @@ void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportCl
 		StyleTransferInferenceContext = MakeShared<int32>(StyleTransferNetwork->CreateInferenceContext());
 		checkf(*StyleTransferInferenceContext != INDEX_NONE, TEXT("Could not create inference context for StyleTransferNetwork"));
 
-		FlushRenderingCommands();
+		UTexture2D* StyleTexture = LoadObject<UTexture2D>(this, TEXT("/Script/Engine.Texture2D'/StyleTransfer/T_StyleImage.T_StyleImage'"));
+		UpdateStyle(StyleTexture);
 		StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(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<float>());
-
-	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);

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<UNeuralNetwork> 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;
 };

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;

Plugins/StyleTransfer/Source/StyleTransfer/StyleTransfer.Build.cs

@@ -42,7 +42,9 @@ public class StyleTransfer : ModuleRules
 				"RenderCore",
 				"Renderer",
 				"Projects",
-				"StyleTransferShaders"
+				"StyleTransferShaders",
+				"PixWinPlugin",
+				"InputDevice",
 			}
 		);
 

Plugins/StyleTransfer/StyleTransfer.uplugin

@@ -32,6 +32,10 @@
     {
       "Name": "NeuralNetworkInference",
       "Enabled": true
+    },
+    {
+      "Name": "PixWinPlugin",
+      "Enabled": true
     }
   ]
 }