<commit_before>#include "UnrealCVPrivate.h"
#include "GTCapturer.h"
#include "ImageWrapper.h"

void InitCaptureComponent(USceneCaptureComponent2D* CaptureComponent)
{
	// Can not use ESceneCaptureSource::SCS_SceneColorHDR, this option will disable post-processing 
	CaptureComponent->CaptureSource = ESceneCaptureSource::SCS_FinalColorLDR;

	CaptureComponent->TextureTarget = NewObject<UTextureRenderTarget2D>();
	CaptureComponent->TextureTarget->InitAutoFormat(640, 480); // TODO: Update this later
	// CaptureComponent->TextureTarget->TargetGamma = 2.2;
	// CaptureComponent->TextureTarget->TargetGamma = 1 / 2.2;
	// CaptureComponent->TextureTarget->TargetGamma = GEngine->GetDisplayGamma();
	// CaptureComponent->TextureTarget->TargetGamma = 1;
	// CaptureComponent->TextureTarget->TargetGamma = 2.2;
	CaptureComponent->TextureTarget->TargetGamma = 1;
	
	CaptureComponent->RegisterComponentWithWorld(GWorld); // What happened for this?
	// CaptureComponent->AddToRoot(); This is not necessary since it has been attached to the Pawn.
}


void SaveExr(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FFloat16Color> FloatImage;
	FloatImage.AddZeroed(Width * Height);
	FTextureRenderTargetResource* RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	RenderTargetResource->ReadFloat16Pixels(FloatImage);

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::EXR);

	ImageWrapper->SetRaw(FloatImage.GetData(), FloatImage.GetAllocatedSize(), Width, Height, ERGBFormat::RGBA, 16);
	const TArray<uint8>& PngData = ImageWrapper->GetCompressed(ImageCompression::Uncompressed);
	FFileHelper::SaveArrayToFile(PngData, *Filename);
}

void SavePng(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FColor> Image;
	Image.AddZeroed(Width * Height);
	FTextureRenderTargetResource* RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	RenderTargetResource->ReadPixels(Image);

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);

	ImageWrapper->SetRaw(Image.GetData(), Image.GetAllocatedSize(), Width, Height, ERGBFormat::BGRA, 8);
	const TArray<uint8>& HDRData = ImageWrapper->GetCompressed(ImageCompression::Uncompressed);
	FFileHelper::SaveArrayToFile(HDRData, *Filename);
}

UMaterial* LoadMaterial(FString InModeName = TEXT(""))
{
	// Load material for visualization
	static TMap<FString, FString>* MaterialPathMap = nullptr;
	if (MaterialPathMap == nullptr)
	{ 
		// MaterialPathMap = NewObject<TMap<FString, FString> >();
		MaterialPathMap = new TMap<FString, FString>();
		// MaterialPathMap->Add(TEXT("depth"), TEXT("Material'/UnrealCV/SceneDepth.SceneDepth'"));
		MaterialPathMap->Add(TEXT("debug"), TEXT("Material'/UnrealCV/debug.debug'"));
		MaterialPathMap->Add(TEXT("depth"), TEXT("Material'/UnrealCV/debug.debug'"));
		// MaterialPathMap->Add(TEXT("depth"), TEXT("Material'/Game/SceneDepth.SceneDepth'"));
	}

	static TMap<FString, UMaterial*>* StaticMaterialMap = nullptr;
	if (StaticMaterialMap == nullptr)
	{
		StaticMaterialMap = new TMap<FString, UMaterial*>();
		for (auto& Elem : *MaterialPathMap)
		{
			FString ModeName = Elem.Key;
			FString MaterialPath = Elem.Value;
			static ConstructorHelpers::FObjectFinder<UMaterial> Material(*MaterialPath); // ConsturctorHelpers is only available for UObject.

			if (Material.Object != NULL)
			{
				StaticMaterialMap->Add(ModeName, (UMaterial*)Material.Object);
			}
		}
	}

	UMaterial* Material = StaticMaterialMap->FindRef(InModeName);
	if (Material == nullptr)
	{
		UE_LOG(LogTemp, Warning, TEXT("Can not recognize visualization mode %s"), *InModeName);
	}
	return Material;
}

