UNSTBL - Refacto

+ Exit CaptureImage faster in case of problem
+ Register cameras manually in Episode SubSystem
+ Make FilePath optional to keep the correct default value in saved games
- Image capture doesn't work, image is empty

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

@@ -204,90 +204,95 @@ void ULuckyDataTransferSubsystem::RegisterSensor(ALuckySensorPawnBase* Sensor)
 	}
 }
 
-bool ULuckyDataTransferSubsystem::WriteImageToDisk(const FString& inPath, const double Timestamp, const FString& Comment)
+bool ULuckyDataTransferSubsystem::WriteImageToDisk(const double Timestamp, const FString& InPath, const FString& Comment)
 {
 	if (SessionID.IsEmpty())
 	{
 		UE_LOG(LogTemp, Warning, TEXT("The Capture Session ID is empty"));
 		return false;
 	}
-
+	
 	// TODO JB: I have this line in my project because I had a discrepancy in editor vs packaged
 	// TODO UKismetSystemLibrary::GetProjectSavedDirectory()
 	// TODO Check which one is working in shipping build
-	FString Path = inPath.IsEmpty() ? FPaths::ProjectSavedDir() : inPath; // swap this for const path
+	FString Path = InPath.IsEmpty() ? FPaths::ProjectSavedDir() : InPath; // swap this for const path
-	
+
-	if (!SensorPawns.IsEmpty())
+	if (SensorPawns.IsEmpty())
 	{
-		for (const ALuckySensorPawnBase* Sensor : SensorPawns)
+		UE_LOG(LogTemp, Warning, TEXT("Attempted to capture image data but no sensors are registered"));
+		return false;
+	}
+	
+	for (const ALuckySensorPawnBase* Sensor : SensorPawns)
+	{
+		if (IsValid(Sensor) && Sensor->SensorInfo.bActive && Sensor->GetCamera() && Sensor->GetCaptureComponent())
 		{
-			if (IsValid(Sensor) && Sensor->SensorInfo.bActive && Sensor->GetCamera() && Sensor->GetCaptureComponent())
+			// TODO Pass the robot name in params - use RobotType and EnumToString
+			FString Robot = TEXT("Robot_Name");
+
+			FString Episode = SessionID;
+
+			ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList)
 			{
-				FString Robot = TEXT("Robot_Name");
+				FTextureResource* Resource = Sensor->RenderTarget->GetResource();
+				FRHITexture* ResourceRHI = Resource->GetTexture2DRHI();
 
-				FString Episode = SessionID;
+				TArray<FColor> OutPixels;
 
-				ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList)
+				// Fail and exit 
+				if (!ensure(ResourceRHI))
 				{
-					FTextureResource* Resource = Sensor->RenderTarget->GetResource();
+					UE_LOG(LogTemp, Warning, TEXT("ResourceRHI not found"));
-					FRHITexture* ResourceRHI = Resource->GetTexture2DRHI();
+					return;
+				}
+				
+				OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY());
+				RHICmdList.ReadSurfaceData(
+					ResourceRHI,
+					FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()),
+					OutPixels,
+					FReadSurfaceDataFlags()
+				);
 
-					TArray<FColor> OutPixels;
+				UE_LOG(LogTemp, Warning, TEXT("Logged pixels: %d"), OutPixels.Num())
 
-					if (ensure(ResourceRHI))
+				FImageWriteTask* ImageTask = new FImageWriteTask();
-					{
-						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())
+				// TODO What's that for?
+				if (Path.Right(1) == "/")
+				{
+					//Path = Path.Left(Path.Len() - 1);
+				}
 
