Agrarsense
Tagger.cpp
Go to the documentation of this file.
1// Copyright (c) 2023 FrostBit Software Lab at the Lapland University of Applied Sciences
2//
3// This work is licensed under the terms of the MIT license.
4// For a copy, see <https://opensource.org/licenses/MIT>.
5
6#include "Tagger.h"
12
14#include "GeoReferencingSystem.h"
15#include "AssetRegistry/AssetRegistryModule.h"
16#include "Components/SkeletalMeshComponent.h"
17#include "Components/StaticMeshComponent.h"
18#include "Engine/SkeletalMesh.h"
19#include "Engine/StaticMesh.h"
20#include "Landscape.h"
21#include "LandscapeComponent.h"
22#include "Engine/World.h"
23#include "Kismet/GameplayStatics.h"
24#include "TimerManager.h"
25
27{
28 PrimaryActorTick.bCanEverTick = false;
29}
30
32{
33 Super::BeginPlay();
34
36 {
37 return;
38 }
39
40 UWorld* World = GetWorld();
41
42 // Create TMap<FString, ELabels>
43 LabelMap = UEnumUtilities::CreateStringEnumMap<ELabels>("/Script/Agrarsense.ELabels");
44
45 ActorSpawnedDelegateHandle = World->AddOnActorSpawnedHandler(FOnActorSpawned::FDelegate::CreateUObject(this, &ATagger::OnActorSpawned));
46
47 // Tag all existing Actors in the World after small delay
48 FTimerHandle Handle;
49 World->GetTimerManager().SetTimer(Handle, FTimerDelegate::CreateLambda([this]
50 {
52 }), 0.200f, false);
53}
54
55void ATagger::EndPlay(const EEndPlayReason::Type EndPlayReason)
56{
57 Super::EndPlay(EndPlayReason);
58
59 TaggedActors.Empty();
61}
62
64{
65 TArray<AActor*> Actors;
66 UGameplayStatics::GetAllActorsOfClass(GetWorld(), AActor::StaticClass(), Actors);
67
68 TaggedActors.Reserve(Actors.Num());
69
70 // Sort actors based on distance to the origin (0,0,0)
71 // This way we can ensure that the actors are tagged in a consistent order each run
72 const FVector Origin(0.0f, 0.0f, 0.0f);
73 Actors.Sort([&](const AActor& A, const AActor& B)
74 {
75 return FVector::DistSquared(A.GetActorLocation(), Origin) < FVector::DistSquared(B.GetActorLocation(), Origin);
76 });
77
78 for (AActor* Actor : Actors)
79 {
80 if (Actor)
81 {
82 TagActor(*Actor);
83 }
84 }
85}
86
87void ATagger::OnActorSpawned(AActor* Actor)
88{
89 if (Actor)
90 {
91 TagActor(*Actor);
92 }
93}
94
95template <typename T>
96static auto CastEnum(T label)
97{
98 return static_cast<typename std::underlying_type<T>::type>(label);
99}
100
101void ATagger::TagActor(AActor& Actor)
102{
103 // Check if this is a Landscape actor
104 // if so, apply same ID and label to all landscape and its components.
105 ALandscape* Landscape = Cast<ALandscape>(&Actor);
106 if (Landscape)
107 {
108 TagLandscape(Landscape, Actor);
109 return;
110 }
111
112 // Iterate all static meshes
113 TArray<UStaticMeshComponent*> StaticMeshComponents;
114 Actor.GetComponents<UStaticMeshComponent>(StaticMeshComponents);
115 for (UStaticMeshComponent* Component : StaticMeshComponents)
116 {
117 if (Component)
118 {
119 ELabels Label = GetLabelFromStaticComponent(Component);
120 int32 InstanceSegmentationID = SetStencilValue(*Component, Actor, Label);
121
122 FTaggedActorData ActorData;
123 ActorData.Actor = &Actor;
124 ActorData.InstanceID = InstanceSegmentationID;
125 ActorData.StaticMeshComponent = Component;
126 ActorData.Label = Label;
127 TaggedActors.Add(ActorData);
128 }
129 }
130
131 // Iterate all skeletal meshes
132 TArray<USkeletalMeshComponent*> SkeletalMeshComponents;
133 Actor.GetComponents<USkeletalMeshComponent>(SkeletalMeshComponents);
134 for (USkeletalMeshComponent* Component : SkeletalMeshComponents)
135 {
136 if (Component)
137 {
139 int32 InstanceSegmentationID = SetStencilValue(*Component, Actor, Label);
140
141 FTaggedActorData ActorData;
142 ActorData.Actor = &Actor;
143 ActorData.InstanceID = InstanceSegmentationID;
144 ActorData.SkeletalMeshComponent = Component;
145 ActorData.Label = Label;
146 TaggedActors.Add(ActorData);
147 }
148 }
149}
150
151int32 ATagger::SetStencilValue(UPrimitiveComponent& Component, AActor& Actor, const ELabels& Label)
152{
153 Component.SetRenderCustomDepth(true);
154
155 const int32 ActorLabelStencilValue = static_cast<int32>(Label);
156 float ActorLabelStencilValueFloat = static_cast<float>(ActorLabelStencilValue);
157
158 Component.SetCustomDepthStencilValue(ActorLabelStencilValue + 1);
159
160 int32& Counter = LabelCounters.FindOrAdd(Label); // default 0 on first use
161 Counter += 1; // assign next id
162 const int32 CurrentIntID = Counter; // 1..n for this label
163 const float CurrentID = static_cast<float>(CurrentIntID);
164
165 // Write custom primitive data (X=per-label instance id, Y=semantic/stencil)
166 Component.SetCustomPrimitiveDataVector4f(4,
167 FVector4f(CurrentID,
168 ActorLabelStencilValueFloat,
169 0.0f,
170 0.0f));
171
172 Component.MarkRenderStateDirty();
173
174 /*UE_LOG(LogTemp, Log, TEXT("Per-label ID %f -> Actor %s (Label %d)"),
175 CurrentID, *Actor.GetName(), static_cast<int32>(Label));*/
176
177 return CurrentIntID;
178}
179
181{
182 if (TaggedActors.IsEmpty())
183 {
184 return;
185 }
186
188 if (!CSVFile)
189 {
190 return;
191 }
192
193 for (const FTaggedActorData& TaggedActorData : TaggedActors)
194 {
195 AActor* Actor = TaggedActorData.Actor;
196 if (!Actor)
197 {
198 continue;
199 }
200
201 WriteComponentToCSV(TaggedActorData);
202 }
203
204 CSVFile->Close();
205}
206
207AActor* ATagger::TryGetActorByInstanceID(int32 LabelIndex, int32 InstanceID)
208{
209 const ELabels Label = static_cast<ELabels>(LabelIndex);
210
211 for (const FTaggedActorData& TaggedActorData : TaggedActors)
212 {
213 if (TaggedActorData.Label == Label && TaggedActorData.InstanceID == InstanceID)
214 {
215 return TaggedActorData.Actor;
216 }
217 }
218
219 return nullptr;
220}
221
223{
224 GeoReferencingSystem = AGeoReferencingSystem::GetGeoReferencingSystem(GetWorld());
225
226 // CSV file settings
227 FCSVFileSettings Settings;
229 Settings.Append = false;
230 Settings.CreateUnique = true;
232
233 FString MapName = GetWorld()->GetMapName();
234 FString FileName = MapName + TEXT("_ObjectLocations");
235
236 if (MapName.Contains("dev"))
237 {
238 Settings.QueueSize = 65536;
239 }
240
241 CSVFile = UCSVFile::CreateCSVFile(FileName, Settings);
242
243 TArray<FString> Header = { TEXT("Name"), TEXT("Label"), TEXT("InstanceID"), TEXT("LocationX"), TEXT("LocationY"),
244 TEXT("LocationZ"), TEXT("Yaw"), TEXT("Pitch"), TEXT("Roll"), TEXT("Bounds")};
246 {
247 Header.Append({ TEXT("Latitude"), TEXT("Longitude"), TEXT("Altitude") });
248 }
249 CSVFile->WriteRow(Header);
250}
251
253{
254 AActor* Actor = Data.Actor;
255 if (!Actor)
256 {
257 return;
258 }
259
260 FString AssetPath;
261 FString MeshName;
262 FVector Location = Actor->GetActorLocation();
263 FRotator Rotation = Actor->GetActorRotation();
264 FBoxSphereBounds Bounds;
265
266 // Try StaticMeshComponent first
267 if (Data.StaticMeshComponent && Data.StaticMeshComponent->GetStaticMesh())
268 {
269 UStaticMesh* StaticMesh = Data.StaticMeshComponent->GetStaticMesh();
270 AssetPath = StaticMesh->GetPathName();
271 MeshName = StaticMesh->GetName();
272 Bounds = StaticMesh->GetBounds();
273 }
274 // If not valid, try SkeletalMeshComponent
275 else if (Data.SkeletalMeshComponent && Data.SkeletalMeshComponent->GetSkeletalMeshAsset())
276 {
277 USkeletalMesh* SkeletalMesh = Data.SkeletalMeshComponent->GetSkeletalMeshAsset();
278 AssetPath = SkeletalMesh->GetPathName();
279 MeshName = SkeletalMesh->GetName();
280 Bounds = SkeletalMesh->GetBounds();
281 }
282 else
283 {
284 return;
285 }
286
287 // Check asset path
288 if (!AssetPath.StartsWith(TEXT("/Game/Agrarsense/Models")))
289 {
290 return;
291 }
292
293 TArray<FString> Row;
294 Row.Reserve(GeoReferencingSystem ? 13 : 11);
295
296 // Mesh name
297 Row.Add(MeshName);
298
299 // Label as: Label (NUM) like Forwarder (18)
300 const UEnum* LabelsEnum = StaticEnum<ELabels>();
301 const int32 LabelIndex = static_cast<int32>(static_cast<uint8>(Data.Label));
302 const FString LabelName = LabelsEnum
303 ? LabelsEnum->GetNameStringByValue((int64)LabelIndex)
304 : TEXT("Unknown");
305 Row.Add(FString::Printf(TEXT("%s (%d)"), *LabelName, LabelIndex));
306
307 // Instance ID
308 Row.Add(FString::FromInt(Data.InstanceID));
309
310 // Transform
311 Row.Add(FString::Printf(TEXT("%.2f"), Location.X));
312 Row.Add(FString::Printf(TEXT("%.2f"), Location.Y));
313 Row.Add(FString::Printf(TEXT("%.2f"), Location.Z));
314 Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Yaw));
315 Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Pitch));
316 Row.Add(FString::Printf(TEXT("%.2f"), Rotation.Roll));
317
318 // Bounds
319 Row.Add(FString::Printf(
320 TEXT("Origin=(%.2f,%.2f,%.2f),BoxExtent=(%.2f,%.2f,%.2f),SphereRadius=%.2f"),
321 Bounds.Origin.X, Bounds.Origin.Y, Bounds.Origin.Z,
322 Bounds.BoxExtent.X, Bounds.BoxExtent.Y, Bounds.BoxExtent.Z,
323 Bounds.SphereRadius));
324
325 // Optional: Georeferenced coords
327 {
328 const FGeographicCoordinates GeoCoords =
330 Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Latitude));
331 Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Longitude));
332 Row.Add(FString::Printf(TEXT("%.8f"), GeoCoords.Altitude));
333 }
334
335 if (CSVFile)
336 {
337 CSVFile->WriteRow(Row);
338 }
339}
340
341void ATagger::TagLandscape(ALandscape* Landscape, AActor& Actor)
342{
343 if (!Landscape)
344 {
345 return;
346 }
347
348 // Do not change. hardcoded values for landscape.
349 // So each landscape section has same ID and label.
350 constexpr float LANDSCAPE_LABEL = 1.0f;
351 constexpr float LANDSCAPE_ID = 0.0f;
352
353 // For some reason we needed to loop UStaticMeshComponents instead of ULandscapeComponents
354 TArray<UStaticMeshComponent*> StaticMeshComponents;
355 Actor.GetComponents<UStaticMeshComponent>(StaticMeshComponents);
356 for (UStaticMeshComponent* Component : StaticMeshComponents)
357 {
358 if (Component)
359 {
360 Component->SetRenderCustomDepth(true);
361 Component->SetCustomDepthStencilValue(static_cast<int32>(ELabels::Terrain) + 1);
362
363#ifdef InstanceSegmentationPass_EXISTS
364 // Set custom primitive data, used for instance segmentation.
365 Component->SetCustomPrimitiveDataVector4f(4,
366 FVector4f(LANDSCAPE_ID,
367 LANDSCAPE_LABEL,
368 0.0f,
369 0.0f));
370#endif
371 }
372 }
373}
374
376{
377 const ELabels* FoundLabel = LabelMap.Find(String);
378 if (FoundLabel)
379 {
380 return *FoundLabel;
381 }
382 else
383 {
384#if WITH_EDITOR
385 UE_LOG(LogTemp, Warning, TEXT("Tagger.cpp: Missing Label: %s"), *String);
386#endif
387 return ELabels::None;
388 }
389}
390
391ELabels ATagger::GetLabelFromStaticComponent(UStaticMeshComponent* StaticMeshComponentPtr)
392{
393 ELabels Label = ELabels::None;
394
395 if (!StaticMeshComponentPtr)
396 {
397 return Label;
398 }
399
400 Label = GetLabelFromTag(StaticMeshComponentPtr);
401 if (Label != ELabels::None)
402 {
403 return Label;
404 }
405
406 UStaticMesh* MeshPtr = StaticMeshComponentPtr->GetStaticMesh();
407 if (MeshPtr)
408 {
409 FString Name;
410 Label = GetLabelFromPath(MeshPtr->GetPathName(), Name);
411
412 if (Label != ELabels::None)
413 {
414 StaticMeshComponentPtr->ComponentTags.Insert(FName(Name), 0);
415 }
416 }
417
418 return Label;
419}
420
421ELabels ATagger::GetLabelFromSkeletalMeshComponent(USkeletalMeshComponent* USkeletalMeshComponentPtr)
422{
423 ELabels Label = ELabels::None;
424
425 if (!USkeletalMeshComponentPtr)
426 {
427 return Label;
428 }
429
430 Label = GetLabelFromTag(USkeletalMeshComponentPtr);
431 if (Label != ELabels::None)
432 {
433 return Label;
434 }
435
436 USkeletalMesh* PhysicsAsset = USkeletalMeshComponentPtr->GetSkeletalMeshAsset();
437 if (PhysicsAsset)
438 {
439 FString Name;
440 Label = GetLabelFromPath(PhysicsAsset->GetPathName(), Name);
441
442 if (Label != ELabels::None)
443 {
444 USkeletalMeshComponentPtr->ComponentTags.Insert(FName(Name), 0);
445 }
446 }
447
448 return Label;
449}
450
451ELabels ATagger::GetLabelFromPath(const FString& Path, FString& FolderName)
452{
453 ELabels Label = ELabels::None;
454
455 // If the Path doesn't contain Models,
456 // the asset is not under /Content/Agrarsense/Models directory
457 // meaing we ignore it.
458 if (!Path.Contains("Models"))
459 {
460 return Label;
461 }
462
463 TArray<FString> StringArray;
464 Path.ParseIntoArray(StringArray, TEXT("/"), false);
465
466 FolderName = StringArray[StringArray.Num() - 2];
467
468 Label = GetLabelFromString(FolderName);
469
470 return Label;
471}
ELabels
Definition: ObjectLabel.h:14
static auto CastEnum(T label)
Definition: Tagger.cpp:96
void ExportObjectLocationsToCSV()
Definition: Tagger.cpp:180
virtual void BeginPlay() override
Definition: Tagger.cpp:31
void CreateCSVFile()
Definition: Tagger.cpp:222
TMap< FString, ELabels > LabelMap
Definition: Tagger.h:115
int32 SetStencilValue(UPrimitiveComponent &Component, AActor &Actor, const ELabels &Label)
Definition: Tagger.cpp:151
ELabels GetLabelFromPath(const FString &Path, FString &FolderName)
Definition: Tagger.cpp:451
ELabels GetLabelFromStaticComponent(UStaticMeshComponent *StaticMeshComponentPtr)
Definition: Tagger.cpp:391
void OnActorSpawned(AActor *Actor)
Definition: Tagger.cpp:87
ELabels GetLabelFromTag(const T *Object)
Definition: Tagger.h:95
void WriteComponentToCSV(const FTaggedActorData &Data)
Definition: Tagger.cpp:252
void TagExistingActors()
Definition: Tagger.cpp:63
AGeoReferencingSystem * GeoReferencingSystem
Definition: Tagger.h:107
AActor * TryGetActorByInstanceID(int32 LabelIndex, int32 InstanceID)
Definition: Tagger.cpp:207
TArray< FTaggedActorData > TaggedActors
Definition: Tagger.h:111
ATagger()
Definition: Tagger.cpp:26
void TagActor(AActor &Actor)
Definition: Tagger.cpp:101
UCSVFile * CSVFile
Definition: Tagger.h:109
FDelegateHandle ActorSpawnedDelegateHandle
Definition: Tagger.h:113
ELabels GetLabelFromSkeletalMeshComponent(USkeletalMeshComponent *USkeletalMeshComponentPtr)
Definition: Tagger.cpp:421
TMap< ELabels, int32 > LabelCounters
Definition: Tagger.h:125
virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override
Definition: Tagger.cpp:55
ELabels GetLabelFromString(const FString &String)
Definition: Tagger.cpp:375
void TagLandscape(ALandscape *Landscape, AActor &Actor)
Definition: Tagger.cpp:341
static bool IsPlayingInMainMenu()
static UCSVFile * CreateCSVFile(const FString &FileNameWithoutExtension, const FCSVFileSettings &Settings)
Definition: CSVFile.cpp:14
void Close()
Definition: CSVFile.cpp:115
void WriteRow(const TArray< FString > &Cells)
Definition: CSVFile.cpp:77
static FGeographicCoordinates UnrealToGeographicCoordinates(AGeoReferencingSystem *GeoReferencingSystem, const FVector &Position)
bool CreateUnique
Definition: CSVFile.h:42
ECSVDelimiter Delimiter
Definition: CSVFile.h:36
int32 QueueSize
Definition: CSVFile.h:48
FCSVFileWriteOptions FileWriteOption
Definition: CSVFile.h:45
ELabels Label
Definition: Tagger.h:40
AActor * Actor
Definition: Tagger.h:28
UStaticMeshComponent * StaticMeshComponent
Definition: Tagger.h:31
int32 InstanceID
Definition: Tagger.h:37
USkeletalMeshComponent * SkeletalMeshComponent
Definition: Tagger.h:34