Skip to content

Commit

Permalink
Seperate serialization code from GTComponent, make the code more clean.
Browse files Browse the repository at this point in the history
  • Loading branch information
@qiuwch
qiuwch committed Jul 10, 2017
1 parent f94d5b3 commit b50e6be
Show file tree
Hide file tree
Showing 6 changed files with 117 additions and 167 deletions.
84 changes: 67 additions & 17 deletions Source/UnrealCV/Private/Commands/CameraHandler.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,11 @@
#include "CineCameraActor.h"
#include "ObjectPainter.h"
#include "ScreenCapture.h"
#include "Serialization.h"

FExecStatus GetPngBinary(const TArray<FString>& Args, const FString& Mode);
FExecStatus GetDepthNpy(const TArray<FString>& Args);
FExecStatus GetNormalNpy(const TArray<FString>& Args);

FString GetDiskFilename(FString Filename)
{
Expand Down Expand Up @@ -107,16 +112,21 @@ void FCameraCommandHandler::RegisterCommands()
// CommandDispatcher->BindCommand("vget /camera/[uint]/buffer", Cmd, "Get buffer of this camera");

Help = "Return raw binary image data, instead of the image filename";
Cmd = FDispatcherDelegate::CreateLambda([this](const TArray<FString>& Args) { return this->GetPngBinary(Args, TEXT("lit")); });
Cmd = FDispatcherDelegate::CreateLambda([](const TArray<FString>& Args) { return GetPngBinary(Args, TEXT("lit")); });
CommandDispatcher->BindCommand("vget /camera/[uint]/lit png", Cmd, Help);
Cmd = FDispatcherDelegate::CreateLambda([this](const TArray<FString>& Args) { return this->GetPngBinary(Args, TEXT("depth")); });

Cmd = FDispatcherDelegate::CreateLambda([](const TArray<FString>& Args) { return GetPngBinary(Args, TEXT("depth")); });
CommandDispatcher->BindCommand("vget /camera/[uint]/depth png", Cmd, Help);
Cmd = FDispatcherDelegate::CreateLambda([this](const TArray<FString>& Args) { return this->GetPngBinary(Args, TEXT("normal")); });

Cmd = FDispatcherDelegate::CreateLambda([](const TArray<FString>& Args) { return GetPngBinary(Args, TEXT("normal")); });
CommandDispatcher->BindCommand("vget /camera/[uint]/normal png", Cmd, Help);

Cmd = FDispatcherDelegate::CreateLambda([this](const TArray<FString>& Args) { return this->GetNpyBinary(Args, TEXT("depth")); });
Cmd = FDispatcherDelegate::CreateStatic(GetDepthNpy);
CommandDispatcher->BindCommand("vget /camera/[uint]/depth npy", Cmd, Help);

Cmd = FDispatcherDelegate::CreateStatic(GetNormalNpy);
CommandDispatcher->BindCommand("vget /camera/[uint]/normal npy", Cmd, Help);

// TODO: object_mask will be handled differently
}

Expand Down Expand Up @@ -406,30 +416,70 @@ FExecStatus FCameraCommandHandler::GetActorLocation(const TArray<FString>& Args)
return FExecStatus::OK(Message);
}

FExecStatus FCameraCommandHandler::GetPngBinary(const TArray<FString>& Args, const FString& ViewMode)
FExecStatus ParseCamera(const TArray<FString>& Args, UGTCaptureComponent* &OutCamera)
{
int32 CameraId = FCString::Atoi(*Args[0]);

UGTCaptureComponent* GTCapturer = FCaptureManager::Get().GetCamera(CameraId);
if (GTCapturer == nullptr)
OutCamera = FCaptureManager::Get().GetCamera(CameraId);
if (OutCamera == nullptr)
{
return FExecStatus::Error(FString::Printf(TEXT("Invalid camera id %d"), CameraId));
}
return FExecStatus::OK();
}

