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 "AssetRegistry/AssetRegistryModule.h"
#include "Components/SkeletalMeshComponent.h"
#include "Components/StaticMeshComponent.h"
#include "Engine/SkeletalMesh.h"
#include "Engine/StaticMesh.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();
    if (!World)
    {
        return;
    }
 
    // 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.500f, false);
}
 
void ATagger::EndPlay(const EEndPlayReason::Type EndPlayReason)
{
    Super::EndPlay(EndPlayReason);
 
    ActorSpawnedDelegateHandle.Reset();
}
 
void ATagger::TagExistingActors()
{
    TArray<AActor*> Actors;
    UGameplayStatics::GetAllActorsOfClass(GetWorld(), AActor::StaticClass(), Actors);
 
    for (AActor* Actor : Actors)
    {
        if (Actor)
        {
            TagActor(*Actor, true);
        }
    }
}
 
void ATagger::OnActorSpawned(AActor* Actor)
{
    if (Actor)
    {
        TagActor(*Actor, true);
    }
}
 
void ATagger::TagActor(AActor& Actor, bool TagForSemanticSegmentation)
{
    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;
    }
 
    // Iterate all static meshes
    TArray<UStaticMeshComponent*> StaticMeshComponents;
    Actor.GetComponents<UStaticMeshComponent>(StaticMeshComponents);
    for (UStaticMeshComponent* Component : StaticMeshComponents)
    {
        if (Component)
        {
            auto Label = GetLabelFromStaticComponent(Component);
            SetStencilValue(*Component, Label, TagForSemanticSegmentation);
        }
    }
 
    // Iterate all skeletal meshes
    TArray<USkeletalMeshComponent*> SkeletalMeshComponents;
    Actor.GetComponents<USkeletalMeshComponent>(SkeletalMeshComponents);
    for (USkeletalMeshComponent* Component : SkeletalMeshComponents)
    {
        if (Component)
        {
            ELabels Label = GetLabelFromSkeletalMeshComponent(Component);
            SetStencilValue(*Component, Label, TagForSemanticSegmentation);
        }
    }
}
 
void ATagger::SetStencilValue(UPrimitiveComponent& Component, const ELabels& Label, const bool SetRenderCustomDepth)
{
    // Only set Custom Depth enabled if Label is not None.
    Component.SetRenderCustomDepth(SetRenderCustomDepth && (Label != ELabels::None));
    Component.SetCustomDepthStencilValue(static_cast<int>(Label) + 1);
}
 
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, 
    // it means the Asset is not under /Content/Agrarsense/Models
    // and 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);
 
    //UE_LOG(LogTemp, Warning, TEXT("FolderName %s"), *FolderName);
    //UE_LOG(LogTemp, Warning, TEXT("AssetPath %s"), *Path);
 
    return Label;
}