UGTCapturer* UGTCapturer::Create(APawn* InPawn, FString Mode)
{
	UGTCapturer* GTCapturer = NewObject<UGTCapturer>();

	GTCapturer->Pawn = InPawn;
	GTCapturer->AddToRoot();

	// DEPRECATED_FORGAME(4.6, "CaptureComponent2D should not be accessed directly, please use GetCaptureComponent2D() function instead. CaptureComponent2D will soon be private and your code will not compile.")
	USceneCaptureComponent2D* CaptureComponent = NewObject<USceneCaptureComponent2D>(); 
	GTCapturer->CaptureComponent = CaptureComponent;
		
	// CaptureComponent needs to be attached to somewhere immediately, otherwise it will be gc-ed
	CaptureComponent->AttachTo(InPawn->GetRootComponent());
	InitCaptureComponent(CaptureComponent);

	UMaterial* Material = LoadMaterial(Mode);
	if (Mode == "lit") // For rendered images
	{
		FEngineShowFlags& ShowFlags = CaptureComponent->ShowFlags;

		ApplyViewMode(VMI_Lit, true, ShowFlags);
		ShowFlags.SetMaterials(true);
		ShowFlags.SetLighting(true);
		ShowFlags.SetPostProcessing(true);
		ShowFlags.SetTonemapper(true);

		// ToneMapper needs to be enabled, or the screen will be very dark
		// TemporalAA needs to be disabled, otherwise the previous frame might contaminate current frame.
		// Check: https://answers.unrealengine.com/questions/436060/low-quality-screenshot-after-setting-the-actor-pos.html for detail
		// Viewport->EngineShowFlags.SetTemporalAA(false);
		ShowFlags.SetTemporalAA(true);
	}
	else if (Mode == "object_mask") // For object mask
	{ 
		FEngineShowFlags& ShowFlags = CaptureComponent->ShowFlags;<commit_msg>BUGFIX: make SceneCaptureComponent support multiple sources.<commit_after>#include "UnrealCVPrivate.h"
#include "GTCapturer.h"
#include "ImageWrapper.h"

void InitCaptureComponent(USceneCaptureComponent2D* CaptureComponent)
{
	// Can not use ESceneCaptureSource::SCS_SceneColorHDR, this option will disable post-processing 
	CaptureComponent->CaptureSource = ESceneCaptureSource::SCS_FinalColorLDR;

	CaptureComponent->TextureTarget = NewObject<UTextureRenderTarget2D>();
	CaptureComponent->TextureTarget->InitAutoFormat(640, 480); // TODO: Update this later
	// GEngine->GetDisplayGamma(), the default gamma is 2.2
	CaptureComponent->TextureTarget->TargetGamma = 1;
	
	CaptureComponent->RegisterComponentWithWorld(GWorld); // What happened for this?
	// CaptureComponent->AddToRoot(); This is not necessary since it has been attached to the Pawn.
}


void SaveExr(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FFloat16Color> FloatImage;
	FloatImage.AddZeroed(Width * Height);
	FTextureRenderTargetResource* RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	RenderTargetResource->ReadFloat16Pixels(FloatImage);

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::EXR);

	ImageWrapper->SetRaw(FloatImage.GetData(), FloatImage.GetAllocatedSize(), Width, Height, ERGBFormat::RGBA, 16);
	const TArray<uint8>& PngData = ImageWrapper->GetCompressed(ImageCompression::Uncompressed);
	FFileHelper::SaveArrayToFile(PngData, *Filename);
}

void SavePng(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FColor> Image;
	Image.AddZeroed(Width * Height);
	FTextureRenderTargetResource* RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	RenderTargetResource->ReadPixels(Image);

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);

	ImageWrapper->SetRaw(Image.GetData(), Image.GetAllocatedSize(), Width, Height, ERGBFormat::BGRA, 8);
	const TArray<uint8>& HDRData = ImageWrapper->GetCompressed(ImageCompression::Uncompressed);
	FFileHelper::SaveArrayToFile(HDRData, *Filename);
}