FExecStatus GetPngBinary(const TArray<FString>& Args, const FString& Mode)
{
UGTCaptureComponent* Camera;
FExecStatus Status = ParseCamera(Args, Camera);
if (Status != FExecStatusType::OK) return Status;

TArray<FColor> ImageData;
int32 Height = 0, Width = 0;
Camera->CaptureImage(Mode, ImageData, Width, Height);
if (ImageData.Num() == 0)
{
return FExecStatus::Error(FString::Printf(TEXT("Failed to read %s data."), *Mode));
}

TArray<uint8> ImgData = GTCapturer->CapturePng(ViewMode);
return FExecStatus::Binary(ImgData);
TArray<uint8> PngBinaryData = SerializationUtils::Image2Png(ImageData, Width, Height);
return FExecStatus::Binary(PngBinaryData);
}

FExecStatus FCameraCommandHandler::GetNpyBinary(const TArray<FString>& Args, const FString& ViewMode)
FExecStatus GetDepthNpy(const TArray<FString>& Args)
{
int32 CameraId = FCString::Atoi(*Args[0]);
UGTCaptureComponent* Camera;
FExecStatus Status = ParseCamera(Args, Camera);
if (Status != FExecStatusType::OK) return Status;

TArray<FFloat16Color> FloatImageData;
int32 Height = 0, Width = 0;
Camera->CaptureFloat16Image("depth", FloatImageData, Width, Height);
if (FloatImageData.Num() == 0)
{
return FExecStatus::Error("Failed to read depth data.");
}

TArray<uint8> NpyBinaryData = SerializationUtils::Array2Npy(FloatImageData, Width, Height, 1);

return FExecStatus::Binary(NpyBinaryData);
}

UGTCaptureComponent* GTCapturer = FCaptureManager::Get().GetCamera(CameraId);
if (GTCapturer == nullptr)
FExecStatus GetNormalNpy(const TArray<FString>& Args)
{
UGTCaptureComponent* Camera;
FExecStatus Status = ParseCamera(Args, Camera);
if (Status != FExecStatusType::OK) return Status;

TArray<FFloat16Color> FloatImageData;
int32 Height = 0, Width = 0;
Camera->CaptureFloat16Image("normal", FloatImageData, Width, Height);
if (FloatImageData.Num() == 0)
{
return FExecStatus::Error(FString::Printf(TEXT("Invalid camera id %d"), CameraId));
return FExecStatus::Error("Failed to read normal data.");
}

TArray<uint8> ImgData = GTCapturer->CaptureNpy(ViewMode);
return FExecStatus::Binary(ImgData);
TArray<uint8> NpyBinaryData = SerializationUtils::Array2Npy(FloatImageData, Width, Height, 3);

return FExecStatus::Binary(NpyBinaryData);

}
5 changes: 0 additions & 5 deletions Source/UnrealCV/Private/Commands/CameraHandler.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,4 @@ class FCameraCommandHandler : public FCommandHandler
/** Get ViewMode data by switching to this viewmode then switch back, can not support multi-camera */
FExecStatus GetObjectInstanceMask(const TArray<FString>& Args);

/** Get raw binary image data instead of filename */
FExecStatus GetPngBinary(const TArray<FString>& Args, const FString& ViewMode);

/** Get raw binary data as an uncompressed numpy array */
FExecStatus GetNpyBinary(const TArray<FString>& Args, const FString& ViewMode);
};
181 changes: 40 additions & 141 deletions Source/UnrealCV/Private/GTCaptureComponent.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
#include "UnrealCVPrivate.h"
#include "GTCaptureComponent.h"
#include "ImageWrapper.h"
#include "ViewMode.h"
#include "cnpy.h"
#include "Serialization.h"

