Tagger.cpp

Go to the documentation of this file.

// Copyright (c) 2023 FrostBit Software Lab at the Lapland University of Applied Sciences
//
// This work is licensed under the terms of the MIT license.
// For a copy, see <https://opensource.org/licenses/MIT>.
 
#include "Tagger.h"
#include "Agrarsense/Utils/EnumUtilities.h"
#include "Agrarsense/Walker/Walker.h"
#include "Agrarsense/InstancedRendering/InstancedActor.h"
#include "Agrarsense/Settings/AgrarsenseSettings.h"
#include "Agrarsense/Logging/CSVFile.h"
 
#include "Agrarsense/Utils/CoordinateConversionUtilities.h"
#include "GeoReferencingSystem.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "Components/SkeletalMeshComponent.h"
#include "Components/StaticMeshComponent.h"
#include "Engine/SkeletalMesh.h"
#include "Engine/StaticMesh.h"
#include "Landscape.h"
#include "LandscapeComponent.h"
#include "Engine/World.h"
#include "Kismet/GameplayStatics.h"
#include "TimerManager.h"
 
ATagger::ATagger()
{
    PrimaryActorTick.bCanEverTick = false;
}
 
void ATagger::BeginPlay()
{
    Super::BeginPlay();
 
    if (UAgrarsenseSettings::IsPlayingInMainMenu()) 
    {
        return;
    }
 
    UWorld* World = GetWorld();
 
    // Create TMap<FString, ELabels>
    LabelMap = UEnumUtilities::CreateStringEnumMap<ELabels>("/Script/Agrarsense.ELabels");
 
    ActorSpawnedDelegateHandle = World->AddOnActorSpawnedHandler(FOnActorSpawned::FDelegate::CreateUObject(this, &ATagger::OnActorSpawned));
 
    // Tag all existing Actors in the World after small delay
    FTimerHandle Handle;
    World->GetTimerManager().SetTimer(Handle, FTimerDelegate::CreateLambda([this]
        {
            TagExistingActors();
    }), 0.200f, false);
}
 
void ATagger::EndPlay(const EEndPlayReason::Type EndPlayReason)
{
    Super::EndPlay(EndPlayReason);
 
    TaggedActors.Empty();
    ActorSpawnedDelegateHandle.Reset();
}
 
void ATagger::TagExistingActors()
{
    TArray<AActor*> Actors;
    UGameplayStatics::GetAllActorsOfClass(GetWorld(), AActor::StaticClass(), Actors);
 
    TaggedActors.Reserve(Actors.Num());
 
    // Sort actors based on distance to the origin (0,0,0)
    // This way we can ensure that the actors are tagged in a consistent order each run
    const FVector Origin(0.0f, 0.0f, 0.0f);
    Actors.Sort([&](const AActor& A, const AActor& B)
        {
            return FVector::DistSquared(A.GetActorLocation(), Origin) < FVector::DistSquared(B.GetActorLocation(), Origin);
        });
 
    for (AActor* Actor : Actors)
    {
        if (Actor)
        {
            TagActor(*Actor);
        }
    }
}
 
void ATagger::OnActorSpawned(AActor* Actor)
{
    if (Actor)
    {
        TagActor(*Actor);
    }
}
 
template <typename T>
static auto CastEnum(T label)
{
    return static_cast<typename std::underlying_type<T>::type>(label);
}
 
void ATagger::TagActor(AActor& Actor)
{
//#ifndef InstanceSegmentationPass_EXISTS
//  AInstancedActor* InstancedActor = Cast<AInstancedActor>(&Actor);
//  if (InstancedActor)
//  {
//      if (InstancedActor->IsCustomDepthSetup())
//      {
//          // Custom depth has been setup beforehand. Return
//          return;
//      }
//  }
//
//  AWalker* Walker = Cast<AWalker>(&Actor);
//  if (Walker)
//  {
//      // If the actor is AWalker, return. These are tagged in editor.
//      return;
//  }
//#endif
 
    // Check if this is a Landscape actor
    // if so, apply same ID and label to all landscape and its components.
    ALandscape* Landscape = Cast<ALandscape>(&Actor);
    if (Landscape)
    {
        TagLandscape(Landscape, Actor);
        return;
    }
 
    // Iterate all static meshes
    TArray<UStaticMeshComponent*> StaticMeshComponents;
    Actor.GetComponents<UStaticMeshComponent>(StaticMeshComponents);
    for (UStaticMeshComponent* Component : StaticMeshComponents)
    {
        if (Component)
        {
            ELabels Label = GetLabelFromStaticComponent(Component);
            float InstanceSegmentationID = SetStencilValue(*Component, Actor, Label);
 
            FTaggedActorData ActorData;
            ActorData.Actor = &Actor;
            ActorData.InstanceID = InstanceSegmentationID;
            ActorData.StaticMeshComponent = Component;
            ActorData.Label = Label;
            TaggedActors.Add(ActorData);
        }
    }
 
    // Iterate all skeletal meshes
    TArray<USkeletalMeshComponent*> SkeletalMeshComponents;
    Actor.GetComponents<USkeletalMeshComponent>(SkeletalMeshComponents);
    for (USkeletalMeshComponent* Component : SkeletalMeshComponents)
    {
        if (Component)
        {
            ELabels Label = GetLabelFromSkeletalMeshComponent(Component);
            float InstanceSegmentationID = SetStencilValue(*Component, Actor, Label);
 
            FTaggedActorData ActorData;
            ActorData.Actor = &Actor;
            ActorData.InstanceID = InstanceSegmentationID;
            ActorData.SkeletalMeshComponent = Component;
            ActorData.Label = Label;
            TaggedActors.Add(ActorData);
        }
    }
}
 
