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"));
}
}
}