-						FImageWriteTask* ImageTask = new FImageWriteTask();
+				const FString RobotName = TEXT("Robot_Name");
-						
+				
-						if (Path.Right(1) == "/")
+				const FString Filename = FString::Printf(
-						{
+				TEXT("%s/TrainingData/%s/%s/%s/%s_%s_%s_%d"),
-							//Path = Path.Left(Path.Len() - 1);
+					*Path,
-						}
+					*RobotName,
+					*Episode,
+					*Sensor->SensorInfo.Name,
+					*RobotName,
+					*Sensor->SensorInfo.Name,
+					*Comment,
+					FMath::Floor<int32>(Timestamp * 1000)
+				);
+				
+				UE_LOG(LogTemp, Warning, TEXT("FileName %s"), *Filename);
 
-						FString Robot = TEXT("Robot_Name");
+				ImageTask->Format = EImageFormat::PNG;
-						
+				ImageTask->Filename = Filename;
-						const FString Filename = FString::Printf(
+				ImageTask->PixelData = MakeUnique<TImagePixelData<FColor>>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64<FColor>(OutPixels));
-						TEXT("LuckRobotData/%s/%s/%s/%s/%s_%s_%s_%d"),
+				ImageTask->bOverwriteFile = true;
-							*Path,
+				ImageTask->CompressionQuality = static_cast<int32>(EImageCompressionQuality::Default);
-							*Robot,
+				
-							*Episode,
+				// Add to write queue (async)
-							*Sensor->SensorInfo.Name,
+				FModuleManager::LoadModuleChecked<IImageWriteQueueModule>("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr<IImageWriteTaskBase>(ImageTask));
-							*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->OnCompleted = [](bool bSuccess) {
-						ImageTask->Filename = Filename;
+					UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed"));
-						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;
-				});
-			}
 		}
 	}
 	

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

@@ -51,9 +51,9 @@ public:
 
 	ULuckyDataTransferSubsystem();
 
-	virtual void Initialize(FSubsystemCollectionBase& Collection);
+	virtual void Initialize(FSubsystemCollectionBase& Collection) override;
 	
-	virtual void Deinitialize();
+	virtual void Deinitialize() override;
 
 	TSharedPtr<IWebSocket> Socket;
 
@@ -126,7 +126,7 @@ public:
 	void RegisterSensor(ALuckySensorPawnBase* Sensor);
 
 	UFUNCTION(BlueprintCallable, Meta = (AutoCreateRefTerm = "Path, Comment"), Category = "Capture")
-	bool WriteImageToDisk(const FString& Path, const double Timestamp, const FString& Comment = "");
+	bool WriteImageToDisk(const double Timestamp, const FString& InPath = "", const FString& Comment = "");
 	//-------------------------------------------------------//
 	
 };

Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp

@@ -184,6 +184,21 @@ void UEpisodeSubSystem::FindRobotPawnFromScene()
 	}
 }
 
+void UEpisodeSubSystem::InitCameras()
+{
+	// TODO Fix the spawning of sensors in Cpp and spawn them using a config?
+	// TODO How people can move the camera themselves?
+
+	// Find all sensors in the scene
+	TArray<AActor*> Sensors;
+	UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors);
+	
+	for (const auto Sensor : Sensors)
+	{
+		if (const auto Camera = Cast<ALuckySensorPawnBase>(Sensor)) Cameras.Add(Camera);
+	}
+}
+
 void UEpisodeSubSystem::ConfigureDataCapture()
 {
 	if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
@@ -192,6 +207,13 @@ void UEpisodeSubSystem::ConfigureDataCapture()
 		//Connect to websocket and create session id
 		DataTransfer->ConnectToWebsocket("ws://127.0.0.1:3000", "");
 		DataTransfer->CreateCaptureSessionID();
+
+		InitCameras();
+		for (const auto& Cam : Cameras)
+		{
+			DataTransfer->RegisterSensor(Cam.Get());
+			Cam->SensorInfo.bActive = true;
+		}
 	}
 }
 
@@ -211,17 +233,23 @@ FObservationPayload UEpisodeSubSystem::CreatePayload()
 		// enter a message here - FString,
 		// TMap of FString (Actuator name or index), and Float (value of actuator)
 		// Camera info struct goes here, don't worry about this for now, just use TArray<FObservationCameraObject>()
