TArray<uint8> UGTCaptureComponent::CapturePng(FString Mode)
{
	// Flush location and rotation
	check(CaptureComponents.Num() != 0);
	USceneCaptureComponent2D* CaptureComponent = CaptureComponents.FindRef(Mode);

	TArray<uint8> ImgData;
	if (CaptureComponent == nullptr)
		return ImgData;

	// Attach this to something, for example, a real camera
	const FRotator PawnViewRotation = Pawn->GetViewRotation();
	if (!PawnViewRotation.Equals(CaptureComponent->GetComponentRotation()))
	{
		CaptureComponent->SetWorldRotation(PawnViewRotation);
	}

	UTextureRenderTarget2D* RenderTarget = CaptureComponent->TextureTarget;
	static IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	static TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FColor> Image;
	FTextureRenderTargetResource* RenderTargetResource;
	Image.AddZeroed(Width * Height);
	RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();

	FReadSurfaceDataFlags ReadSurfaceDataFlags;
	ReadSurfaceDataFlags.SetLinearToGamma(false); // This is super important to disable this!
												  // Instead of using this flag, we will set the gamma to the correct value directly
	RenderTargetResource->ReadPixels(Image, ReadSurfaceDataFlags);
	ImageWrapper->SetRaw(Image.GetData(), Image.GetAllocatedSize(), Width, Height, ERGBFormat::BGRA, 8);
	ImgData = ImageWrapper->GetCompressed();

	return ImgData;
}
void SavePng(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
	SCOPE_CYCLE_COUNTER(STAT_SavePng);

	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FColor> Image;
	FTextureRenderTargetResource* RenderTargetResource;
	Image.AddZeroed(Width * Height);
	{
		SCOPE_CYCLE_COUNTER(STAT_GetResource);
		RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	}
	{
		SCOPE_CYCLE_COUNTER(STAT_ReadPixels);
		FReadSurfaceDataFlags ReadSurfaceDataFlags;
		ReadSurfaceDataFlags.SetLinearToGamma(false); // This is super important to disable this!
		// Instead of using this flag, we will set the gamma to the correct value directly
		RenderTargetResource->ReadPixels(Image, ReadSurfaceDataFlags);
	}
	TArray<uint8> ImgData = SerializationUtils::Image2Png(Image, Width, Height);
	FFileHelper::SaveArrayToFile(ImgData, *Filename);
}
TArray<uint8> UGTCaptureComponent::CaptureNpyUint8(FString Mode, int32 Channels)
{
	// Flush location and rotation
	check(CaptureComponents.Num() != 0);
	USceneCaptureComponent2D* CaptureComponent = CaptureComponents.FindRef(Mode);

	TArray<uint8> NpyData;
	if (CaptureComponent == nullptr) {
    UE_LOG(LogUnrealCV, Error, TEXT("Component for mode %s not found. Returning empty array."), *Mode);
    return NpyData;
  }

	// Attach this to something, for example, a real camera
	const FRotator PawnViewRotation = Pawn->GetViewRotation();
	if (!PawnViewRotation.Equals(CaptureComponent->GetComponentRotation()))
	{
		CaptureComponent->SetWorldRotation(PawnViewRotation);
	}

	UTextureRenderTarget2D* RenderTarget = CaptureComponent->TextureTarget;
	int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
	TArray<FColor> ImageData;
	FTextureRenderTargetResource* RenderTargetResource;
	ImageData.AddUninitialized(Width * Height);
	RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
	FReadSurfaceDataFlags ReadSurfaceDataFlags;
	ReadSurfaceDataFlags.SetLinearToGamma(false); // This is super important to disable this!
	// Instead of using this flag, we will set the gamma to the correct value directly
	RenderTargetResource->ReadPixels(ImageData, ReadSurfaceDataFlags);

	// Check the byte order of data
	// Compress image data to npy array
	// Generate a header for the numpy array
	NpyData = NpySerialization(ImageData, Width, Height, Channels);

	return NpyData;
}
void FVisualizeTexture::GenerateContent(const FSceneRenderTargetItem& RenderTargetItem, const FPooledRenderTargetDesc& Desc)
{
	// otherwise StartFrame() wasn't called
	check(ViewRect != FIntRect(0, 0, 0, 0))

	FTexture2DRHIRef VisTexture = (FTexture2DRHIRef&)RenderTargetItem.ShaderResourceTexture;

	if(!IsValidRef(VisTexture) || !Desc.IsValid())
	{
		// todo: improve
		return;
	}

	FIntRect VisualizeTextureRect = ComputeVisualizeTextureRect(Desc.Extent);

	FIntPoint Size = VisualizeTextureRect.Size();

	// clamp to reasonable value to prevent crash
	Size.X = FMath::Max(Size.X, 1);
	Size.Y = FMath::Max(Size.Y, 1);

	FPooledRenderTargetDesc OutputDesc(FPooledRenderTargetDesc::Create2DDesc(Size, PF_B8G8R8A8, TexCreate_None, TexCreate_RenderTargetable | TexCreate_ShaderResource, false));
	
	GRenderTargetPool.FindFreeElement(OutputDesc, VisualizeTextureContent, TEXT("VisualizeTexture"));

	if(!VisualizeTextureContent)
	{
		return;
	}

	const FSceneRenderTargetItem& DestRenderTarget = VisualizeTextureContent->GetRenderTargetItem();

	RHISetRenderTarget(DestRenderTarget.TargetableTexture, FTextureRHIRef());	
	RHIClear(true, FLinearColor(1,1,0,1), false, 0.0f, false, 0, FIntRect());
	RHISetBlendState(TStaticBlendState<>::GetRHI());
	RHISetRasterizerState(TStaticRasterizerState<>::GetRHI());
	RHISetDepthStencilState(TStaticDepthStencilState<false,CF_Always>::GetRHI());

	FIntPoint RTExtent = GSceneRenderTargets.GetBufferSizeXY();

	FVector2D Tex00 = FVector2D(0, 0);
	FVector2D Tex11 = FVector2D(1, 1);

	uint32 LocalVisualizeTextureInputMapping = UVInputMapping;

	if(!Desc.Is2DTexture())
	{
		LocalVisualizeTextureInputMapping = 1;
	}

	// set UV
	switch(LocalVisualizeTextureInputMapping)
	{
		// UV in left top
		case 0:
			Tex11 = FVector2D((float)ViewRect.Width() / RTExtent.X, (float)ViewRect.Height() / RTExtent.Y);
			break;

		// whole texture
		default:
			break;
	}

	bool bIsDefault = StencilSRVSrc == GBlackTexture->TextureRHI;
	bool bDepthStencil = Desc.Is2DTexture() && Desc.Format == PF_DepthStencil;

	//clear if this is a new different Stencil buffer, or it's not a stencil buffer and we haven't switched to the default yet.
	bool bNeedsClear = bDepthStencil && (StencilSRVSrc != RenderTargetItem.TargetableTexture);
	bNeedsClear |= !bDepthStencil && !bIsDefault;
	if (bNeedsClear)
	{
		StencilSRVSrc = nullptr;
		StencilSRV.SafeRelease();
	}

	//always set something into the StencilSRV slot for platforms that require a full resource binding, even if
	//dynamic branching will cause them not to be used.	
	if(bDepthStencil && !StencilSRVSrc)
	{
		StencilSRVSrc = RenderTargetItem.TargetableTexture;
		StencilSRV = RHICreateShaderResourceView((FTexture2DRHIRef&) RenderTargetItem.TargetableTexture, 0, 1, PF_X24_G8);
	}
	else if(!StencilSRVSrc)
	{
		StencilSRVSrc = GBlackTexture->TextureRHI;
		StencilSRV = RHICreateShaderResourceView((FTexture2DRHIRef&) GBlackTexture->TextureRHI, 0, 1, PF_B8G8R8A8);
	}	

	FVisualizeTextureData VisualizeTextureData(RenderTargetItem, Desc);

	bool bDepthTexture = (Desc.TargetableFlags & TexCreate_DepthStencilTargetable) != 0;
	
	VisualizeTextureData.RGBMul = RGBMul;
	VisualizeTextureData.AMul = AMul;
	VisualizeTextureData.Tex00 = Tex00;
	VisualizeTextureData.Tex11 = Tex11;
	VisualizeTextureData.bSaturateInsteadOfFrac = (Flags & 1) != 0;
	VisualizeTextureData.InputValueMapping = bDepthTexture ? 1 : 0;
	VisualizeTextureData.ArrayIndex = ArrayIndex;
	VisualizeTextureData.CustomMip = CustomMip;
	VisualizeTextureData.StencilSRV = StencilSRV;

	{	
		SCOPED_DRAW_EVENT(VisualizeTexture, DEC_SCENE_ITEMS);
		// continue rendering to HDR if necessary
		RenderVisualizeTexture(VisualizeTextureData);
	}
	RHICopyToResolveTarget(DestRenderTarget.TargetableTexture, DestRenderTarget.ShaderResourceTexture, false, FResolveParams());

	VisualizeTextureDesc = Desc;

	// save to disk
	if(bSaveBitmap)
	{
		bSaveBitmap = false;

		uint32 MipAdjustedExtentX = FMath::Clamp(Desc.Extent.X >> CustomMip, 0, Desc.Extent.X);
		uint32 MipAdjustedExtentY = FMath::Clamp(Desc.Extent.Y >> CustomMip, 0, Desc.Extent.Y);
		FIntPoint Extent(MipAdjustedExtentX, MipAdjustedExtentY);

		FReadSurfaceDataFlags ReadDataFlags;
		ReadDataFlags.SetLinearToGamma(false);
		ReadDataFlags.SetOutputStencil(bOutputStencil);		
		ReadDataFlags.SetMip(CustomMip);

		FTextureRHIRef Texture = RenderTargetItem.TargetableTexture ? RenderTargetItem.TargetableTexture : RenderTargetItem.ShaderResourceTexture;

		check(Texture);

		TArray<FColor> Bitmap;
		
		
		
		
		RHIReadSurfaceData(Texture, FIntRect(0, 0, Extent.X, Extent.Y), Bitmap, ReadDataFlags);

		// if the format and texture type is supported
		if(Bitmap.Num())
		{
			// Create screenshot folder if not already present.
			IFileManager::Get().MakeDirectory(*FPaths::ScreenShotDir(), true);

			const FString ScreenFileName(FPaths::ScreenShotDir() / TEXT("VisualizeTexture"));

			uint32 ExtendXWithMSAA = Bitmap.Num() / Extent.Y;

			// Save the contents of the array to a bitmap file. (24bit only so alpha channel is dropped)
			FFileHelper::CreateBitmap(*ScreenFileName, ExtendXWithMSAA, Extent.Y, Bitmap.GetTypedData());	

			UE_LOG(LogConsoleResponse, Display, TEXT("Content was saved to \"%s\""), *FPaths::ScreenShotDir());
		}
		else
		{
			UE_LOG(LogConsoleResponse, Error, TEXT("Failed to save BMP for VisualizeTexture, format or texture type is not supported"));
		}
	}
}