float ATagger::SetStencilValue(UPrimitiveComponent& Component, AActor& Actor, const ELabels& Label)
{
    Component.SetRenderCustomDepth(true);
 
    const int32 ActorLabelStencilValue = static_cast<int32>(Label) + 1;
    float ActorLabelStencilValueFloat = static_cast<float>(ActorLabelStencilValue);
 
#if ENGINE_MAJOR_VERSION < 5 || (ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION <= 5)
    ActorLabelStencilValueFloat = static_cast<float>(ActorLabelStencilValue - 1);
#endif
 
    Component.SetCustomDepthStencilValue(ActorLabelStencilValue);
 
    float CurrentID = ID;
 
    // Set custom primitive data, used for instance segmentation.
    Component.SetCustomPrimitiveDataVector4(4,
        FVector4(CurrentID,
            ActorLabelStencilValueFloat,
            0.0f,
            0.0f));
 
    // Increase ID counter
    ID += 1.0f;
 
    return CurrentID;
}
 
void ATagger::ExportObjectLocationsToCSV()
{
    if (TaggedActors.IsEmpty()) 
    {
        return;
    }
 
    CreateCSVFile();
    if (!CSVFile)
    {
        return;
    }
 
    for (const FTaggedActorData& TaggedActorData : TaggedActors)
    {
        AActor* Actor = TaggedActorData.Actor;
        if (!Actor)
        {
            continue;
        }
 
        WriteComponentToCSV(TaggedActorData);
    }
 
    CSVFile->Close();
}
 
AActor* ATagger::GetActorByInstanceID(float InstanceID)
{
    for (const FTaggedActorData& TaggedActorData : TaggedActors)
    {
        if (TaggedActorData.InstanceID == InstanceID)
        {
            return TaggedActorData.Actor;
        }
    }
 
    return nullptr;
}
 
void ATagger::CreateCSVFile()
{
    GeoReferencingSystem = AGeoReferencingSystem::GetGeoReferencingSystem(GetWorld());
 
    // CSV file settings
    FCSVFileSettings Settings;
    Settings.Delimiter = ECSVDelimiter::Semicolon;
    Settings.Append = false;
    Settings.CreateUnique = true;
    Settings.FileWriteOption = FCSVFileWriteOptions::Queue;
 
    FString MapName = GetWorld()->GetMapName();
    FString FileName = MapName + TEXT("_ObjectLocations");
 
    if (MapName.Contains("dev"))
    {
        Settings.QueueSize = 65536;
    }
 
    CSVFile = UCSVFile::CreateCSVFile(FileName, Settings);
 
    TArray<FString> Header = { TEXT("Name"), TEXT("InstanceID"), TEXT("LocationX"), TEXT("LocationY"),
        TEXT("LocationZ"), TEXT("Yaw"), TEXT("Pitch"), TEXT("Roll"), TEXT("Bounds")};
    if (GeoReferencingSystem)
    {
        Header.Append({ TEXT("Latitude"), TEXT("Longitude"), TEXT("Altitude") });
    }
    CSVFile->WriteRow(Header);
}
 
void ATagger::WriteComponentToCSV(const FTaggedActorData Data)
{
    AActor* Actor = Data.Actor;
 
    if (!Actor)
    {
        return;
    }
 
    FString AssetPath;
    FString MeshName;
    FVector Location = Actor->GetActorLocation();
    FRotator Rotation = Actor->GetActorRotation();
    FBoxSphereBounds Bounds;
 
    // Try StaticMeshComponent first
    if (Data.StaticMeshComponent && Data.StaticMeshComponent->GetStaticMesh())
    {
        UStaticMesh* StaticMesh = Data.StaticMeshComponent->GetStaticMesh();
        AssetPath = StaticMesh->GetPathName();
        MeshName = StaticMesh->GetName();
        Bounds = StaticMesh->GetBounds();
    }
    // If not valid, try SkeletalMeshComponent
    else if (Data.SkeletalMeshComponent && Data.SkeletalMeshComponent->GetSkeletalMeshAsset())
    {
        USkeletalMesh* SkeletalMesh = Data.SkeletalMeshComponent->GetSkeletalMeshAsset();
        AssetPath = SkeletalMesh->GetPathName();
        MeshName = SkeletalMesh->GetName();
        Bounds = SkeletalMesh->GetBounds();
    }
    else
    {
        return;
    }
 
    // Check asset path
    if (!AssetPath.StartsWith(TEXT("/Game/Agrarsense/Models")))
    {
        return;
    }
 
    TArray<FString> Row;
    Row.Reserve(GeoReferencingSystem ? 12 : 9);
 
    Row.Add(MeshName);
    Row.Add(FString::Printf(TEXT("%.0f"), Data.InstanceID));
 
    // TODO add semantic label to the CSV
    Row.Add(FString::Printf(TEXT("%.2f"), Location.X));
    Row.Add(FString::Printf(TEXT("%.2f"), Location.Y));
    Row.Add(FString::Printf(TEXT("%.2f"), Location.Z));
    Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Yaw));
    Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Pitch));
    Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Roll));
 
    // Bounds
    Row.Add(FString::Printf(TEXT("Origin=(%.2f,%.2f,%.2f),BoxExtent=(%.2f,%.2f,%.2f),SphereRadius=%.2f"),
            Bounds.Origin.X, Bounds.Origin.Y, Bounds.Origin.Z,
            Bounds.BoxExtent.X, Bounds.BoxExtent.Y, Bounds.BoxExtent.Z,
            Bounds.SphereRadius));
 
    if (GeoReferencingSystem)
    {
        const FGeographicCoordinates GeoCoords = UCoordinateConversionUtilities::UnrealToGeographicCoordinates(GeoReferencingSystem, Location);
        Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Latitude));
        Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Longitude));
        Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Altitude));
    }
 
    if (CSVFile)
    {
        CSVFile->WriteRow(Row);
    }
}
 