-		// What about episode success?
+		// What about episode success? -> can be stated after the result is known
 		// How to invalidate data
+
+		// Anuj -> How many frames do we need to store in a single parquet chunk
+		// Exact data structure with correct data types
 	};
 }
 
 void UEpisodeSubSystem::SendEpisodeData(const FObservationPayload& Payload) const
 {
+	// PayloadBuffer.Add(Payload)
+	// Every X frames -> Write parquet chunk
+	
 	if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
 	{
 		// Here generate the path for each image?
-		// DataTransfer->WriteImageToDisk(BaseImageDataPath, 0.f);
+		DataTransfer->WriteImageToDisk(CurrentRobot->PhysicsSceneProxy->GetMujocoData().time);
 
 		// Don't send data if socket is disconnected
 		if (!DataTransfer->Socket->IsConnected()) return;

Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp

@@ -54,11 +54,6 @@ FTransform URobotPilotSO100Component::GetReachableTransform()
 	RewardAxis.Z = 0; // Nullify Z to keep a 2D vector -> ensure the geometry roll/pitch are 0
 	const FRotator TowardPivotRotation = UKismetMathLibrary::MakeRotFromXZ(RewardAxis, FVector::UpVector);
 
-	// Debug
-	// DrawDebugLine(this->GetWorld(), ArmPivotLocation + ArmWorldRotation.GetForwardVector() * 70, ArmPivotLocation, FColor::Green, true);
-	// DrawDebugLine(this->GetWorld(), ArmPivotLocation + FVector::UpVector * 70, ArmPivotLocation, FColor::Red, true);
-	// DrawDebugLine(this->GetWorld(), RandomLocation, RandomLocation + TowardPivotRotation.Quaternion().GetForwardVector() * -50 , FColor::Blue, true);
-	
 	// Return the Object Transform
 	return FTransform(TowardPivotRotation, RandomLocation);
 }

Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp

@@ -20,7 +20,6 @@ void ARobotPawn::BeginPlay()
 void ARobotPawn::InitRobot()
 {
 	InitPilotComponent();
-	InitCamera();
 }
 
 void ARobotPawn::InitPilotComponent()
@@ -70,17 +69,3 @@ void ARobotPawn::InitPilotComponent()
 	}
 }
 
-void ARobotPawn::InitCamera()
-{
-	// TODO Fix the spawning of sensors in Cpp and spawn them using a config?
-	// TODO How people can move the camera themselves?
-
-	// Find all sensors in the scene
-	TArray<AActor*> Sensors;
-	UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors);
-	
-	for (const auto Sensor : Sensors)
-	{
-		if (const auto Camera = Cast<ALuckySensorPawnBase>(Sensor)) Cameras.Add(Camera);
-	}
-}

Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h

@@ -6,6 +6,7 @@
 #include "EpisodeSubSystem.generated.h"
 
 
+class ALuckySensorPawnBase;
 class ATextRenderActor;
 class AMujocoStaticMeshActor;
 class ARobotPawn;
@@ -85,6 +86,11 @@ private:
 	UPROPERTY()
 	TObjectPtr<ARobotPawn> CurrentRobot;
 
+	// ---------------------
+	// ------ SENSORS ------
+	// ---------------------
+	void InitCameras();
+	TArray<TObjectPtr<ALuckySensorPawnBase>> Cameras;
 
 
 

Source/LuckyWorldV2/Public/Robot/RobotPawn.h

@@ -40,11 +40,4 @@ public:
 	URobotPilotComponent* RobotPilotComponent = nullptr;
 	UFUNCTION(BlueprintCallable)
 	void InitPilotComponent(); // This should have Robot type as parameter?
-
-	// ---------------------
-	// ------ SENSORS ------
-	// ---------------------
-	void InitCamera();
-	UPROPERTY(EditAnywhere, BlueprintReadWrite)
-	TArray<TObjectPtr<ALuckySensorPawnBase>> Cameras;
 };