System and method for model-free, one-shot object pose estimation via coordinate regression
A computer implemented method and system using an object-agnostic model for predicting a pose of an object in an image receives a query image having a target object therein; receives a set of reference images of the target object from different viewpoints; encodes, using a vision transformer, the received query image and the received set of reference images to generate a set of token features for the received query image and a set of token features for the received set of reference images; extracts, using a transformer decoder, information from the set of token features for the encoded reference images with respect to a set of token features for the received query image; processes, using a prediction head, the combined set of token features to generate a 2D-3D mapping and a confidence map of the query image; and processes the 2D-3D mapping and confidence map to determine the pose of the target object in the query image.
1 . A computer-implemented method of training a machine learning model for regression on pixel-level annotations in images, the machine learning model comprising an image encoder, a feature mixer, and a decoder, the method comprising:
pre-training the image encoder and the decoder for cross-view completion on images;
constructing training tuples, each training tuple comprising a first image and one or more second images, wherein the first image is associated with dense pixel-level annotations, and wherein each of the one or more second images is associated with sparse pixel-level annotations;
generating, by the image encoder, a set of first image tokens from the first image;
generating, by the image encoder, one or more sets of second image tokens from the one or more second images;
generating, by the feature mixer, one or more sets of augmented second image tokens by augmenting each set of second image tokens with encodings of the respective sparse pixel-level annotations associated with the respective set of second image tokens, wherein the augmenting includes mixing the respective set of second image tokens and the encodings of the respective sparse pixel-level annotations;
processing, by the decoder, the set of first image tokens and the one or more sets of augmented second image tokens to generate prediction data of the machine learning model for the first image, wherein the processing comprises receiving a set of augmented second image tokens of the one or more sets of augmented second image tokens, wherein the prediction data comprises predictions for each image pixel of the first image and confidences for the predictions; and
fine-tuning the machine learning model, wherein the fine-tuning comprises adjusting parameters of the image encoder, the feature mixer, and the decoder to minimize a loss function, wherein the loss function is based on the prediction data and the dense pixel-level annotations.
2 . The method as claimed in claim 1 , wherein the feature mixer comprises a first pipeline of image-level decoders for processing the respective second set of tokens, and a second pipeline of point-level decoders for processing the respective encodings, wherein each image-level decoder comprises a first cross-attention layer, wherein each point-level decoder comprises a second cross-attention layer, wherein the mixing the respective set of second image tokens and the respective encodings comprises each first cross-attention layer receiving information from the second pipeline and each second cross-attention block receiving information from the first pipeline.
3 . The method as claimed in claim 2 , wherein the feature mixer further comprises first linear projection modules configured for processing the information from the second pipeline, and second linear projection modules configured for processing the information from the first pipeline, to account for different set sizes of the encodings and the set of second image tokens.
4 . The method as claimed in claim 3 , wherein the decoder comprises a plurality of cross-attention layers, wherein the number of cross-attention layers matches the number of second images, wherein the processing, by the decoder, the set of first image tokens and the one or more sets of augmented second image tokens comprises processing the set of first image tokens as input of the decoder and providing each set of augmented second image tokens to a respective cross-attention layer of the plurality of cross-attention layers, or
wherein the decoder comprises a single cross-attention layer, wherein the processing, by the decoder, the set of first image tokens and the one or more sets of augmented second image tokens comprises providing one set of augmented second image tokens to the single cross-attention layer at a time, wherein for each set of augmented second image tokens, intermediate prediction data are generated and the prediction data of the machine learning model is based on selecting, as the prediction data for an image pixel, an intermediate prediction for the image pixel having a highest confidence value of the confidences.
5 . The method as claimed in claim 4 , wherein the sparse pixel-level annotations and the dense pixel-level annotations each comprise annotations of image pixels with 3D coordinates of real-world object features corresponding to the image pixels, wherein the sparse pixel-level annotations are scattered over the respective image when compared with the dense pixel-level annotations, and wherein the prediction data for each image pixel of the first image comprises a 3D coordinate for the image pixel.
6 . The method as claimed in claim 5 , further comprising generating the encodings of the pixel-level annotations by a trainable pixel annotation encoder, wherein the pixel annotation encoder is fed with representations of the pixel-level annotations based on an embedding of 3D coordinates in the hypercube [−1,1] d .
7 . The method as claimed in claim 6 , wherein the embedding is an injective projection and has an inverse, wherein the inverse of the embedding is well-defined over the hypercube [−1,1] d .
8 . The method as claimed in claim 7 , wherein the embedding is defined by φ(x))=(ψ(x), (ψ(y), ψ(z)) with
ψ
(
x
)
=
[
cos
(
f
1
x
)
,
sin
(
f
1
x
)
,
cos
(
f
2
x
)
,
sin
(
f
2
x
)
,
…
]
,
ψ
(
y
)
=
[
cos
(
f
1
y
)
,
sin
(
f
1
y
)
,
cos
(
f
2
y
)
,
sin
(
f
2
y
)
,
…
]
,
ψ
(
z
)
=
[
cos
(
f
1
z
)
,
sin
(
f
1
z
)
,
cos
(
f
2
z
)
,
sin
(
f
2
z
)
,
…
]
,
wherein
f
i
=
f
0
γ
i
-
1
,
i
∈
{
1
,
…
,
d
6
}
,
f
0
>
0
,
and
γ
>
0
.
9 . The method as claimed in claim 4 , wherein the pixel-level annotation relates to optical flow, relates to information for identifying instances of objects, or relates to information for segmentation of views.
10 . The method as claimed in claim 9 , wherein the constructing training tuples comprises selecting the one or more second images from an image database based on selecting easy inliers, hard inliers, and hard outliers, wherein the easy and hard inliers are determined based on a viewpoint angle between camera poses employed to capture the respective images, and the hard outliers are selected as being images most similar to the first image.
11 . The method as claimed in claim 9 , wherein the pre-training of the image encoder along with the decoder for cross-view completion comprises pre-training a pipeline comprising an encoder and a decoder on a pair of pre-training images comprising a first pre-training image and a second pre-training image, wherein the pipeline is trained to reconstruct masked portions of the first pre-training image based on the second pre-training image.
12 . The method as claimed in claim 1 , further comprising using the trained machine learning model for performing one or more of:
inferring a camera pose of a camera used for capturing an unannotated query image;
generating a 3D-reconstruction of a scene depicted in the unannotated query image;
performing 3D completion of a sparsely annotated query image, wherein the sparsely annotated query image is employed as the second image; and
performing dense depth prediction for an unannotated query image.