135 lines
3.8 KiB
C++
135 lines
3.8 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "CoreMinimal.h"
|
|
#include "Async/AsyncWork.h"
|
|
#include "UObject/WeakObjectPtr.h"
|
|
|
|
#define UE_API METAHUMANIDENTITY_API
|
|
|
|
/**
|
|
* Predictive solvers configuration.
|
|
*/
|
|
struct FPredictiveSolversTaskConfig
|
|
{
|
|
FString TemplateDescriptionJson;
|
|
FString ConfigurationJson;
|
|
TWeakObjectPtr<class UDNAAsset> DNAAsset;
|
|
TArray<uint8> PredictiveSolverGlobalTeethTrainingData;
|
|
TArray<uint8> PredictiveSolverTrainingData;
|
|
bool bTrainPreviewSolvers = true;
|
|
};
|
|
|
|
/**
|
|
* Predictive solvers training result.
|
|
*/
|
|
struct FPredictiveSolversResult
|
|
{
|
|
TArray<uint8> PredictiveWithoutTeethSolver;
|
|
TArray<uint8> PredictiveSolvers;
|
|
bool bSuccess = false;
|
|
};
|
|
|
|
/**
|
|
* Predictive solvers worker that actually does calculations.
|
|
*/
|
|
class FPredictiveSolversWorker : public FNonAbandonableTask
|
|
{
|
|
public:
|
|
using SolverProgressFunc = TFunction<void(float)>;
|
|
using SolverCompletedFunc = TFunction<void(void)>;
|
|
|
|
FPredictiveSolversWorker(bool bInIsAsync, const FPredictiveSolversTaskConfig& InConfig, SolverProgressFunc InOnProgress, SolverCompletedFunc InOnCompleted, std::atomic<bool>& bInIsCancelled, std::atomic<float>& InProgress);
|
|
|
|
inline TStatId GetStatId() const
|
|
{
|
|
RETURN_QUICK_DECLARE_CYCLE_STAT(FPredictiveSolversWorker, STATGROUP_ThreadPoolAsyncTasks);
|
|
}
|
|
|
|
void DoWork();
|
|
|
|
private:
|
|
friend class FPredictiveSolversTask;
|
|
void RunTraining();
|
|
|
|
bool bIsAsync = true;
|
|
FPredictiveSolversTaskConfig Config;
|
|
SolverProgressFunc OnProgress;
|
|
SolverCompletedFunc OnCompleted;
|
|
std::atomic<bool>& bIsCancelled;
|
|
std::atomic<float>& Progress;
|
|
float LastProgress = 0.0f;
|
|
std::atomic<bool> bIsDone = false;
|
|
FPredictiveSolversResult Result;
|
|
};
|
|
|
|
DECLARE_DELEGATE_OneParam(FOnPredictiveSolversCompleted, const FPredictiveSolversResult&);
|
|
DECLARE_DELEGATE_OneParam(FOnPredictiveSolversProgress, float);
|
|
|
|
/**
|
|
* Predictive solver task that creates new worker for predictive solver calculations.
|
|
*/
|
|
class FPredictiveSolversTask
|
|
{
|
|
public:
|
|
UE_API FPredictiveSolversTask(const FPredictiveSolversTaskConfig& InConfig);
|
|
|
|
UE_API FPredictiveSolversResult StartSync();
|
|
UE_API void StartAsync();
|
|
|
|
/** Only triggered when task is executed asynchronously */
|
|
UE_API FOnPredictiveSolversCompleted& OnCompletedCallback();
|
|
UE_API FOnPredictiveSolversProgress& OnProgressCallback();
|
|
|
|
UE_API bool IsDone() const;
|
|
UE_API bool WasCancelled() const;
|
|
|
|
UE_API void Cancel();
|
|
UE_API void Stop();
|
|
|
|
/** Obtains current training progress in range [0..1]. Return true if it was successful (task is still active), false otherwise */
|
|
UE_API bool PollProgress(float& OutProgress) const;
|
|
|
|
private:
|
|
FPredictiveSolversTaskConfig Config;
|
|
TUniquePtr<FAsyncTask<class FPredictiveSolversWorker>> Task;
|
|
std::atomic<bool> bCancelled = false;
|
|
std::atomic<float> Progress = 0.0f;
|
|
std::atomic<bool> bSkipCallback = false;
|
|
FOnPredictiveSolversCompleted OnCompletedDelegate;
|
|
FOnPredictiveSolversProgress OnProgressDelegate;
|
|
|
|
UE_API void OnProgress_Thread(float InProgress);
|
|
UE_API void OnCompleted_Thread();
|
|
};
|
|
|
|
/**
|
|
* Singleton responsible for managing and owning predictive solver tasks.
|
|
*/
|
|
class FPredictiveSolversTaskManager
|
|
{
|
|
public:
|
|
static UE_API FPredictiveSolversTaskManager& Get();
|
|
|
|
FPredictiveSolversTaskManager() = default;
|
|
~FPredictiveSolversTaskManager() = default;
|
|
|
|
FPredictiveSolversTaskManager(FPredictiveSolversTaskManager const&) = delete;
|
|
FPredictiveSolversTaskManager& operator=(const FPredictiveSolversTaskManager&) = delete;
|
|
|
|
/** Creates new solver tasks and adds to the list of active tasks */
|
|
UE_API FPredictiveSolversTask* New(const FPredictiveSolversTaskConfig& InConfig);
|
|
|
|
/** Removes given task from the list and nullifies the reference */
|
|
UE_API bool Remove(FPredictiveSolversTask*& InOutTask);
|
|
|
|
/** Stops all active tasks */
|
|
UE_API void StopAll();
|
|
|
|
private:
|
|
TArray<TUniquePtr<FPredictiveSolversTask>> Tasks;
|
|
};
|
|
|
|
#undef UE_API
|