Efficient pose estimation through iterative refinement

1. A method, comprising:

processing input data with a feature extraction stage of a machine learning model to generate a feature map;

applying an attention map to the feature map to generate an augmented feature map;

selecting a batch normalization layer, from a plurality of batch normalization layers, based on a loop count;

processing the augmented feature map with a refinement stage of the machine learning model to generate a refined feature map, wherein processing the augmented feature map with the refinement stage of the machine learning model comprises applying the batch normalization layer to the augmented feature map;

processing the refined feature map with a first regression stage of the machine learning model to generate multi-dimensional task output data; and

processing the refined feature map with an attention stage of the machine learning model to generate an updated attention map.

2. The method of claim 1 , further comprising processing the multi-dimensional task output data with a three-dimensional position generator to generate a three-dimensional position estimate.

3. The method of claim 2 , further comprising processing the three-dimensional position estimate with the multi-dimensional task output data to generate a two-dimensional position estimate.

4. The method of claim 3 , wherein the three-dimensional position generator comprises a body element pose model.

5. The method of claim 4 , wherein:

the three-dimensional position estimate comprises one or more body element joint positions in a three-dimensional space,

the two-dimensional position estimate comprises the one or more body element joint positions in a two-dimensional space, and

the multi-dimensional task output data comprises one or more of:

a rotation estimate,

a translation estimate, or

a scaling estimate.

6. The method of claim 5 , wherein the body element pose model comprises a hand model.

7. The method of claim 5 , further comprising processing the refined feature map with a second regression stage of the machine learning model to generate a task uncertainty estimate.

8. The method of claim 7 , further comprising:

determining an uncertainty associated with the three-dimensional position estimate based on the task uncertainty estimate; and

determining an uncertainty associated with the two-dimensional position estimate based on the task uncertainty estimate.

9. The method of claim 7 , further comprising:

processing the task uncertainty estimate with a loop decision stage of the machine learning model to generate a loop decision value; and

determining whether to exit processing of the machine learning model based on the loop decision value.

10. The method of claim 9 , further comprising determining to exit processing of the machine learning model based on the loop decision value.

11. The method of claim 9 , further comprising:

determining not to exit processing of the machine learning model based on the loop decision value; and

determining to exit processing of the machine learning model based on a loop count.

12. The method of claim 5 , further comprising displaying the two-dimensional position estimate on a display screen of an extended reality device.

13. The method of claim 1 , further comprising:

determining a training loss based on a training loss function, wherein the training loss function comprises:

an uncertainty loss component; and

a task loss component; and

updating the machine learning model based on the training loss.

14. A processing system, comprising:

one or more processors individually or collectively configured to execute instructions stored on one or more memories and to cause the processing system to:

process input data with a feature extraction stage of a machine learning model to generate a feature map;

apply an attention map to the feature map to generate an augmented feature map;

select a batch normalization layer from a plurality of batch normalization layers based on a loop count;

process the augmented feature map with a refinement stage of the machine learning model to generate a refined feature map, wherein in order to process the augmented feature map with the refinement stage of the machine learning model, the one or more processors are further configured to cause the processing system to apply the batch normalization layer to the augmented feature map;

process the refined feature map with a first regression stage of the machine learning model to generate multi-dimensional task output data; and

process the refined feature map with an attention stage of the machine learning model to generate an updated attention map.

15. The processing system of claim 14 , wherein the one or more processors are further configured to cause the processing system to process the multi-dimensional task output data with a three-dimensional position generator to generate a three-dimensional position estimate.

16. The processing system of claim 15 , wherein the one or more processors are further configured to cause the processing system to process the three-dimensional position estimate with the multi-dimensional task output data to generate a two-dimensional position estimate.

17. The processing system of claim 16 , wherein the three-dimensional position generator comprises a body element pose model.

18. The processing system of claim 17 , wherein:

the three-dimensional position estimate comprises one or more body element joint positions in a three-dimensional space,

the two-dimensional position estimate comprises the one or more body element joint positions in a two-dimensional space, and

the multi-dimensional task output data comprises one or more of:

a rotation estimate,

a translation estimate, or

a scaling estimate.

19. The processing system of claim 18 , wherein the body element pose model comprises a hand model.

20. The processing system of claim 18 , wherein the one or more processors are further configured to cause the processing system to display the two-dimensional position estimate on a display screen of an extended reality device.

21. The processing system of claim 16 , wherein the one or more processors are further configured to cause the processing system to process the refined feature map with a second regression stage of the machine learning model to generate a task uncertainty estimate.

22. The processing system of claim 21 , wherein the one or more processors are further configured to cause the processing system to:

determine an uncertainty associated with the three-dimensional position estimate based on the task uncertainty estimate; and

determine an uncertainty associated with the two-dimensional position estimate based on the task uncertainty estimate.

23. The processing system of claim 21 , wherein the one or more processors are further configured to cause the processing system to:

process the task uncertainty estimate with a loop decision stage of the machine learning model to generate a loop decision value; and

determine whether to exit processing of the machine learning model based on the loop decision value.

24. The processing system of claim 23 , wherein the one or more processors are further configured to cause the processing system to determine to exit processing of the machine learning model based on the loop decision value.

25. The processing system of claim 23 , wherein the one or more processors are further configured to cause the processing system to:

determine not to exit processing of the machine learning model based on the loop decision value; and

determine to exit processing of the machine learning model based on a loop count.

26. The processing system of claim 14 , wherein the one or more processors are further configured to cause the processing system to:

determine a training loss based on a training loss function, wherein the training loss function comprises:

an uncertainty loss component; and

a task loss component; and

update the machine learning model based on the training loss.

27. A non-transitory computer-readable medium comprising computer-executable instructions that, when executed by one or more processors of a processing system, cause the processing system to perform a method, the method comprising:

processing input data with a feature extraction stage of a machine learning model to generate a feature map;

applying an attention map to the feature map to generate an augmented feature map;

selecting a batch normalization layer, from a plurality of batch normalization layers, based on a loop count;

processing the augmented feature map with a refinement stage of the machine learning model to generate a refined feature map, wherein processing the augmented feature map with the refinement stage of the machine learning model comprises applying the batch normalization layer to the augmented feature map;

processing the refined feature map with a first regression stage of the machine learning model to generate multi-dimensional task output data; and

processing the refined feature map with an attention stage of the machine learning model to generate an updated attention map.

28. A processing system, comprising:

means for processing input data with a feature extraction stage of a machine learning model to generate a feature map;

means for applying an attention map to the feature map to generate an augmented feature map;

means for selecting a batch normalization layer, from a plurality of batch normalization layers, based on a loop count;

means for processing the augmented feature map with a refinement stage of the machine learning model to generate a refined feature map, wherein means for processing the augmented feature map with the refinement stage of the machine learning model comprise means for applying the batch normalization layer to the augmented feature map;

means for processing the refined feature map with a first regression stage of the machine learning model to generate multi-dimensional task output data; and

means for processing the refined feature map with an attention stage of the machine learning model to generate an updated attention map.