added style transfer plugin
This commit is contained in:
parent
7eb843d6c2
commit
97c5d9ed57
|
@ -285,6 +285,10 @@
|
|||
{
|
||||
"Name": "TopDownArena",
|
||||
"Enabled": true
|
||||
},
|
||||
{
|
||||
"Name": "StyleTransfer",
|
||||
"Enabled": true
|
||||
}
|
||||
],
|
||||
"EpicSampleNameHash": "451731683"
|
||||
|
|
Binary file not shown.
|
@ -0,0 +1,30 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#include "/Engine/Public/Platform.ush"
|
||||
|
||||
|
||||
RWTexture2D<float4> OutputTexture;
|
||||
Buffer<float> InputTensor;
|
||||
|
||||
[numthreads(1, 1, 1)]
|
||||
void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID)
|
||||
{
|
||||
const uint GlobalIndex = DispatchThreadID.x;
|
||||
|
||||
uint TensorVolume;
|
||||
InputTensor.GetDimensions(TensorVolume);
|
||||
|
||||
if (GlobalIndex >= TensorVolume)
|
||||
{
|
||||
return;
|
||||
}
|
||||
uint2 TextureSize = 0;
|
||||
OutputTexture.GetDimensions(TextureSize.x, TextureSize.y);
|
||||
const uint PixelIndex = GlobalIndex / 3;
|
||||
const uint ChannelIndex = GlobalIndex % 3;
|
||||
const uint MajorIndex = PixelIndex / TextureSize.x;
|
||||
const uint MinorIndex = PixelIndex % TextureSize.x;
|
||||
const uint2 TextureCoords = uint2(MajorIndex, MinorIndex);
|
||||
|
||||
OutputTexture[TextureCoords][ChannelIndex] = InputTensor[GlobalIndex];
|
||||
}
|
|
@ -0,0 +1,27 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#include "/Engine/Public/Platform.ush"
|
||||
|
||||
Texture2D<float4> InputTexture;
|
||||
RWBuffer<float> OutputUAV;
|
||||
|
||||
[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, THREADGROUP_SIZE_Z)]
|
||||
void SceneColorToInputTensorCS(in const uint3 DispatchThreadID : SV_DispatchThreadID)
|
||||
{
|
||||
uint2 TexelCoordinate = DispatchThreadID.xy;
|
||||
|
||||
const uint GlobalIndex = TexelCoordinate.x * TexelCoordinate.y * 3;
|
||||
|
||||
uint TensorVolume;
|
||||
OutputUAV.GetDimensions(TensorVolume);
|
||||
|
||||
if (GlobalIndex >= TensorVolume)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
float4 TextureValue = InputTexture[DispatchThreadID.xy];
|
||||
OutputUAV[GlobalIndex + 0] = TextureValue.r;
|
||||
OutputUAV[GlobalIndex + 1] = TextureValue.g;
|
||||
OutputUAV[GlobalIndex + 2] = TextureValue.b;
|
||||
}
|
|
@ -0,0 +1,23 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#include "StyleTransferModule.h"
|
||||
|
||||
#include "ShaderCore.h"
|
||||
#include "Interfaces/IPluginManager.h"
|
||||
#include "Logging/LogMacros.h"
|
||||
|
||||
DEFINE_LOG_CATEGORY(LogStyleTransfer)
|
||||
|
||||
#define LOCTEXT_NAMESPACE "FStyleTransferModule"
|
||||
|
||||
void FStyleTransferModule::StartupModule()
|
||||
{
|
||||
}
|
||||
|
||||
void FStyleTransferModule::ShutdownModule()
|
||||
{
|
||||
}
|
||||
|
||||
#undef LOCTEXT_NAMESPACE
|
||||
|
||||
IMPLEMENT_MODULE(FStyleTransferModule, StyleTransfer)
|
|
@ -0,0 +1,158 @@
|
|||
#include "StyleTransferSceneViewExtension.h"
|
||||
|
||||
#include "CoreGlobals.h"
|
||||
#include "SceneView.h"
|
||||
|
||||
#include "GlobalShader.h"
|
||||
#include "PipelineStateCache.h"
|
||||
#include "RenderTargetPool.h"
|
||||
#include "Shader.h"
|
||||
|
||||
#include "RHI.h"
|
||||
#include "SceneView.h"
|
||||
#include "ScreenPass.h"
|
||||
#include "CommonRenderResources.h"
|
||||
#include "NeuralNetwork.h"
|
||||
#include "RenderGraphEvent.h"
|
||||
#include "PostProcess/PostProcessing.h"
|
||||
#include "Containers/DynamicRHIResourceArray.h"
|
||||
#include "PostProcess/PostProcessMaterial.h"
|
||||
#include "OutputTensorToSceneColorCS.h"
|
||||
#include "SceneColorToInputTensorCS.h"
|
||||
|
||||
template <class OutType, class InType>
|
||||
OutType CastNarrowingSafe(InType InValue)
|
||||
{
|
||||
if (InValue > TNumericLimits<OutType>::Max())
|
||||
{
|
||||
return TNumericLimits<OutType>::Max();
|
||||
}
|
||||
if (InValue < TNumericLimits<OutType>::Min())
|
||||
{
|
||||
return TNumericLimits<OutType>::Min();
|
||||
}
|
||||
return static_cast<OutType>(InValue);
|
||||
}
|
||||
|
||||
|
||||
FStyleTransferSceneViewExtension::FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork)
|
||||
: FSceneViewExtensionBase(AutoRegister)
|
||||
, StyleTransferNetwork(InStyleTransferNetwork)
|
||||
, LinkedViewportClient(AssociatedViewportClient)
|
||||
{
|
||||
ensure(InStyleTransferNetwork->GetDeviceType() == ENeuralDeviceType::GPU);
|
||||
}
|
||||
|
||||
void FStyleTransferSceneViewExtension::SetupViewFamily(FSceneViewFamily& InViewFamily)
|
||||
{
|
||||
InferenceContext = StyleTransferNetwork->CreateInferenceContext();
|
||||
}
|
||||
|
||||
void FStyleTransferSceneViewExtension::SubscribeToPostProcessingPass(EPostProcessingPass PassId,
|
||||
FAfterPassCallbackDelegateArray&
|
||||
InOutPassCallbacks, bool bIsPassEnabled)
|
||||
{
|
||||
if (PassId == EPostProcessingPass::Tonemap)
|
||||
{
|
||||
InOutPassCallbacks.Add(
|
||||
FAfterPassCallbackDelegate::CreateRaw(
|
||||
this, &FStyleTransferSceneViewExtension::PostProcessPassAfterTonemap_RenderThread));
|
||||
}
|
||||
}
|
||||
|
||||
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))
|
||||
{
|
||||
return SceneColor;
|
||||
}
|
||||
|
||||
if (!SceneColor.IsValid())
|
||||
{
|
||||
return SceneColor;
|
||||
}
|
||||
|
||||
RDG_EVENT_SCOPE(GraphBuilder, "StyleTransfer");
|
||||
|
||||
FScreenPassRenderTarget BackBufferRenderTarget;
|
||||
|
||||
// If the override output is provided it means that this is the last pass in post processing.
|
||||
if (InOutInputs.OverrideOutput.IsValid())
|
||||
{
|
||||
BackBufferRenderTarget = InOutInputs.OverrideOutput;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Reusing the same output description for our back buffer as SceneColor when it's not overriden
|
||||
FRDGTextureDesc OutputDesc = SceneColor.Texture->Desc;
|
||||
OutputDesc.Flags |= TexCreate_RenderTargetable;
|
||||
FLinearColor ClearColor(0., 0., 0., 0.);
|
||||
OutputDesc.ClearValue = FClearValueBinding(ClearColor);
|
||||
|
||||
FRDGTexture* BackBufferRenderTargetTexture = GraphBuilder.CreateTexture(
|
||||
OutputDesc, TEXT("BackBufferRenderTargetTexture"));
|
||||
BackBufferRenderTarget = FScreenPassRenderTarget(BackBufferRenderTargetTexture, SceneColor.ViewRect,
|
||||
ERenderTargetLoadAction::EClear);
|
||||
}
|
||||
|
||||
//Get input and output viewports. Backbuffer could be targeting a different region than input viewport
|
||||
const FScreenPassTextureViewport SceneColorViewport(SceneColor);
|
||||
const FScreenPassTextureViewport BackBufferViewport(BackBufferRenderTarget);
|
||||
|
||||
FScreenPassRenderTarget SceneColorRenderTarget(SceneColor, ERenderTargetLoadAction::ELoad);
|
||||
|
||||
checkSlow(View.bIsViewInfo);
|
||||
const FViewInfo& ViewInfo = static_cast<const FViewInfo&>(View);
|
||||
|
||||
/*AddDrawScreenPass(GraphBuilder, RDG_EVENT_NAME("ProcessOCIOColorSpaceXfrm"), ViewInfo, BackBufferViewport,
|
||||
SceneColorViewport, OCIOPixelShader, Parameters);*/
|
||||
|
||||
auto SceneColorToInputTensorParameters = GraphBuilder.AllocParameters<FSceneColorToInputTensorCS::FParameters>();
|
||||
FNeuralTensor InputTensor = StyleTransferNetwork->GetInputTensor();
|
||||
SceneColorToInputTensorParameters->TensorVolume = CastNarrowingSafe<uint32>(InputTensor.Num());
|
||||
SceneColorToInputTensorParameters->InputTexture = SceneColorRenderTarget.Texture;
|
||||
SceneColorToInputTensorParameters->OutputUAV = InputTensor.GetBufferUAVRef();
|
||||
FIntVector SceneColorToInputTensorGroupCount;
|
||||
|
||||
TShaderMapRef<FSceneColorToInputTensorCS> SceneColorToInputTensorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
|
||||
ClearUnusedGraphResources(SceneColorToInputTensorCS, SceneColorToInputTensorParameters);
|
||||
GraphBuilder.AddPass(
|
||||
RDG_EVENT_NAME("SceneColorToInputTensor"),
|
||||
SceneColorToInputTensorParameters,
|
||||
ERDGPassFlags::Compute,
|
||||
[SceneColorToInputTensorCS, SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount](FRHICommandList& RHICommandList)
|
||||
{
|
||||
FComputeShaderUtils::Dispatch(RHICommandList, SceneColorToInputTensorCS,
|
||||
*SceneColorToInputTensorParameters, SceneColorToInputTensorGroupCount);
|
||||
}
|
||||
);
|
||||
|
||||
StyleTransferNetwork->GetInputDataPointerMutableForContext(InferenceContext, 0);
|
||||
StyleTransferNetwork->Run(GraphBuilder, InferenceContext);
|
||||
|
||||
auto OutputTensorToSceneColorParameters = GraphBuilder.AllocParameters<FOutputTensorToSceneColorCS::FParameters>();
|
||||
FNeuralTensor OutputTensor = StyleTransferNetwork->GetOutputTensor(0);
|
||||
OutputTensorToSceneColorParameters->InputTensor = OutputTensor.GetBufferSRVRef();
|
||||
OutputTensorToSceneColorParameters->OutputTexture = GraphBuilder.CreateUAV(SceneColorRenderTarget.Texture);
|
||||
FIntVector OutputTensorToSceneColorGroupCount;
|
||||
|
||||
TShaderMapRef<FOutputTensorToSceneColorCS> OutputTensorToSceneColorCS(GetGlobalShaderMap(GMaxRHIFeatureLevel));
|
||||
ClearUnusedGraphResources(OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters);
|
||||
GraphBuilder.AddPass(
|
||||
RDG_EVENT_NAME("OutputTensorToSceneColor"),
|
||||
OutputTensorToSceneColorParameters,
|
||||
ERDGPassFlags::Compute,
|
||||
[OutputTensorToSceneColorCS, OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount](FRHICommandList& RHICommandList)
|
||||
{
|
||||
FComputeShaderUtils::Dispatch(RHICommandList, OutputTensorToSceneColorCS,
|
||||
*OutputTensorToSceneColorParameters, OutputTensorToSceneColorGroupCount);
|
||||
}
|
||||
);
|
||||
|
||||
return MoveTemp(BackBufferRenderTarget);
|
||||
}
|
|
@ -0,0 +1,60 @@
|
|||
// Fill out your copyright notice in the Description page of Project Settings.
|
||||
|
||||
|
||||
#include "StyleTransferSubsystem.h"
|
||||
|
||||
#include "NeuralNetwork.h"
|
||||
#include "StyleTransferSceneViewExtension.h"
|
||||
|
||||
void UStyleTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection)
|
||||
{
|
||||
Super::Initialize(Collection);
|
||||
|
||||
StyleTransferNetwork = NewObject<UNeuralNetwork>();
|
||||
|
||||
FString ONNXModelFilePath = TEXT("SOME_PARENT_FOLDER/SOME_ONNX_FILE_NAME.onnx");
|
||||
if (StyleTransferNetwork->Load(ONNXModelFilePath))
|
||||
{
|
||||
if (StyleTransferNetwork->IsGPUSupported())
|
||||
{
|
||||
StyleTransferNetwork->SetDeviceType(ENeuralDeviceType::GPU);
|
||||
}
|
||||
else
|
||||
{
|
||||
StyleTransferNetwork->SetDeviceType(ENeuralDeviceType::CPU);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
UE_LOG(LogTemp, Warning, TEXT("StyleTransferNetwork could not loaded from %s."), *ONNXModelFilePath);
|
||||
}
|
||||
|
||||
StylePredictionNetwork = NewObject<UNeuralNetwork>();
|
||||
|
||||
ONNXModelFilePath = TEXT("SOME_PARENT_FOLDER/SOME_ONNX_FILE_NAME.onnx");
|
||||
if (StylePredictionNetwork->Load(ONNXModelFilePath))
|
||||
{
|
||||
if (StylePredictionNetwork->IsGPUSupported())
|
||||
{
|
||||
StylePredictionNetwork->SetDeviceType(ENeuralDeviceType::GPU);
|
||||
}
|
||||
else
|
||||
{
|
||||
StylePredictionNetwork->SetDeviceType(ENeuralDeviceType::CPU);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
UE_LOG(LogTemp, Warning, TEXT("StylePredictionNetwork could not loaded from %s."), *ONNXModelFilePath);
|
||||
}
|
||||
}
|
||||
|
||||
void UStyleTransferSubsystem::StartStylizingViewport(FViewportClient* ViewportClient)
|
||||
{
|
||||
StyleTransferSceneViewExtension = FSceneViewExtensions::NewExtension<FStyleTransferSceneViewExtension>(ViewportClient, StyleTransferNetwork);
|
||||
}
|
||||
|
||||
void UStyleTransferSubsystem::UpdateStyle(FNeuralTensor StyleImage)
|
||||
{
|
||||
|
||||
}
|
|
@ -0,0 +1,17 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "CoreMinimal.h"
|
||||
#include "Modules/ModuleManager.h"
|
||||
|
||||
DECLARE_LOG_CATEGORY_EXTERN(LogStyleTransfer, Log, All);
|
||||
|
||||
class FStyleTransferModule : public IModuleInterface
|
||||
{
|
||||
public:
|
||||
|
||||
/** IModuleInterface implementation */
|
||||
virtual void StartupModule() override;
|
||||
virtual void ShutdownModule() override;
|
||||
};
|
|
@ -0,0 +1,29 @@
|
|||
#pragma once
|
||||
#include "SceneViewExtension.h"
|
||||
|
||||
class UNeuralNetwork;
|
||||
|
||||
class FStyleTransferSceneViewExtension : public FSceneViewExtensionBase
|
||||
{
|
||||
public:
|
||||
using Ptr = TSharedPtr<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
|
||||
using Ref = TSharedRef<FStyleTransferSceneViewExtension, ESPMode::ThreadSafe>;
|
||||
|
||||
FStyleTransferSceneViewExtension(const FAutoRegister& AutoRegister, FViewportClient* AssociatedViewportClient, UNeuralNetwork* InStyleTransferNetwork);
|
||||
|
||||
// - ISceneViewExtension
|
||||
virtual void SubscribeToPostProcessingPass(EPostProcessingPass Pass, FAfterPassCallbackDelegateArray& InOutPassCallbacks, bool bIsPassEnabled) override;
|
||||
FScreenPassTexture PostProcessPassAfterTonemap_RenderThread(FRDGBuilder& GraphBuilder, const FSceneView& View,
|
||||
const FPostProcessMaterialInputs& InOutInputs);
|
||||
virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override;
|
||||
virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override {};
|
||||
virtual void BeginRenderViewFamily(FSceneViewFamily& InViewFamily) override {};
|
||||
// --
|
||||
|
||||
private:
|
||||
TObjectPtr<UNeuralNetwork> StyleTransferNetwork;
|
||||
|
||||
FViewportClient* LinkedViewportClient;
|
||||
|
||||
int32 InferenceContext = -1;
|
||||
};
|
|
@ -0,0 +1,37 @@
|
|||
// Fill out your copyright notice in the Description page of Project Settings.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "CoreMinimal.h"
|
||||
#include "NeuralTensor.h"
|
||||
#include "StyleTransferSceneViewExtension.h"
|
||||
#include "Subsystems/GameInstanceSubsystem.h"
|
||||
#include "UObject/Object.h"
|
||||
#include "StyleTransferSubsystem.generated.h"
|
||||
|
||||
/**
|
||||
*
|
||||
*/
|
||||
UCLASS()
|
||||
class STYLETRANSFER_API UStyleTransferSubsystem : public UGameInstanceSubsystem
|
||||
{
|
||||
GENERATED_BODY()
|
||||
|
||||
public:
|
||||
// - UGameInstanceSubsystem
|
||||
virtual void Initialize(FSubsystemCollectionBase& Collection) override;
|
||||
// --
|
||||
|
||||
void StartStylizingViewport(FViewportClient* ViewportClient);
|
||||
|
||||
void UpdateStyle(FNeuralTensor StyleImage);
|
||||
|
||||
private:
|
||||
FStyleTransferSceneViewExtension::Ptr StyleTransferSceneViewExtension;
|
||||
|
||||
UPROPERTY()
|
||||
TObjectPtr<UNeuralNetwork> StyleTransferNetwork;
|
||||
|
||||
UPROPERTY()
|
||||
TObjectPtr<UNeuralNetwork> StylePredictionNetwork;
|
||||
};
|
|
@ -0,0 +1,56 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
using System.IO;
|
||||
using UnrealBuildTool;
|
||||
|
||||
public class StyleTransfer : ModuleRules
|
||||
{
|
||||
public StyleTransfer(ReadOnlyTargetRules Target) : base(Target)
|
||||
{
|
||||
PCHUsage = ModuleRules.PCHUsageMode.UseExplicitOrSharedPCHs;
|
||||
|
||||
PublicIncludePaths.AddRange(
|
||||
new string[]
|
||||
{
|
||||
}
|
||||
);
|
||||
|
||||
string EngineDir = Path.GetFullPath(Target.RelativeEnginePath);
|
||||
PrivateIncludePaths.AddRange(
|
||||
new string[]
|
||||
{
|
||||
Path.Combine(EngineDir, "Source/Runtime/Renderer/Private"),
|
||||
}
|
||||
);
|
||||
|
||||
|
||||
PublicDependencyModuleNames.AddRange(
|
||||
new string[]
|
||||
{
|
||||
"Core",
|
||||
"NeuralNetworkInference",
|
||||
}
|
||||
);
|
||||
|
||||
|
||||
PrivateDependencyModuleNames.AddRange(
|
||||
new string[]
|
||||
{
|
||||
"CoreUObject",
|
||||
"Engine",
|
||||
"RHI",
|
||||
"RenderCore",
|
||||
"Renderer",
|
||||
"Projects",
|
||||
"StyleTransferShaders"
|
||||
}
|
||||
);
|
||||
|
||||
|
||||
DynamicallyLoadedModuleNames.AddRange(
|
||||
new string[]
|
||||
{
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
|
@ -0,0 +1,15 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#include "OutputTensorToSceneColorCS.h"
|
||||
#include "Utils.h"
|
||||
|
||||
|
||||
|
||||
void FOutputTensorToSceneColorCS::ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
|
||||
{
|
||||
FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
|
||||
}
|
||||
|
||||
IMPLEMENT_GLOBAL_SHADER(FOutputTensorToSceneColorCS,
|
||||
"/Plugins/StyleTransfer/Shaders/Private/OutputTensorToSceneColor.usf",
|
||||
"OutputTensorToSceneColorCS", SF_Compute); // Path defined in StyleTransferModule.cpp
|
|
@ -0,0 +1,19 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#include "SceneColorToInputTensorCS.h"
|
||||
#include "Utils.h"
|
||||
|
||||
|
||||
|
||||
void FSceneColorToInputTensorCS::ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
|
||||
{
|
||||
FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
|
||||
|
||||
OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_X"), THREADGROUP_SIZE_X);
|
||||
OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Y"), THREADGROUP_SIZE_Y);
|
||||
OutEnvironment.SetDefine(TEXT("THREADGROUP_SIZE_Z"), 1);
|
||||
}
|
||||
|
||||
IMPLEMENT_GLOBAL_SHADER(FSceneColorToInputTensorCS,
|
||||
"/Plugins/StyleTransfer/Shaders/Private/SceneColorToInputTensor.usf",
|
||||
"SceneColorToInputTensorCS", SF_Compute); // Path defined in StyleTransferModule.cpp
|
|
@ -0,0 +1,33 @@
|
|||
#include "StyleTransferShaders.h"
|
||||
|
||||
#include "Interfaces/IPluginManager.h"
|
||||
|
||||
DEFINE_LOG_CATEGORY(LogStyleTransferShaders);
|
||||
|
||||
#define LOCTEXT_NAMESPACE "FStyleTransferShadersModule"
|
||||
|
||||
void FStyleTransferShadersModule::StartupModule()
|
||||
{
|
||||
// Add shader directory
|
||||
const TSharedPtr<IPlugin> StyleTransferPlugin = IPluginManager::Get().FindPlugin(TEXT("StyleTransfer"));
|
||||
if (StyleTransferPlugin.IsValid())
|
||||
{
|
||||
const FString RealShaderDirectory = StyleTransferPlugin->GetBaseDir() / TEXT("Shaders"); // TEXT("../../../Plugins/StyleTransfer/Shaders");
|
||||
const FString VirtualShaderDirectory = TEXT("/Plugins/StyleTransfer/Shaders");
|
||||
AddShaderSourceDirectoryMapping(VirtualShaderDirectory, RealShaderDirectory);
|
||||
}
|
||||
else
|
||||
{
|
||||
UE_LOG(LogStyleTransferShaders, Warning,
|
||||
TEXT("FStyleTransferModule::StartupModule(): StyleTransferPlugin was nullptr, shaders directory not added."));
|
||||
}
|
||||
}
|
||||
|
||||
void FStyleTransferShadersModule::ShutdownModule()
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
#undef LOCTEXT_NAMESPACE
|
||||
|
||||
IMPLEMENT_MODULE(FStyleTransferShadersModule, StyleTransferShaders)
|
|
@ -0,0 +1,34 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
// GPU/RHI/shaders
|
||||
#include "GlobalShader.h"
|
||||
#include "RHI.h"
|
||||
#include "ProfilingDebugging/RealtimeGPUProfiler.h"
|
||||
#include "RenderGraphUtils.h"
|
||||
#include "ShaderParameterUtils.h"
|
||||
|
||||
|
||||
|
||||
class STYLETRANSFERSHADERS_API FOutputTensorToSceneColorCS : public FGlobalShader
|
||||
{
|
||||
DECLARE_GLOBAL_SHADER(FOutputTensorToSceneColorCS);
|
||||
SHADER_USE_PARAMETER_STRUCT(FOutputTensorToSceneColorCS, FGlobalShader)
|
||||
|
||||
BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
|
||||
SHADER_PARAMETER(uint32, TensorVolume)
|
||||
SHADER_PARAMETER_RDG_BUFFER_SRV(Buffer<float>, InputTensor)
|
||||
SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D, OutputTexture)
|
||||
END_SHADER_PARAMETER_STRUCT()
|
||||
|
||||
// - FShader
|
||||
static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters) {
|
||||
return GetMaxSupportedFeatureLevel(Parameters.Platform) >= ERHIFeatureLevel::SM5;
|
||||
}
|
||||
|
||||
static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment);
|
||||
// --
|
||||
|
||||
private:
|
||||
};
|
|
@ -0,0 +1,40 @@
|
|||
// Copyright Epic Games, Inc. All Rights Reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
// GPU/RHI/shaders
|
||||
#include "GlobalShader.h"
|
||||
#include "RHI.h"
|
||||
#include "ProfilingDebugging/RealtimeGPUProfiler.h"
|
||||
#include "RenderGraphUtils.h"
|
||||
#include "ShaderParameterUtils.h"
|
||||
|
||||
|
||||
|
||||
class STYLETRANSFERSHADERS_API FSceneColorToInputTensorCS : public FGlobalShader
|
||||
{
|
||||
DECLARE_GLOBAL_SHADER(FSceneColorToInputTensorCS);
|
||||
SHADER_USE_PARAMETER_STRUCT(FSceneColorToInputTensorCS, FGlobalShader)
|
||||
|
||||
static const uint32 THREADGROUP_SIZE_X = 8;
|
||||
static const uint32 THREADGROUP_SIZE_Y = 8;
|
||||
|
||||
BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
|
||||
// Input variables
|
||||
SHADER_PARAMETER(uint32, TensorVolume)
|
||||
// SRV/UAV variables
|
||||
SHADER_PARAMETER_RDG_BUFFER_UAV(RWBuffer<float>, OutputUAV)
|
||||
// Optional SRV/UAV variables
|
||||
SHADER_PARAMETER_RDG_TEXTURE(Texture2D, InputTexture)
|
||||
END_SHADER_PARAMETER_STRUCT()
|
||||
|
||||
// - FShader
|
||||
static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters) {
|
||||
return GetMaxSupportedFeatureLevel(Parameters.Platform) >= ERHIFeatureLevel::SM5;
|
||||
}
|
||||
|
||||
static void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment);
|
||||
// --
|
||||
|
||||
private:
|
||||
};
|
|
@ -0,0 +1,13 @@
|
|||
#pragma once
|
||||
|
||||
#include "CoreMinimal.h"
|
||||
#include "Modules/ModuleManager.h"
|
||||
|
||||
DECLARE_LOG_CATEGORY_EXTERN(LogStyleTransferShaders, Log, All);
|
||||
|
||||
class FStyleTransferShadersModule : public IModuleInterface
|
||||
{
|
||||
public:
|
||||
virtual void StartupModule() override;
|
||||
virtual void ShutdownModule() override;
|
||||
};
|
|
@ -0,0 +1,29 @@
|
|||
using UnrealBuildTool;
|
||||
|
||||
public class StyleTransferShaders : ModuleRules
|
||||
{
|
||||
public StyleTransferShaders(ReadOnlyTargetRules Target) : base(Target)
|
||||
{
|
||||
PCHUsage = ModuleRules.PCHUsageMode.UseExplicitOrSharedPCHs;
|
||||
|
||||
PublicDependencyModuleNames.AddRange(
|
||||
new string[]
|
||||
{
|
||||
"Core",
|
||||
"RenderCore",
|
||||
}
|
||||
);
|
||||
|
||||
PrivateDependencyModuleNames.AddRange(
|
||||
new string[]
|
||||
{
|
||||
"CoreUObject",
|
||||
"Engine",
|
||||
"Slate",
|
||||
"SlateCore",
|
||||
"RHI",
|
||||
"Projects"
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
|
@ -0,0 +1,37 @@
|
|||
{
|
||||
"FileVersion": 3,
|
||||
"Version": 1,
|
||||
"VersionName": "1.0",
|
||||
"FriendlyName": "StyleTransfer",
|
||||
"Description": "Neural Style Transfer for rendering in different styles",
|
||||
"Category": "Other",
|
||||
"CreatedBy": "Manuel Wagner",
|
||||
"CreatedByURL": "https://singinwhale.com",
|
||||
"DocsURL": "",
|
||||
"MarketplaceURL": "",
|
||||
"CanContainContent": true,
|
||||
"IsBetaVersion": false,
|
||||
"IsExperimentalVersion": false,
|
||||
"Installed": false,
|
||||
"Modules": [
|
||||
{
|
||||
"Name": "StyleTransfer",
|
||||
"Type": "Runtime",
|
||||
"LoadingPhase": "PostConfigInit",
|
||||
"WhitelistPlatforms": [
|
||||
"Win64"
|
||||
]
|
||||
},
|
||||
{
|
||||
"Name": "StyleTransferShaders",
|
||||
"Type": "Runtime",
|
||||
"LoadingPhase": "PostConfigInit"
|
||||
}
|
||||
],
|
||||
"Plugins": [
|
||||
{
|
||||
"Name": "NeuralNetworkInference",
|
||||
"Enabled": true
|
||||
}
|
||||
]
|
||||
}
|
|
@ -41,7 +41,8 @@ public class LyraGame : ModuleRules
|
|||
"CommonLoadingScreen",
|
||||
"Niagara",
|
||||
"AsyncMixin",
|
||||
"ControlFlows"
|
||||
"ControlFlows",
|
||||
"StyleTransfer"
|
||||
}
|
||||
);
|
||||
|
||||
|
|
Loading…
Reference in New Issue