Created image write on async thread

This commit is contained in:
Noah 2025-05-01 13:18:42 -04:00 committed by Jean-Baptiste Briant (JB)
parent 6bef120e16
commit b3e0389675
9 changed files with 295 additions and 5 deletions

Binary file not shown.

View File

@ -29,7 +29,11 @@ public class LuckyDataTransfer : ModuleRules
"WebSockets",
"Json",
"JsonUtilities",
"ImageWriteQueue"
"ImageWriteQueue",
"FunctionalTesting",
"RenderCore",
"RHI",
"RHICore"
// ... add other public dependencies that you statically link with here ...
}
);
@ -41,7 +45,8 @@ public class LuckyDataTransfer : ModuleRules
"CoreUObject",
"Engine",
"Slate",
"SlateCore", "ImageWriteQueue",
"SlateCore",
"ImageWriteQueue"
// ... add private dependencies that you statically link with here ...
}
);

View File

@ -2,18 +2,45 @@
#include "LuckyDataTransferSubsystem.h"
#include "AutomationBlueprintFunctionLibrary.h"
#include "ImageUtils.h"
#include "RenderingThread.h"
#include "RenderUtils.h"
#include "RenderGraphUtils.h"
#include "RHI.h"
#include "RHICommandList.h"
#include "ImageWriteQueue.h"
#include "ImageWriteTask.h"
#include "ImagePixelData.h"
#include "JsonUtilities.h"
#include "JsonObjectConverter.h"
#include "ReviewComments.h"
#include "WebSocketsModule.h"
#include "Kismet/KismetStringLibrary.h"
#include "Camera/CameraActor.h"
#include "Camera/CameraComponent.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Kismet/GameplayStatics.h"
#include "Kismet/KismetMathLibrary.h"
#include "Kismet/KismetRenderingLibrary.h"
#include "Slate/SceneViewport.h"
#include "Virtualization/VirtualizationTypes.h"
ULuckyDataTransferSubsystem::ULuckyDataTransferSubsystem()
{
}
void ULuckyDataTransferSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
Super::Initialize(Collection);
if (UGameInstance* GI = GetWorld()->GetGameInstance())
{
}
}
void ULuckyDataTransferSubsystem::Deinitialize()
@ -123,6 +150,12 @@ FPayload ULuckyDataTransferSubsystem::InterpretData(const FString& Message)
return Payload;
}
void IncomingMessage(const FString& Message)
{
const FString outMessage = Message.IsEmpty() ? TEXT("no valid message") : Message;
UE_LOG(LogTemp, Warning, TEXT("Incoming message: %s"), *outMessage);
}
void ULuckyDataTransferSubsystem::CommandReady(const FPayload& Payload)
{
if (OnCommandReady.IsBound())
@ -137,6 +170,12 @@ bool ULuckyDataTransferSubsystem::MakeObservationPayload(const FObservationPaylo
return CreateJsonPayload_Observation(Data);
}
FString ULuckyDataTransferSubsystem::CreateCaptureSessionID()
{
SessionID = FGuid::NewGuid().ToString().Left(5);
return SessionID;
}
bool ULuckyDataTransferSubsystem::CreateJsonPayload_Observation(const FObservationPayload& Data)
{
bool bSuccess = false;
@ -150,4 +189,99 @@ bool ULuckyDataTransferSubsystem::CreateJsonPayload_Observation(const FObservati
}
return bSuccess;
}
void ULuckyDataTransferSubsystem::RegisterSensor(ALuckySensorPawnBase* Sensor)
{
if (IsValid(Sensor))
{
SensorPawns.Add(Sensor);
}
}
bool ULuckyDataTransferSubsystem::WriteImageToDisk(const FString& inPath, const double Timestamp, const FString& Comment)
{
if (SessionID.IsEmpty())
{
UE_LOG(LogTemp, Warning, TEXT("The Capture Session ID is empty"));
return false;
}
FString Path = inPath.IsEmpty() ? TEXT("C:/LuckyRobotsImages") : inPath;
if (!SensorPawns.IsEmpty())
{
for (const ALuckySensorPawnBase* Sensor : SensorPawns)
{
if (IsValid(Sensor) && Sensor->SensorInfo.bActive && Sensor->GetCamera() && Sensor->GetCaptureComponent())
{
FString Robot = TEXT("Robot_Name");
FString Episode = SessionID;
ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList)
{
FTextureResource* Resource = Sensor->RenderTarget->GetResource();
FRHITexture* ResourceRHI = Resource->GetTexture2DRHI();
TArray<FColor> OutPixels;
if (ensure(ResourceRHI))
{
OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY());
RHICmdList.ReadSurfaceData(
ResourceRHI,
FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()),
OutPixels,
FReadSurfaceDataFlags()
);
UE_LOG(LogTemp, Warning, TEXT("Logged pixels: %d"), OutPixels.Num())
FImageWriteTask* ImageTask = new FImageWriteTask();
if (Path.Right(1) == "/")
{
//Path = Path.Left(Path.Len() - 1);
}
FString Robot = TEXT("Robot_Name");
const FString Filename = FString::Printf(
TEXT("%s/%s/%s/%s/%s_%s_%s_%d"),
*Path,
*Robot,
*Episode,
*Sensor->SensorInfo.Name,
*Robot,
*Sensor->SensorInfo.Name,
*Comment,
FMath::Floor<int32>(Timestamp * 1000)
);
//UE_LOG(LogTemp, Warning, TEXT("Evan requested a longer string describing the inner workings of the following string which describes in great detail the file path for the image you've just written to disk. It is: %s"), *Filename);
ImageTask->Format = EImageFormat::PNG;
ImageTask->Filename = Filename;
ImageTask->PixelData = MakeUnique<TImagePixelData<FColor>>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64<FColor>(OutPixels));
ImageTask->bOverwriteFile = true;
ImageTask->CompressionQuality = (int32)EImageCompressionQuality::Default;
//Add to write queue (async)
FModuleManager::LoadModuleChecked<IImageWriteQueueModule>("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr<IImageWriteTaskBase>(ImageTask));
ImageTask->OnCompleted = [](bool bSuccess) {
UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed"));
};
return true;
}
return false;
});
}
}
}
return false;
}

