TimmHess / UnrealImageCapture

A small tutorial repository on capturing images with semantic annotation from UnrealEngine to disk.
MIT License
224 stars 49 forks source link

Save image from target render every frame. #37

Open shekharsuman3 opened 6 months ago

shekharsuman3 commented 6 months ago

Hello Everyone, I am trying to capture and save image to disc every frame from a camera sensor in unreal engine 5.2. At every tick i try to capture and save image to disc. i tried to use code from this repo. But when i try to save every frame, scene freezes and same image is rendered at every frame. but if i capture image at alternate frame it works. Here are my code. I would really appreciate any help. // PinHoleCineCamera.cpp

include "PinHoleCineCamera.h"

include "Engine.h"

include "Engine/SceneCapture2D.h"

include "Runtime/Engine/Classes/Components/SceneCaptureComponent2D.h"

include "CineCameraComponent.h"

include "Kismet/GameplayStatics.h"

include "ShowFlags.h"

include "RHICommandList.h"

include "ImageWrapper/Public/IImageWrapper.h"

include "ImageWrapper/Public/IImageWrapperModule.h"

include "ImageUtils.h"

include "Modules/ModuleManager.h"

include "Misc/FileHelper.h"

include "Engine/TextureRenderTarget2D.h"

include "Engine/Texture2D.h"

include "Misc/FileHelper.h"

include "IImageWrapper.h"

include "HAL/PlatformProcess.h"

include "RenderingThread.h"

include "ImageUtils.h"

include "Engine/World.h"

include "Engine/GameViewportClient.h"

include "IImageWrapperModule.h"

APinHoleCineCamera::APinHoleCineCamera(const FObjectInitializer& ObjectInitializer) : ACineCameraActor(ObjectInitializer) { //RootComponent = CreateDefaultSubobject(TEXT("RootComponent")); UCineCameraComponent CineCameraComponent1 = GetCineCameraComponent(); SceneCaptureComponent = CreateDefaultSubobject(TEXT("SceneCaptureComponent")); SceneCaptureComponent->AttachToComponent(CineCameraComponent1, FAttachmentTransformRules::KeepRelativeTransform); // Set the SceneCaptureComponent properties SceneCaptureComponent->ProjectionType = ECameraProjectionMode::Perspective; //SceneCaptureComponent->OrthoWidth = 0.0f; // Set OrthoWidth to 0 for perspective projection SceneCaptureComponent->FOVAngle = 93.665; // Set Camera Properties SceneCaptureComponent->CaptureSource = ESceneCaptureSource::SCS_FinalColorHDR; SceneCaptureComponent->ShowFlags.SetTemporalAA(true); } void APinHoleCineCamera::AttachPinHoleCameraToVehicle(AActor Vehicle, PINHOLECINECAMERACfg SensorCfg, float FixedDeltaTime) {

SensorCfg = SensorCfg_;
FixedDeltaTime = FixedDeltaTime_;
if (Vehicle) {
    SetActorRelativeLocation(SensorCfg.TranslationSensorToVehicle);
    SetActorRelativeRotation(SensorCfg.RotationSensorToVehicle);
    if (AttachToActor(Vehicle, FAttachmentTransformRules::KeepRelativeTransform)) {
        UE_LOG(LogTemp, Log, TEXT("PinHole: Sensor attached to EGO vehicle."));
    }
    else {
        UE_LOG(LogTemp, Error, TEXT("PinHole: Error during sensor attachment to EGO vehicle."));
    }
}
else {
    UE_LOG(LogTemp, Error, TEXT("PinHole: No parent vehicle found."));
}
// Make sure that output folder exists
if (!FPaths::DirectoryExists(SensorCfg.PathSave)) {
    FPlatformFileManager::Get().GetPlatformFile().CreateDirectoryTree(*SensorCfg.PathSave);
}  
UE_LOG(LogTemp, Log, TEXT("[PinHoleCineCamera] Pin Hole camera settings applied"));

}

