diff --git a/Content/Levels/kitchenLevel/BP-Sensor.uasset b/Content/Levels/kitchenLevel/BP-Sensor.uasset new file mode 100644 index 00000000..bb01a116 Binary files /dev/null and b/Content/Levels/kitchenLevel/BP-Sensor.uasset differ diff --git a/Content/Levels/kitchenLevel/kitchenLevel1.umap b/Content/Levels/kitchenLevel/kitchenLevel1.umap index 417621ea..62a4b7f5 100644 Binary files a/Content/Levels/kitchenLevel/kitchenLevel1.umap and b/Content/Levels/kitchenLevel/kitchenLevel1.umap differ diff --git a/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckyDataTransferSubsystem.cpp b/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckyDataTransferSubsystem.cpp index c0a62a36..63e32093 100644 --- a/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckyDataTransferSubsystem.cpp +++ b/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 - - if (!SensorPawns.IsEmpty()) + FString Path = InPath.IsEmpty() ? FPaths::ProjectSavedDir() : InPath; // swap this for const path + + 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 OutPixels; - ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList) + // Fail and exit + if (!ensure(ResourceRHI)) { - FTextureResource* Resource = Sensor->RenderTarget->GetResource(); - FRHITexture* ResourceRHI = Resource->GetTexture2DRHI(); + UE_LOG(LogTemp, Warning, TEXT("ResourceRHI not found")); + return; + } + + OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY()); + RHICmdList.ReadSurfaceData( + ResourceRHI, + FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()), + OutPixels, + FReadSurfaceDataFlags() + ); - TArray OutPixels; + UE_LOG(LogTemp, Warning, TEXT("Logged pixels: %d"), OutPixels.Num()) - if (ensure(ResourceRHI)) - { - OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY()); - RHICmdList.ReadSurfaceData( - ResourceRHI, - FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()), - OutPixels, - FReadSurfaceDataFlags() - ); + FImageWriteTask* ImageTask = new FImageWriteTask(); - 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(); - - if (Path.Right(1) == "/") - { - //Path = Path.Left(Path.Len() - 1); - } + const FString RobotName = TEXT("Robot_Name"); + + const FString Filename = FString::Printf( + TEXT("%s/TrainingData/%s/%s/%s/%s_%s_%s_%d"), + *Path, + *RobotName, + *Episode, + *Sensor->SensorInfo.Name, + *RobotName, + *Sensor->SensorInfo.Name, + *Comment, + FMath::Floor(Timestamp * 1000) + ); + + UE_LOG(LogTemp, Warning, TEXT("FileName %s"), *Filename); - FString Robot = TEXT("Robot_Name"); - - const FString Filename = FString::Printf( - TEXT("LuckRobotData/%s/%s/%s/%s/%s_%s_%s_%d"), - *Path, - *Robot, - *Episode, - *Sensor->SensorInfo.Name, - *Robot, - *Sensor->SensorInfo.Name, - *Comment, - FMath::Floor(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>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64(OutPixels)); + ImageTask->bOverwriteFile = true; + ImageTask->CompressionQuality = static_cast(EImageCompressionQuality::Default); + + // Add to write queue (async) + FModuleManager::LoadModuleChecked("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr(ImageTask)); - ImageTask->Format = EImageFormat::PNG; - ImageTask->Filename = Filename; - ImageTask->PixelData = MakeUnique>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64(OutPixels)); - ImageTask->bOverwriteFile = true; - ImageTask->CompressionQuality = (int32)EImageCompressionQuality::Default; - - //Add to write queue (async) - FModuleManager::LoadModuleChecked("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr(ImageTask)); - - ImageTask->OnCompleted = [](bool bSuccess) { - UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed")); - }; - - return true; - } - - return false; - }); - } + ImageTask->OnCompleted = [](bool bSuccess) { + UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed")); + }; + }); } } diff --git a/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckyDataTransferSubsystem.h b/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckyDataTransferSubsystem.h index 120913ad..4f7ee3dd 100644 --- a/Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckyDataTransferSubsystem.h +++ b/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 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 = ""); //-------------------------------------------------------// }; diff --git a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp b/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp index b3398ab4..b313574d 100644 --- a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp +++ b/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 Sensors; + UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors); + + for (const auto Sensor : Sensors) + { + if (const auto Camera = Cast(Sensor)) Cameras.Add(Camera); + } +} + void UEpisodeSubSystem::ConfigureDataCapture() { if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem()) @@ -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() - // 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()) { // 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; diff --git a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp index 851d11bb..9ecd187d 100644 --- a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp +++ b/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); } diff --git a/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp b/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp index 2f0848f2..8335967b 100644 --- a/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp +++ b/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 Sensors; - UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors); - - for (const auto Sensor : Sensors) - { - if (const auto Camera = Cast(Sensor)) Cameras.Add(Camera); - } -} diff --git a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h b/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h index 92cd95be..755de0b2 100644 --- a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h +++ b/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 CurrentRobot; + // --------------------- + // ------ SENSORS ------ + // --------------------- + void InitCameras(); + TArray> Cameras; diff --git a/Source/LuckyWorldV2/Public/Robot/RobotPawn.h b/Source/LuckyWorldV2/Public/Robot/RobotPawn.h index 623f08cb..1d1df691 100644 --- a/Source/LuckyWorldV2/Public/Robot/RobotPawn.h +++ b/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> Cameras; };