View File

@ -0,0 +1,57 @@
// Fill out your copyright notice in the Description page of Project Settings.
#include "LuckySensorPawnBase.h"
#include "Camera/CameraComponent.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Engine/TextureRenderTarget2D.h"
#include "GameFramework/SpringArmComponent.h"
// Sets default values
ALuckySensorPawnBase::ALuckySensorPawnBase()
{
// Set this pawn to call Tick() every frame. You can turn this off to improve performance if you don't need it.
PrimaryActorTick.bCanEverTick = true;
SpringArm = CreateDefaultSubobject<USpringArmComponent>(TEXT("SpringArm"));
SpringArm->bDoCollisionTest = false;
RootComponent = SpringArm;
Camera = CreateDefaultSubobject<UCameraComponent>(TEXT("Camera"));
Camera->SetupAttachment(SpringArm);
Camera->SetRelativeTransform(FTransform());
CaptureComponent = CreateDefaultSubobject<USceneCaptureComponent2D>(TEXT("CameraComponent"));
CaptureComponent->SetupAttachment(RootComponent);
RenderTarget = CreateDefaultSubobject<UTextureRenderTarget2D>(TEXT("RenderTarget"));
RenderTarget->UpdateResourceImmediate(true);
RenderTarget->ResizeTarget(1, 1);
CaptureComponent->CaptureSource = SCS_FinalColorLDR;
}
// Called when the game starts or when spawned
void ALuckySensorPawnBase::BeginPlay()
{
CaptureComponent->TextureTarget = RenderTarget;
Super::BeginPlay();
}
// Called every frame
void ALuckySensorPawnBase::Tick(float DeltaTime)
{
Super::Tick(DeltaTime);
}
// Called to bind functionality to input
void ALuckySensorPawnBase::SetupPlayerInputComponent(UInputComponent* PlayerInputComponent)
{
Super::SetupPlayerInputComponent(PlayerInputComponent);
}

View File

@ -4,7 +4,7 @@
#include "CoreMinimal.h"
#include "IWebSocket.h"
//#include "LuckyWriteThread.h"
#include "LuckySensorPawnBase.h"
#include "ObservationData.h"
#include "Subsystems/WorldSubsystem.h"
#include "LuckyDataTransferSubsystem.generated.h"
@ -107,7 +107,26 @@ public:
UFUNCTION(BlueprintCallable, Category = "Websocket")
bool MakeObservationPayload(const FObservationPayload& Data);
//---------------------------------------------------------//
//---Image Capture----------------------------------------//
UFUNCTION(BlueprintCallable, Category = "Websocket")
FString CreateCaptureSessionID();
UPROPERTY(BlueprintReadOnly, Category = "Websocket")
FString SessionID;
UPROPERTY(BlueprintReadWrite, Category = "Capture")
TMap<FString, ACameraActor*> ActiveCameras;
UPROPERTY(BlueprintReadWrite, Category = "Capture")
TArray<ALuckySensorPawnBase*> SensorPawns;
UFUNCTION(BlueprintCallable, Category = "Capture")
void RegisterSensor(ALuckySensorPawnBase* Sensor);
UFUNCTION(BlueprintCallable, Meta = (AutoCreateRefTerm = "Path, Comment"), Category = "Capture")
bool WriteImageToDisk(const FString& Path, const double Timestamp, const FString& Comment);
//-------------------------------------------------------//
protected:
//LuckyWriteThread* WriteThread = nullptr;
};

View File

@ -0,0 +1,75 @@
// Fill out your copyright notice in the Description page of Project Settings.
#pragma once
#include "CoreMinimal.h"
#include "Camera/CameraActor.h"
#include "GameFramework/Pawn.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Engine/TextureRenderTarget2D.h"
#include "GameFramework/SpringArmComponent.h"
#include "LuckySensorPawnBase.generated.h"
USTRUCT(BlueprintType)
struct FSensorInfo
{
GENERATED_USTRUCT_BODY()
public:
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
FString Name = FString();
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
FVector Anchor = FVector::ZeroVector;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
float Distance = 0.f;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
FRotator Rotation = FRotator::ZeroRotator;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
bool bActive = false;
};
UCLASS()
class LUCKYDATATRANSFER_API ALuckySensorPawnBase : public APawn
{
GENERATED_BODY()
public:
// Sets default values for this pawn's properties
ALuckySensorPawnBase();
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
FSensorInfo SensorInfo = FSensorInfo();
UFUNCTION(BlueprintCallable, Category="Sensor")
UCameraComponent* GetCamera() const { return Camera; }
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
UTextureRenderTarget2D* RenderTarget = nullptr;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
USceneCaptureComponent2D* CaptureComponent = nullptr;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
USpringArmComponent* SpringArm = nullptr;
UFUNCTION(BlueprintCallable, Category="Sensor")
USceneCaptureComponent2D* GetCaptureComponent() const { return CaptureComponent; }
protected:
// Called when the game starts or when spawned
virtual void BeginPlay() override;
UPROPERTY(BlueprintReadWrite, EditAnywhere, Category="Sensor")
UCameraComponent* Camera = nullptr;
public:
// Called every frame
virtual void Tick(float DeltaTime) override;
// Called to bind functionality to input
virtual void SetupPlayerInputComponent(class UInputComponent* PlayerInputComponent) override;
};