void APinHoleCineCamera::GetPinHoleCineCameraData() { const uint32_t Timestamp = std::round(FixedDeltaTime Frame 1000); Frame = Frame + 1U; // Read pixels once RenderFence is completed if (!RenderRequestQueue.IsEmpty()) { // Peek the next RenderRequest from queue FRenderRequestStruct* nextRenderRequest = nullptr; RenderRequestQueue.Peek(nextRenderRequest); if (nextRenderRequest) { //nullptr check if (nextRenderRequest->RenderFence.IsFenceComplete()) { // Check if rendering is done, indicated by RenderFence // Load the image wrapper module IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked(FName("ImageWrapper")); FString FrameNumberString = FString::Printf(TEXT("%06d.png"), Timestamp); FString FilePath = FPaths::Combine(SensorCfg.PathSave, FrameNumberString); static TSharedPtr imageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG); //EImageFormat::PNG //EImageFormat::JPEG imageWrapper->SetRaw(nextRenderRequest->Image.GetData(), nextRenderRequest->Image.GetAllocatedSize(), FrameWidth, FrameHeight, ERGBFormat::BGRA, 8); const TArray64& ImgData = imageWrapper->GetCompressed(5); //const TArray& ImgData = static_cast<TArray<uint8, FDefaultAllocator>> (imageWrapper->GetCompressed(5)); RunAsyncImageSaveTask(ImgData, FilePath); }
} } }

void APinHoleCineCamera::CaptureNonBlocking() { if (!IsValid(SceneCaptureComponent)) { UE_LOG(LogTemp, Error, TEXT("CaptureColorNonBlocking: CaptureComponent was not valid!")); return; } SceneCaptureComponent->TextureTarget->TargetGamma = GEngine->GetDisplayGamma(); FTextureRenderTargetResource* renderTargetResource = SceneCaptureComponent->TextureTarget->GameThread_GetRenderTargetResource();

UE_LOG(LogTemp, Warning, TEXT("Got display gamma"));
struct FReadSurfaceContext { FRenderTarget* SrcRenderTarget;
                             TArray<FColor>* OutData;
                             FIntRect Rect;
                             FReadSurfaceDataFlags Flags;};

UE_LOG(LogTemp, Warning, TEXT("Inited ReadSurfaceContext"));
// Init new RenderRequest
FRenderRequestStruct* renderRequest = new FRenderRequestStruct();
UE_LOG(LogTemp, Warning, TEXT("inited renderrequest"));

// Setup GPU command
FReadSurfaceContext readSurfaceContext = {
    renderTargetResource,
    &(renderRequest->Image),
    FIntRect(0,0,renderTargetResource->GetSizeXY().X, renderTargetResource->GetSizeXY().Y),
    FReadSurfaceDataFlags(RCM_UNorm, CubeFace_MAX)
};
UE_LOG(LogTemp, Warning, TEXT("GPU Command complete"));
// Above 4.22 use this
ENQUEUE_RENDER_COMMAND(SceneDrawCompletion)(
    [readSurfaceContext](FRHICommandListImmediate& RHICmdList) {
        RHICmdList.ReadSurfaceData(
            readSurfaceContext.SrcRenderTarget->GetRenderTargetTexture(),
            readSurfaceContext.Rect,
            *readSurfaceContext.OutData,
            readSurfaceContext.Flags
        );
    });

// Notifiy new task in RenderQueue
RenderRequestQueue.Enqueue(renderRequest);
// Set RenderCommandFence
renderRequest->RenderFence.BeginFence();
renderRequest->RenderFence.Wait();
if (renderRequest->RenderFence.IsFenceComplete()) {
    UE_LOG(LogTemp, Warning, TEXT("fencing complete"));
    GetPinHoleCineCameraData();
    // Delete the first element from RenderQueue
    RenderRequestQueue.Pop();
    delete renderRequest;
}
else {
    UE_LOG(LogTemp, Error, TEXT("fencing not complete"));
    FPlatformProcess::Sleep(0.01f);
    CaptureNonBlocking();
}

} void APinHoleCineCamera::RunAsyncImageSaveTask(TArray64 Image, FString ImageName) { (new FAutoDeleteAsyncTask(Image, ImageName))->StartBackgroundTask(); }

AsyncSaveImageToDiskTask::AsyncSaveImageToDiskTask(TArray64 Image, FString ImageName) { ImageCopy = Image; FileName = ImageName; }

