LuckyRobots/LuckyWorld
11
2
Fork 0
Code Issues Pull Requests Actions Packages Projects 1 Releases Wiki Activity

UNSTBL - Refacto

Browse Source 
+ 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
This commit is contained in:
Jb win
2025-05-06 16:23:17 +07:00
parent 022f7dd518
commit 42fc0e5b74
9 changed files with 107 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
Content/Levels/kitchenLevel/BP-Sensor.uasset Normal file
View File

Binary file not shown.

BIN
Content/Levels/kitchenLevel/kitchenLevel1.umap
View File

Binary file not shown.

131
Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Private/LuckyDataTransferSubsystem.cpp
View File

@ -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<FColor> 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<FColor> 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<int32>(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<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 = static_cast<int32>(EImageCompressionQuality::Default);
// Add to write queue (async)
FModuleManager::LoadModuleChecked<IImageWriteQueueModule>("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr<IImageWriteTaskBase>(ImageTask));
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;
});
}
ImageTask->OnCompleted = [](bool bSuccess) {
UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed"));
};
});
}
}

6
Plugins/LuckyDataTransfer/Source/LuckyDataTransfer/Public/LuckyDataTransferSubsystem.h
View File

@ -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 = "");
//-------------------------------------------------------//
};

32
Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp
View File

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

5
Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp
View File

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

15
Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp
View File

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

6
Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h
View File

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

7
Source/LuckyWorldV2/Public/Robot/RobotPawn.h
View File

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