UMaterial* LoadMaterial(FString InModeName = TEXT(""))
{
	// Load material for visualization
	static TMap<FString, FString>* MaterialPathMap = nullptr;
	if (MaterialPathMap == nullptr)
	{ 
		MaterialPathMap = new TMap<FString, FString>();
		MaterialPathMap->Add(TEXT("depth"), TEXT("Material'/UnrealCV/SceneDepth.SceneDepth'"));
		MaterialPathMap->Add(TEXT("debug"), TEXT("Material'/UnrealCV/debug.debug'"));
	}

	static TMap<FString, UMaterial*>* StaticMaterialMap = nullptr;
	if (StaticMaterialMap == nullptr)
	{
		StaticMaterialMap = new TMap<FString, UMaterial*>();
		for (auto& Elem : *MaterialPathMap)
		{
			FString ModeName = Elem.Key;
			FString MaterialPath = Elem.Value;
			ConstructorHelpers::FObjectFinder<UMaterial> Material(*MaterialPath); // ConsturctorHelpers is only available for UObject.

			if (Material.Object != NULL)
			{
				StaticMaterialMap->Add(ModeName, (UMaterial*)Material.Object);
			}
		}
	}

	UMaterial* Material = StaticMaterialMap->FindRef(InModeName);
	if (Material == nullptr)
	{
		UE_LOG(LogTemp, Warning, TEXT("Can not recognize visualization mode %s"), *InModeName);
	}
	return Material;
}

UGTCapturer* UGTCapturer::Create(APawn* InPawn, FString Mode)
{
	UGTCapturer* GTCapturer = NewObject<UGTCapturer>();

	GTCapturer->Pawn = InPawn;
	GTCapturer->AddToRoot();

	// DEPRECATED_FORGAME(4.6, "CaptureComponent2D should not be accessed directly, please use GetCaptureComponent2D() function instead. CaptureComponent2D will soon be private and your code will not compile.")
	USceneCaptureComponent2D* CaptureComponent = NewObject<USceneCaptureComponent2D>(); 
	GTCapturer->CaptureComponent = CaptureComponent;
		
	// CaptureComponent needs to be attached to somewhere immediately, otherwise it will be gc-ed
	CaptureComponent->AttachTo(InPawn->GetRootComponent());
	InitCaptureComponent(CaptureComponent);

	UMaterial* Material = LoadMaterial(Mode);
	if (Mode == "lit") // For rendered images
	{
		FEngineShowFlags& ShowFlags = CaptureComponent->ShowFlags;

		ApplyViewMode(VMI_Lit, true, ShowFlags);
		ShowFlags.SetMaterials(true);
		ShowFlags.SetLighting(true);
		ShowFlags.SetPostProcessing(true);
		ShowFlags.SetTonemapper(true);

		// ToneMapper needs to be enabled, otherwise the screen will be very dark
		// TemporalAA needs to be disabled, otherwise the previous frame might contaminate current frame.
		// Check: https://answers.unrealengine.com/questions/436060/low-quality-screenshot-after-setting-the-actor-pos.html for detail
		// Viewport->EngineShowFlags.SetTemporalAA(false);
		ShowFlags.SetTemporalAA(true);
	}
	else if (Mode == "object_mask") // For object mask
	{ 
		FEngineShowFlags& ShowFlags = CaptureComponent->ShowFlags;
{'fraction_non_alphanumeric': 0.11517113783533765, 'fraction_numerical': 0.0077705827937095285, 'mean_word_length': 6.811416184971098, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 111, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2477062', 'n_tokens_mistral': 3433, 'n_tokens_neox': 3221, 'n_words': 649}