LuckyWorld/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Lib/LuckyWorldFunctions.cpp

// Fill out your copyright notice in the Description page of Project Settings.


#include "Lib/LuckyWorldFunctions.h"
#include "EngineUtils.h"
#include "IImageWrapperModule.h"
#include "IImageWrapper.h"
#include "Modules/ModuleManager.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Modules/ModuleManager.h"
#include "Misc/FileHelper.h"
#include "Engine/Texture2D.h"
#include "RenderUtils.h"
#include "Serialization/JsonWriter.h"
#include "Serialization/JsonSerializer.h"
#include "Dom/JsonObject.h"


APickAndPlaceManager* ULuckyWorldFunctions::GetPickAndPlaceManager(UObject* WorldContextObject)
{
	if (IsValid(WorldContextObject))
	{
		for (TActorIterator<APickAndPlaceManager> Itr(WorldContextObject->GetWorld()); Itr; ++Itr)
		{
			return *Itr;
		}
	}
	return nullptr;
}

bool ULuckyWorldFunctions::SaveRenderTargetToDisk(UTextureRenderTarget2D* RenderTarget, const FString& Filename, bool bAsPNG, bool bFlipVertically)
{
	if (!IsValid(RenderTarget)) return false;

	FRenderTarget* RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	TArray<FColor> Bitmap;
	RenderTargetResource->ReadPixels(Bitmap);

	if (Bitmap.Num() <= 0) return false;

	if (bFlipVertically)
	{
		for (int32 Row = 0; Row < RenderTarget->SizeY / 2; Row++)
		{
			int32 IndexA = Row * RenderTarget->SizeX;
			int32 IndexB = (RenderTarget->SizeY - Row - 1) * RenderTarget->SizeX;
			for (int32 Col = 0; Col < RenderTarget->SizeX; Col++)
			{
				Bitmap.SwapMemory(IndexA + Col, IndexB + Col);
			}
		}
	}

	FString FullPath = Filename;
	FPaths::NormalizeFilename(FullPath);
	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	EImageFormat Format = bAsPNG ? EImageFormat::PNG : EImageFormat::JPEG;
	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(Format);

	ImageWrapper->SetRaw(Bitmap.GetData(), Bitmap.GetAllocatedSize(), RenderTarget->SizeX, RenderTarget->SizeY, ERGBFormat::BGRA, 8);
	const TArray64<uint8>& ImageData = bAsPNG
		? ImageWrapper->GetCompressed(100)
		: ImageWrapper->GetCompressed(75);

	return FFileHelper::SaveArrayToFile(ImageData, *FullPath);
}

bool ULuckyWorldFunctions::SaveTextureToDisk(UTexture2D* Texture, const FString& Filename, bool bAsPNG, bool bFlipVertically)
{
	if (IsValid(Texture) || !Texture->GetPlatformData() || Texture->GetPlatformData()->Mips.Num() == 0)
	{
		return false;
	}

	FTexture2DMipMap& Mip = Texture->GetPlatformData()->Mips[0];
	const void* Data = Mip.BulkData.LockReadOnly();

	int32 Width = Mip.SizeX;
	int32 Height = Mip.SizeY;
	TArray<FColor> Pixels;
	Pixels.AddUninitialized(Width * Height);

	if (Texture->GetPlatformData()->PixelFormat == PF_B8G8R8A8)
	{
		FMemory::Memcpy(Pixels.GetData(), Data, Width * Height * sizeof(FColor));
	}
	else
	{
		Mip.BulkData.Unlock();
		return false;
	}

	Mip.BulkData.Unlock();

	if (bFlipVertically)
	{
		for (int32 Row = 0; Row < Height / 2; Row++)
		{
			int32 IndexA = Row * Width;
			int32 IndexB = (Height - Row - 1) * Width;
			for (int32 Col = 0; Col < Width; Col++)
			{
				Pixels.SwapMemory(IndexA + Col, IndexB + Col);
			}
		}
	}

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	EImageFormat Format = bAsPNG ? EImageFormat::PNG : EImageFormat::JPEG;
	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(Format);

	ImageWrapper->SetRaw(Pixels.GetData(), Pixels.GetAllocatedSize(), Width, Height, ERGBFormat::BGRA, 8);
	const TArray64<uint8>& ImageData = bAsPNG
		? ImageWrapper->GetCompressed(100)
		: ImageWrapper->GetCompressed(75);

	FString FullPath = FPaths::ProjectSavedDir() / TEXT("CapturedImages") / Filename;
	FPaths::NormalizeFilename(FullPath);
	return FFileHelper::SaveArrayToFile(ImageData, *FullPath);
}

FString ULuckyWorldFunctions::ConvertMapToJsonString(const TMap<FString, FString>& InMap)
{
	TSharedRef<FJsonObject> JsonObject = MakeShared<FJsonObject>();

	for (const TPair<FString, FString>& Pair : InMap)
	{
		JsonObject->SetStringField(Pair.Key, Pair.Value);
	}

	FString OutputString;
	TSharedRef<TJsonWriter<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>> Writer =
		TJsonWriterFactory<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>::Create(&OutputString);

	if (FJsonSerializer::Serialize(JsonObject, Writer))
	{
		return OutputString;
	}

	return TEXT("{}");
}

FString ULuckyWorldFunctions::FVectorToJsonString(const FVector& Vector)
{
	TSharedRef<FJsonObject> JsonObject = MakeShared<FJsonObject>();
	JsonObject->SetNumberField(TEXT("x"), Vector.X);
	JsonObject->SetNumberField(TEXT("y"), Vector.Y);
	JsonObject->SetNumberField(TEXT("z"), Vector.Z);

	FString OutputString;

	TSharedRef<TJsonWriter<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>> Writer =
		TJsonWriterFactory<TCHAR, TCondensedJsonPrintPolicy<TCHAR>>::Create(&OutputString);

	if (FJsonSerializer::Serialize(JsonObject, Writer))
	{
		return OutputString;
	}

	return TEXT("{}");
}