Unifying densepose and 3D body mesh reconstruction
Methods and systems are disclosed for generating a 3D body mesh. The system receives an image that includes a depiction of a real-world object in a real-world environment. The system applies a first machine learning model to a portion of the image that depicts the real-world object to predict a tensor of heatmaps representing vertex positions of a plurality of triangles of a 3D mesh corresponding to the real-world object. First and second heatmaps of the tensor represent respectively first and second groups of possible coordinates for a first vertex of a first triangle of the plurality of triangles. The system generates the 3D mesh based on the selected subset of the tensor of heatmaps.
1 . A method comprising:
receiving an image that includes a depiction of a real-world object;
predicting a tensor of two-dimensional (2D) heatmaps by applying a first machine learning model to a portion of the image, each 2D heatmap of the tensor of 2D heatmaps exclusively representing possible positions for a corresponding single vertex of an individual triangle of a plurality of triangles of a three-dimensional (3D) mesh corresponding to the real-world object, the first machine learning model trained to predict a part of the 3D mesh that is not visible within the image including at least one of occluded vertices, self-occluded vertices, or truncated vertices, a first heatmap of the tensor of heatmaps representing a first group of possible coordinates exclusively for a single first vertex of a first triangle of the plurality of triangles, a second heatmap of the tensor of heatmaps representing a second group of possible coordinates exclusively for a single second vertex of the same first triangle as the first heatmap;
selecting a subset of the tensor of heatmaps from the tensor of heatmaps, the subset of the tensor of heatmaps comprising a first coordinate of the first group of possible coordinates and a second coordinate of the second group of possible coordinates; and
generating the 3D mesh based on the selected subset of the tensor of heatmaps, the 3D mesh being combined with an inpainted mesh to form a reconstructed 3D mesh.
2 . The method of claim 1 , further comprising:
predicting a plurality of bounding boxes for each of a plurality of real-world objects depicted in the image by applying a second machine learning model to the image; and
selecting a first bounding box from the plurality of bounding boxes to identify the portion of the image.
3 . The method of claim 1 , wherein the tensor of heatmaps represent a sparsepose corresponding to a range of possible vertices of a densepose.
4 . The method of claim 1 , wherein the 3D mesh comprises an initial 3D mesh, further comprising:
predicting a set of 3D key points corresponding to a skeleton of the real-world object by applying a second machine learning model to the portion of the image;
generating, in parallel with the initial 3D mesh, an intermediate 3D mesh using the set of 3D key points representing a per-vertex depth of the real-world object;
fusing the initial 3D mesh with the intermediate 3D mesh to form a reconstructed 3D mesh, wherein generating the intermediate 3D mesh comprises predicting an inverse depth for each vertex of the intermediate 3D mesh, the inverse depth correlated with an object scale of the real-world object, wherein predicting the inverse depth for each vertex comprises computing the inverse depth of a vertex based on an estimated inverse depth of a parent skeleton joint and a per-vertex inverse-depth offset predicted by the first machine learning model; and
un-projecting a vertex to a three-dimensional (3D) position using a focal length of a camera model and the predicted inverse depth.
5 . The method of claim 1 , further comprising training the first machine learning model using training data comprising a plurality of training images depicting real-world objects and corresponding ground-truth denseposes and training 3D meshes corresponding to the training real-world objects, each ground-truth densepose representing coordinates of respective 3D meshes of the training real-world objects, wherein training the first machine learning model comprises applying weak supervision for vertex depths when direct mesh supervision is unavailable.
6 . The method of claim 5 , further comprising training the first machine learning model by computing a plurality of losses.
7 . The method of claim 6 , wherein a first loss of the plurality of losses comprises a cross-entropy loss, further comprising:
obtaining a first training image of the plurality of training images and a first ground-truth densepose corresponding to a first training real-world object depicted in the first training image;
computing barycentric coordinates of an individual triangle of a training 3D mesh corresponding to a particular pixel of the first training real-world object using the first ground-truth densepose;
estimating an individual tensor of heatmaps corresponding to the first training real-world object by applying the first machine learning model to the first training image;
generating predicted posterior coordinates using a selection of a group of coordinates of the individual tensor of heatmaps corresponding to the particular pixel of the first training real-world object;
computing a deviation comprising the cross-entropy loss between the predicted posterior coordinates and the barycentric coordinates; and
updating parameters of the first machine learning model based on the deviation.
8 . The method of claim 7 , wherein each heatmap of the individual tensor of heatmaps represents a respective group of possible coordinates for an individual vertex of an individual triangle of an individual mesh corresponding to the first training real-world object, further comprising updating scores associated with the respective group of possible coordinates for the individual vertex based on the cross-entropy loss.
9 . The method of claim 8 , wherein a second loss of the plurality of losses comprises a vertex localization loss, further comprising:
integrating scores of the group of possible coordinates for each vertex of the individual mesh corresponding to the first training real-world object to generate estimated coordinates for each vertex of the individual mesh;
computing a deviation between a particular vertex of the estimated coordinates of the individual mesh and a corresponding vertex of a training 3D mesh of the first training real-world object; and
computing a first portion of the second loss based on the computed deviation between the particular vertex and the corresponding vertex.
10 . The method of claim 9 , further comprising:
computing a first edge based on a difference between a first pair of the estimated coordinates of a particular triangle of the individual mesh;
computing a second edge based on a difference between a second pair of coordinates of a corresponding triangle of the training 3D mesh; and
computing a second portion of the second loss based on a deviation between the first edge and the second edge.
11 . The method of claim 7 , wherein a second loss of the plurality of losses comprises a part segmentation loss, further comprising:
identifying a plurality of parts of the first training real-world object, each of the plurality of parts comprising a respective collection of triangles;
selecting a particular part of the plurality of parts;
retrieving a group of coordinates of the individual tensor of heatmaps corresponding to the particular part;
computing a sum of scores associated with the retrieved group of coordinates;
comparing the sum of the scores to a threshold; and
updating parameters of the first machine learning model based on a result of comparing the sum of the scores to the threshold to maximize the sum of the scores for the retrieved group of coordinates.
12 . The method of claim 7 , wherein a second loss of the plurality of losses comprises a spatial consistency loss, further comprising:
obtaining a ground-truth pixel value of the first training image for a given set of barycentric coordinates corresponding to the individual triangle of the training 3D mesh;
computing an estimated pixel value based on the predicted posterior coordinates; and
computing the spatial consistency loss based on a deviation between the ground-truth pixel value and the estimated pixel value.
13 . The method of claim 7 , further comprising:
applying a second machine learning model to the portion of the image to estimate offsets between vertices of the 3D mesh.
14 . The method of claim 13 , further comprising training the second machine learning model based on an offset loss by:
computing a first offset between first and second vertices of the training 3D mesh;
computing a second offset between first and second vertices of an individual mesh generated using the individual tensor of heatmaps; and
computing the offset loss based on a deviation between the first and second offsets.
15 . The method of claim 13 , wherein the second machine learning model estimates a depth of the 3D mesh.
16 . The method of claim 1 , wherein the subset of the tensor of heatmaps represents a set of vertices of the plurality of triangles each associated with a score that transgresses a threshold score, further comprising:
predicting a continuous UV field in an interior of the first triangle by interpolating between 2D vertex positions of the first triangle using barycentric coordinates, wherein the continuous UV field comprises a weighted combination of vertex UV values.
17 . The method of claim 1 , wherein generating the 3D mesh comprises:
adopting a perspective camera model where each vertex of the 3D mesh lies on a ray crossing an image plane at a two-dimensional (2D) position predicted from the tensor of 2D heatmaps; and
estimating a depth of each vertex along the ray.
18 . A system comprising:
at least one processor; and
at least one memory component having instructions stored thereon that, when executed by the at least one processor, cause the at least one processor to perform operations comprising:
receiving an image that includes a depiction of a real-world object;
predicting a tensor of two-dimensional (2D) heatmaps by applying a first machine learning model to a portion of the image, each 2D heatmap of the tensor of 2D heatmaps exclusively representing possible positions for a corresponding single vertex of an individual triangle of a plurality of triangles of a three-dimensional (3D) mesh corresponding to the real-world object, the first machine learning model trained to predict a part of the 3D mesh that is not visible within the image including at least one of occluded vertices, self-occluded vertices, or truncated vertices, a first heatmap of the tensor of heatmaps representing a first group of possible coordinates exclusively for a single first vertex of a first triangle of the plurality of triangles, a second heatmap of the tensor of heatmaps representing a second group of possible coordinates exclusively for a single second vertex of the same first triangle as the first heatmap;
selecting a subset of the tensor of heatmaps from the tensor of heatmaps, the subset of the tensor of heatmaps comprising a first coordinate of the first group of possible coordinates and a second coordinate of the second group of possible coordinates; and
generating the 3D mesh based on the selected subset of the tensor of heatmaps, the 3D mesh being combined with an inpainted mesh to form a reconstructed 3D mesh.
19 . The system of claim 18 , the operations further comprising:
localizing vertices of the 3D mesh in two dimensions by applying a differentiable soft-argmax operation to the tensor of 2D heatmaps to obtain a two-dimensional (2D) position for each vertex.
20 . A non-transitory computer-readable storage medium having stored thereon instructions that, when executed by at least one processor, cause the at least one processor to perform operations comprising:
receiving an image that includes a depiction of a real-world object;
predicting a tensor of two-dimensional (2D) heatmaps by applying a first machine learning model to a portion of the image, each 2D heatmap of the tensor of 2D heatmaps exclusively representing possible positions for a corresponding single vertex of an individual triangle of a plurality of triangles of a three-dimensional (3D) mesh corresponding to the real-world object, the first machine learning model trained to predict a part of the 3D mesh that is not visible within the image including at least one of occluded vertices, self-occluded vertices, or truncated vertices, a first heatmap of the tensor of heatmaps representing a first group of possible coordinates exclusively for a single first vertex of a first triangle of the plurality of triangles, a second heatmap of the tensor of heatmaps representing a second group of possible coordinates exclusively for a single second vertex of the same first triangle as the first heatmap;
selecting a subset of the tensor of heatmaps from the tensor of heatmaps, the subset of the tensor of heatmaps comprising a first coordinate of the first group of possible coordinates and a second coordinate of the second group of possible coordinates; and
generating the 3D mesh based on the selected subset of the tensor of heatmaps, the 3D mesh being combined with an inpainted mesh to form a reconstructed 3D mesh.