DECLARE_CYCLE_STAT(TEXT("SaveExr"), STAT_SaveExr, STATGROUP_UnrealCV);
DECLARE_CYCLE_STAT(TEXT("SavePng"), STAT_SavePng, STATGROUP_UnrealCV);
Expand Down Expand Up @@ -40,48 +39,31 @@ void SaveExr(UTextureRenderTarget2D* RenderTarget, FString Filename)
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);
{
SCOPE_CYCLE_COUNTER(STAT_SaveFile);
FFileHelper::SaveArrayToFile(PngData, *Filename);
}
TArray<uint8> ExrData = SerializationUtils::Image2Exr(FloatImage, Width, Height);
FFileHelper::SaveArrayToFile(ExrData, *Filename);
}

void SavePng(UTextureRenderTarget2D* RenderTarget, FString Filename)
{
SCOPE_CYCLE_COUNTER(STAT_SavePng);
static IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
static IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);

int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
TArray<FColor> Image;
FTextureRenderTargetResource* RenderTargetResource;
Image.AddZeroed(Width * Height);
{
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);
}
{
SCOPE_CYCLE_COUNTER(STAT_ImageWrapper);
ImageWrapper->SetRaw(Image.GetData(), Image.GetAllocatedSize(), Width, Height, ERGBFormat::BGRA, 8);
}
const TArray<uint8>& ImgData = ImageWrapper->GetCompressed(ImageCompression::Uncompressed);
{
SCOPE_CYCLE_COUNTER(STAT_SaveFile);
FFileHelper::SaveArrayToFile(ImgData, *Filename);
}
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);
}

UMaterial* UGTCaptureComponent::GetMaterial(FString InModeName = TEXT(""))
Expand Down Expand Up @@ -231,137 +213,54 @@ FAsyncRecord* UGTCaptureComponent::Capture(FString Mode, FString InFilename)
return AsyncRecord;
}

TArray<uint8> UGTCaptureComponent::CapturePng(FString Mode)
void UGTCaptureComponent::CaptureImage(const FString& Mode, TArray<FColor>& OutImageData, int32& OutWidth, int32& OutHeight)
{
// 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);
UE_LOG(LogUnrealCV, Warning, TEXT("Can not find a camera to capture %s"), *Mode);
OutWidth = 0; OutHeight = 0;
return;
}

UTextureRenderTarget2D* RenderTarget = CaptureComponent->TextureTarget;
static IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
static IImageWrapperPtr ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);
int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
TArray<FColor> Image;
FTextureRenderTargetResource* RenderTargetResource;
Image.AddZeroed(Width * Height);
OutImageData.Empty();
OutImageData.AddZeroed(RenderTarget->SizeX * RenderTarget->SizeY);
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(ImageCompression::Uncompressed);

return ImgData;
// Instead of using this flag, we will set the gamma to the correct value directly
RenderTargetResource->ReadPixels(OutImageData, ReadSurfaceDataFlags);
OutWidth = RenderTarget->SizeX;
OutHeight = RenderTarget->SizeY;
}


TArray<uint8> NpySerialization(TArray<FFloat16Color> ImageData, int32 Width, int32 Height, int32 Channel)
void UGTCaptureComponent::CaptureFloat16Image(const FString& Mode, TArray<FFloat16Color>& OutImageData, int32& OutWidth, int32& OutHeight)
{
float *TypePointer = nullptr; // Only used for determing the type

std::vector<int> Shape;
Shape.push_back(Height);
Shape.push_back(Width);
if (Channel != 1) Shape.push_back(Channel);

std::vector<char> NpyHeader = cnpy::create_npy_header(TypePointer, Shape);

// Append the actual data
// FIXME: A slow implementation to convert TArray<FFloat16Color> to binary.
/* A small size test
std::vector<char> NpyData;
for (int i = 0; i < 3 * 3 * 3; i++)
{
NpyData.push_back(i);
}
*/
// std::vector<char> NpyData;
std::vector<float> FloatData;
float DebugMin = 10e10, DebugMax = 0.0;

for (int i = 0; i < ImageData.Num(); i++)
{
if (Channel == 1)
{
float v = ImageData[i].R;
FloatData.push_back(ImageData[i].R);
// debug: Check the range of data
if (v < DebugMin) DebugMin = v;
if (v > DebugMax) DebugMax = v;
}
if (Channel == 3)
{
FloatData.push_back(ImageData[i].R);
FloatData.push_back(ImageData[i].G);
FloatData.push_back(ImageData[i].B); // TODO: Is this a correct order in numpy?
}
}
check(FloatData.size() == Width * Height * Channel);
// Convert to binary array
const unsigned char* bytes = reinterpret_cast<const unsigned char*>(&FloatData[0]);

// https://stackoverflow.com/questions/22629728/what-is-the-difference-between-char-and-unsigned-char
// https://stackoverflow.com/questions/11022099/convert-float-vector-to-byte-vector-and-back
std::vector<unsigned char> NpyData(bytes, bytes + sizeof(float) * FloatData.size());

NpyHeader.insert(NpyHeader.end(), NpyData.begin(), NpyData.end());

// FIXME: Find a more efficient implementation
TArray<uint8> BinaryData;
for (char Element : NpyHeader)
{
BinaryData.Add(Element);
}
return BinaryData;
}

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

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

