Merge pull request 'Write Image on disk + Transfer' (#58) from Noah_DataTransfer_1.0 into main

Reviewed-on: #58

Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/LuckyDataTransfer.Build.cs

@@ -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 ...	
 			}
 			);

Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckyDataTransferSubsystem.cpp

@@ -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;
@@ -151,3 +190,98 @@ 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;
+}

Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckySensorPawnBase.cpp

@@ -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);
+}
+

Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckyDataTransferSubsystem.h

@@ -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"
@@ -108,6 +108,25 @@ public:
 	bool MakeObservationPayload(const FObservationPayload& Data);
 	//---------------------------------------------------------//
 
-protected:
+
-	//LuckyWriteThread* WriteThread = nullptr;
+	//---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);
+	//-------------------------------------------------------//
+	
 };

Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckySensorPawnBase.h

@@ -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;
+};