void ATagger::TagLandscape(ALandscape* Landscape, AActor& Actor)
{
    if (!Landscape) 
    {
        return;
    }
 
    constexpr float LANDSCAPE_ID = 0.0f;
 
    // For some reason we needed to loop UStaticMeshComponents instead of ULandscapeComponents
    TArray<UStaticMeshComponent*> StaticMeshComponents;
    Actor.GetComponents<UStaticMeshComponent>(StaticMeshComponents);
    for (UStaticMeshComponent* Component : StaticMeshComponents)
    {
        if (Component)
        {
            Component->SetRenderCustomDepth(true);
            Component->SetCustomDepthStencilValue(static_cast<int32>(ELabels::Terrain) + 1);
 
#ifdef InstanceSegmentationPass_EXISTS
            // Set custom primitive data, used for instance segmentation.
            Component->SetCustomPrimitiveDataVector4(4,
                FVector4(LANDSCAPE_ID,
                    (float)CastEnum(ELabels::Terrain),
                    0.0f,
                    0.0f));
#endif
        }
    }
}
 
ELabels ATagger::GetLabelFromString(const FString& String)
{
    const ELabels* FoundLabel = LabelMap.Find(String);
    if (FoundLabel)
    {
        return *FoundLabel;
    }
    else
    {
#if WITH_EDITOR
        UE_LOG(LogTemp, Warning, TEXT("Tagger.cpp: Missing Label: %s"), *String);
#endif
        return ELabels::None;
    }
}
 
ELabels ATagger::GetLabelFromStaticComponent(UStaticMeshComponent* StaticMeshComponentPtr)
{
    ELabels Label = ELabels::None;
 
    if (!StaticMeshComponentPtr)
    {
        return Label;
    }
 
    Label = GetLabelFromTag(StaticMeshComponentPtr);
    if (Label != ELabels::None)
    {
        return Label;
    }
 
    UStaticMesh* MeshPtr = StaticMeshComponentPtr->GetStaticMesh();
    if (MeshPtr)
    {
        FString Name;
        Label = GetLabelFromPath(MeshPtr->GetPathName(), Name);
 
        if (Label != ELabels::None)
        {
            StaticMeshComponentPtr->ComponentTags.Insert(FName(Name), 0);
        }
    }
 
    return Label;
}
 
ELabels ATagger::GetLabelFromSkeletalMeshComponent(USkeletalMeshComponent* USkeletalMeshComponentPtr)
{
    ELabels Label = ELabels::None;
 
    if (!USkeletalMeshComponentPtr)
    {
        return Label;
    }
 
    Label = GetLabelFromTag(USkeletalMeshComponentPtr);
    if (Label != ELabels::None)
    {
        return Label;
    }
 
    USkeletalMesh* PhysicsAsset = USkeletalMeshComponentPtr->GetSkeletalMeshAsset();
    if (PhysicsAsset)
    {
        FString Name;
        Label = GetLabelFromPath(PhysicsAsset->GetPathName(), Name);
 
        if (Label != ELabels::None)
        {
            USkeletalMeshComponentPtr->ComponentTags.Insert(FName(Name), 0);
        }
    }
 
    return Label;
}
 
ELabels ATagger::GetLabelFromPath(const FString& Path, FString& FolderName)
{
    ELabels Label = ELabels::None;
 
    // If the Path doesn't contain Models,
    // the asset is not under /Content/Agrarsense/Models directory
    // meaing we ignore it.
    if (!Path.Contains("Models"))
    {
        return Label;
    }
 
    TArray<FString> StringArray;
    Path.ParseIntoArray(StringArray, TEXT("/"), false);
 
    FolderName = StringArray[StringArray.Num() - 2];
 
    Label = GetLabelFromString(FolderName);
 
    return Label;
}