// Attach this to something, for example, a real camera
const FRotator PawnViewRotation = Pawn->GetViewRotation();
if (!PawnViewRotation.Equals(CaptureComponent->GetComponentRotation()))
{
CaptureComponent->SetWorldRotation(PawnViewRotation);
UE_LOG(LogUnrealCV, Warning, TEXT("Can not find a camera to capture %s"), *Mode);
OutWidth = 0; OutHeight = 0;
return;
}

UTextureRenderTarget2D* RenderTarget = CaptureComponent->TextureTarget;
int32 Width = RenderTarget->SizeX, Height = RenderTarget->SizeY;
TArray<FFloat16Color> ImageData;

FTextureRenderTargetResource* RenderTargetResource;
ImageData.AddZeroed(Width * Height);
OutImageData.Empty();
OutImageData.AddZeroed(RenderTarget->SizeX * RenderTarget->SizeY);
RenderTargetResource = RenderTarget->GameThread_GetRenderTargetResource();
RenderTargetResource->ReadFloat16Pixels(ImageData);

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

return NpyData;
RenderTargetResource->ReadFloat16Pixels(OutImageData);
OutWidth = RenderTarget->SizeX;
OutHeight = RenderTarget->SizeY;
return;
}


void UGTCaptureComponent::TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction)
// void UGTCaptureComponent::Tick(float DeltaTime) // This tick function should be called by the scene instead of been
{
Expand Down
3 changes: 2 additions & 1 deletion Source/UnrealCV/Private/libs/cnpy.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@
// Released under MIT License
// Simplied by Weichao Qiu (qiuwch@gmail.com) from https://github.com/rogersce/cnpy
#include "UnrealCVPrivate.h"
#include "cnpy.h"
#include <vector>
#include<complex>
#include <complex>

using byte = unsigned char;

Expand Down
9 changes: 6 additions & 3 deletions Source/UnrealCV/Public/GTCaptureComponent.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,15 @@ class UNREALCV_API UGTCaptureComponent : public USceneComponent // , public FTic

/** Save image to a file */
FAsyncRecord* Capture(FString Mode, FString Filename);

/** Read binary data in png format */
TArray<uint8> CapturePng(FString Mode);
// TArray<uint8> CapturePng(FString Mode);

/** Read binary data in uncompressed numpy array */
TArray<uint8> CaptureNpy(FString Mode);
// TArray<uint8> CaptureNpy(FString Mode);

void CaptureImage(const FString& Mode, TArray<FColor>& OutImageData, int32& OutWidth, int32& OutHeight);
void CaptureFloat16Image(const FString& Mode, TArray<FFloat16Color>& OutImageData, int32& OutWidth, int32& OutHeight);
private:
const bool bIsTicking = true;

Expand Down

0 comments on commit b50e6be

Please sign in to comment.