AsyncSaveImageToDiskTask::~AsyncSaveImageToDiskTask() { UE_LOG(LogTemp, Warning, TEXT("AsyncTaskDone")); }

void AsyncSaveImageToDiskTask::DoWork() { FFileHelper::SaveArrayToFile(ImageCopy, FileName); //UE_LOG(LogTemp, Error, TEXT("Stored Image: %s"), FileName); } // PinHoleCineCamera.h

pragma once

include "CoreMinimal.h"

include "GameFramework/Actor.h"

include "CineCameraActor.h"

include "Containers/Queue.h"

include "Engine/SceneCapture2D.h"

include "Components/SceneCaptureComponent2D.h"

include "PinHoleCineCamera.generated.h"

struct PINHOLECINECAMERACfg { // Rotation from sensor coordinate system to vehicle coordinate system. FRotator RotationSensorToVehicle; // Translation from sensor coordinate system to vehicle coordinate system. FVector TranslationSensorToVehicle; // Path where to save the csv logger file. FString PathSave; // Sensor frequency float Frequency; };

USTRUCT() struct FRenderRequestStruct { GENERATED_BODY() TArray Image; FRenderCommandFence RenderFence; FRenderRequestStruct() {} };

UCLASS() class AVL_UE5_API APinHoleCineCamera : public ACineCameraActor{ GENERATED_BODY() public: APinHoleCineCamera(const FObjectInitializer& ObjectInitializer); protected: //virtual void BeginPlay() override; virtual void BeginPlay() override { Super::BeginPlay(); };

public: //virtual void Tick(float DeltaTime) override; void AttachPinHoleCameraToVehicle(AActor Vehicle, PINHOLECINECAMERACfg SensorCfg, float FixedDeltaTime); void GetPinHoleCineCameraData(); UFUNCTION(BlueprintCallable, Category = "ImageCapture") void CaptureNonBlocking(); //bool IsPreviousFenceComplete() const; UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Capture") int FrameWidth = 4112; UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Capture") int FrameHeight = 2176;
protected: int ImgCounter = 0; void RunAsyncImageSaveTask(TArray64 Image, FString ImageName); // RenderRequest Queue TQueue<FRenderRequestStruct
> RenderRequestQueue; private: UPROPERTY(EditAnywhere, Category = "Capture") USceneCaptureComponent2D* SceneCaptureComponent; PINHOLECINECAMERACfg SensorCfg; FString TestSequenceNumber; float FixedDeltaTime; int32_t Frame = 0;
};

class AsyncSaveImageToDiskTask : public FNonAbandonableTask { public: AsyncSaveImageToDiskTask(TArray64 Image, FString ImageName); ~AsyncSaveImageToDiskTask(); FORCEINLINE TStatId GetStatId() const { RETURN_QUICK_DECLARE_CYCLE_STAT(AsyncSaveImageToDiskTask, STATGROUP_ThreadPoolAsyncTasks); } protected: TArray64 ImageCopy; FString FileName = ""; public: void DoWork(); };

TimmHess commented 6 months ago

Hey,

What do you mean by freeze? Properly freeze and crash, or just very slow but you still get logs and images?

I only skimmed the code... The thing that looked most suspicious to me is this (below): Depending on the behavior of RenderFence.Wait() it looks like either an infinite loop or perhaps interferes with the game-thread because your game thread waits long times for the render to finish. Have you tried peeking the render queue in the tick function? That should allow you to remove the wait(). It might happen that the queue stacks up and eventually you run out of RAM. In that case you might want to either reduce the capturing framerate or if possible use time dilation to artifically slow your application and maintain higher framerate.

// Notifiy new task in RenderQueue
RenderRequestQueue.Enqueue(renderRequest);
// Set RenderCommandFence
renderRequest->RenderFence.BeginFence();
renderRequest->RenderFence.Wait();
if (renderRequest->RenderFence.IsFenceComplete()) {
    UE_LOG(LogTemp, Warning, TEXT("fencing complete"));
    GetPinHoleCineCameraData();
    // Delete the first element from RenderQueue
    RenderRequestQueue.Pop();
    delete renderRequest;
}
else {
    UE_LOG(LogTemp, Error, TEXT("fencing not complete"));
    FPlatformProcess::Sleep(0.01f);
    CaptureNonBlocking();
}
shekharsuman3 commented 6 months ago

Hey thanks for your quick response. I got your point. i did it to check what is causing the issue. Its working now but it drops my fps to 8 from 80. These are my code. Even the sensor warmup as you can see in the code where i am not doing anything brings my fps to 14. I would really appreciate if you can tell me whats the problem here.

void APinHoleCineCamera::SensorWarmup() {
    CaptureNonBlocking();
    // Read pixels once RenderFence is completed
    if (!RenderRequestQueue.IsEmpty()) {
        // Peek the next RenderRequest from queue
        FRenderRequestStruct* nextRenderRequest = nullptr;
        RenderRequestQueue.Peek(nextRenderRequest);
        if (nextRenderRequest) { //nullptr check
            if (nextRenderRequest->RenderFence.IsFenceComplete()) { // Check if rendering is done, indicated by RenderFence
                // Delete the first element from RenderQueue
                RenderRequestQueue.Pop();
                delete nextRenderRequest;
            } else {
                UE_LOG(LogTemp, Error, TEXT("[WarmUp]: Render fence not complete"));
            }
        }
    }
}

void APinHoleCineCamera::CaptureNonBlocking() {
    if (!IsValid(SceneCaptureComponent)) {
        UE_LOG(LogTemp, Error, TEXT("CaptureColorNonBlocking: CaptureComponent was not valid!"));
        return;
    }
    SceneCaptureComponent->TextureTarget->TargetGamma = GEngine->GetDisplayGamma();
    FTextureRenderTargetResource* renderTargetResource = SceneCaptureComponent->TextureTarget->GameThread_GetRenderTargetResource();

    UE_LOG(LogTemp, Warning, TEXT("Got display gamma"));
    struct FReadSurfaceContext {
        FRenderTarget* SrcRenderTarget;
        TArray<FColor>* OutData;
        FIntRect Rect;
        FReadSurfaceDataFlags Flags;
    };

    UE_LOG(LogTemp, Warning, TEXT("Inited ReadSurfaceContext"));
    // Init new RenderRequest
    FRenderRequestStruct* renderRequest = new FRenderRequestStruct();
    UE_LOG(LogTemp, Warning, TEXT("inited renderrequest"));

    // Setup GPU command
    FReadSurfaceContext readSurfaceContext = {
        renderTargetResource,
        &(renderRequest->Image),
        FIntRect(0,0,renderTargetResource->GetSizeXY().X, renderTargetResource->GetSizeXY().Y),
        FReadSurfaceDataFlags(RCM_UNorm, CubeFace_MAX)
    };
    UE_LOG(LogTemp, Warning, TEXT("GPU Command complete"));
    // Above 4.22 use this
    ENQUEUE_RENDER_COMMAND(SceneDrawCompletion)(
        [readSurfaceContext](FRHICommandListImmediate& RHICmdList) {
            RHICmdList.ReadSurfaceData(
                readSurfaceContext.SrcRenderTarget->GetRenderTargetTexture(),
                readSurfaceContext.Rect,
                *readSurfaceContext.OutData,
                readSurfaceContext.Flags
            );
        });

    // Notifiy new task in RenderQueue
    RenderRequestQueue.Enqueue(renderRequest);
    // Set RenderCommandFence
    renderRequest->RenderFence.BeginFence();
}

void APinHoleCineCamera::SaveNonBlocking() {

    // Read pixels once RenderFence is completed
    if (!RenderRequestQueue.IsEmpty()) {
        // Peek the next RenderRequest from queue
        FRenderRequestStruct* nextRenderRequest = nullptr;
        RenderRequestQueue.Peek(nextRenderRequest);
        if (nextRenderRequest) { //nullptr check
            if (nextRenderRequest->RenderFence.IsFenceComplete()) { // Check if rendering is done, indicated by RenderFence
                const uint32_t Timestamp = std::round(FixedDeltaTime * Frame * 1000);
                Frame = Frame + 1U;
                // Load the image wrapper module 
                UE_LOG(LogTemp, Log, TEXT("render fence complete"));
                IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
                FString FrameNumberString = FString::Printf(TEXT("%06d.png"), Timestamp);
                FString FilePath = FPaths::Combine(SensorCfg.PathSave, FrameNumberString);
                // Create an image wrapper for BMP format
                static TSharedPtr<IImageWrapper> imageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG); //EImageFormat::PNG //EImageFormat::JPEG
                imageWrapper->SetRaw(nextRenderRequest->Image.GetData(), nextRenderRequest->Image.GetAllocatedSize(), FrameWidth, FrameHeight, ERGBFormat::BGRA, 8);
                TArray64<uint8> ImgData;
                bool success = imageWrapper->GetRaw(ERGBFormat::BGRA, 8, ImgData);
                if (success) {
                    RunAsyncImageSaveTask(ImgData, FilePath);
                } else {
                    UE_LOG(LogTemp, Error, TEXT("Error getting raw image from ImageWrapper"));
                }              
                // Delete the first element from RenderQueue
                RenderRequestQueue.Pop();
                delete nextRenderRequest;
            }
            else {
                UE_LOG(LogTemp, Error, TEXT("Render fence not complete attempting again.."));
            }
        }
    }
}

void APinHoleCineCamera::RunAsyncImageSaveTask(TArray64<uint8> Image, FString ImageName) {
    (new FAutoDeleteAsyncTask<AsyncSaveImageToDiskTask>(Image, ImageName))->StartBackgroundTask();
}

AsyncSaveImageToDiskTask::AsyncSaveImageToDiskTask(TArray64<uint8> Image, FString ImageName) {
    ImageCopy = Image;
    FileName = ImageName;
}

AsyncSaveImageToDiskTask::~AsyncSaveImageToDiskTask() {
    UE_LOG(LogTemp, Warning, TEXT("AsyncTaskDone"));
}

void AsyncSaveImageToDiskTask::DoWork() {
    // Save compress Image
    TArray<FColor> ImageDataAsColor;
    for (int32 i = 0; i < ImageCopy.Num(); i += 4) // Assuming BGRA format
    {
        FColor PixelColor(ImageCopy[i + 2], ImageCopy[i + 1], ImageCopy[i], 255); // BGRA to RGBA
        ImageDataAsColor.Add(PixelColor);
    }
    TArray<uint8> CompressedImage;
    FImageUtils::CompressImageArray(ImageWidth, ImageHeight, ImageDataAsColor, CompressedImage);
    FFileHelper::SaveArrayToFile(CompressedImage, *FileName);

    // SAVE Raw image 
    // Create an FImage object with the appropriate size
    //FImage Image(ImageWidth, ImageHeight, 1, ERawImageFormat::BGRA8);
    //// Copy your image data to the FImage object
    //FMemory::Memcpy(Image.RawData.GetData(), ImageCopy.GetData(), ImageCopy.Num());
    //// Save the image
    //bool bSuccess = FImageUtils::SaveImageAutoFormat(*FileName, Image);
    //UE_LOG(LogTemp, Error, TEXT("Stored Image: %s"), *FileName);
}
TimmHess commented 6 months ago

Previously, PR #20 pointed out that ReadSufaceData is actually not async and one should use FRHIGPUTextureReadback instead. I accepted that PR, and wonder why it's apparently not in the code.

Anyhow.. that could be the issue. Even when not writing to disk you already fetch the data into RAM. Depending on your resolution that can be quite a bit of data.

shekharsuman3 commented 3 months ago

Hi Timm, Thanks for your response. Now i am using FRHIGPUTextureReadback but it seems it works only with standard resolution. But when i try to use it for 41122176 resolution the pixels in generated image are misalligned. Do you also get this error. Can you please try 41122176 resolution. but it works for 51202880 and 38402160. it also works for 8k. Thanks.

MarvinSt commented 2 months ago

@shekharsuman3 The alignment is perhaps due to the fact that the bytes are packed in a certain way? I am not sure about the details, but in UE5.3 the function returns an additional parameter called OutRowPitchInPixels, which are essentially the amount of pixels per row stored in the databuffer: void *RawData = RenderRequest->Readback.Lock(OutRowPitchInPixels); I'm trying to achieve a similar thing as you and noticed this variable. I have not confirmed if it can deviate from the image width.