Generating depth images for image data

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a machine learning model configured to generate a predicted depth image, comprising receiving data representing training samples that include a plurality of image pairs, each image pair includes a target image and a reference image both capturing a particular scene from different orientations; for each of the plurality of image pairs, generating a compressed cost volume for the image pair; providing the compressed cost volume as an input to the machine learning model; generating, using the machine learning model, output data representing a predicted disparity map for the compressed cost volume; and generating a total loss using the predicted disparity map for the compressed cost volume, the total loss includes a boundary loss, an occlusion loss, and a transfer loss; and updating the plurality of parameters of the machine learning model by minimizing the total losses.

What is claimed is: 
     
       1. A method for training a machine learning model configured to generate a prediction of a depth map for an input image, wherein the machine learning model includes a plurality of parameters to be updated during the training, the method comprising:
 receiving data representing training samples that comprise a plurality of image pairs, each image pair comprises a target image capturing a particular scene from a first orientation by a first image sensor and a reference image capturing the particular scene from a second orientation by a second image sensor; 
 for each of the plurality of image pairs:
 generating a compressed cost volume for the image pair, wherein the compressed cost volume comprises one or more cost values selected for each pixel of a reference image in the image pair; 
 providing the compressed cost volume as an input to the machine learning model; 
 generating, using the machine learning model, output data representing a predicted disparity map for the compressed cost volume; and 
 generating a total loss using the predicted disparity map for the compressed cost volume, the total loss comprising a boundary loss, an occlusion loss, and a transfer loss; and 
 
 updating the plurality of parameters of the machine learning model by minimizing the total losses. 
 
     
     
       2. The method of  claim 1 , wherein generating the compressed cost volume for the image pair comprises:
 generating a cost volume for the image pair, wherein the cost volume includes a plurality of cost values for each pixel of the reference image of the image pair at different disparity levels; 
 for each pixel of the reference image:
 generating a waveform for the pixel along an axis of disparity level based on the plurality of cost values for the pixel; and 
 determining a plurality of local minima of the plurality of cost values in the waveform; and 
 
 generating the compressed cost volume for the reference image by selecting one or more local minima of the plurality of local minima and their corresponding disparity levels. 
 
     
     
       3. The method of  claim 1 , further comprising generating the boundary loss for the image pair, the generating comprising:
 generating a boundary map for the image pair using a trained machine learning model, wherein the boundary map includes a probability for each pixel of the reference image indicating a likelihood of the pixel being a boundary pixel; and 
 generating a boundary loss for the image pair based on the predicted disparity map and the boundary map, comprising:
 generating depth gradients based on the predicted disparity map; 
 generating an enhanced boundary confidence map based on the boundary map; and 
 generating the boundary loss by multiplying the depth gradients and the enhanced boundary confidence map. 
 
 
     
     
       4. The method of  claim 1 , further comprising generating the occlusion loss for the image pair, the generating comprising:
 exchanging the target image and the reference image of the image pair to generate a swapped image pair; 
 generating, using the machine learning model, data representing a predicted second disparity map for the swapped image pair; and 
 generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map. 
 
     
     
       5. The method of  claim 4 , wherein generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map, further comprises:
 for each pixel in the reference image of the image pair:
 determining a disparity discrepancy based on disparity values for the pixel from the predicted disparity map and the predicted second disparity map; 
 determining whether the disparity discrepancy is less than or equal to a threshold value; 
 in response to determining that the disparity discrepancy is less than or equal to the threshold value, calculating a pixel-wise loss for the pixel to be aggregated to the occlusion loss; or 
 in response to determining that the disparity discrepancy is greater than the threshold value, setting the pixel-wise loss to be zero. 
 
 
     
     
       6. The method of  claim 1 , further comprising generating the transfer loss for the image pair, the generating comprising:
 processing the predicted disparity map for the image pair using a neural network to generate a first feature map; 
 processing a ground truth disparity map for another image pair using the neural network to generate a second feature map; 
 generating data representing a first Gram matrix for the first feature map; 
 generating data representing a second Gram matrix for the second feature map; and 
 generating the transfer loss for the image pair based on a sum of squared distances between the first Gram matrix and the second Gram matrix. 
 
     
     
       7. The method of  claim 1 , further comprising:
 deploying the trained machine learning model with the updated plurality of parameters on an edge device; and 
 generating a predicted disparity map for an input image pair by performing inference operations of the trained machine learning model for processing the input image pair. 
 
     
     
       8. A system comprising one or more computers and one or more storage devices storing instructions that when executed by one or more computers cause the one or more computers to perform respective operations for training a machine learning model configured to generate a prediction of a depth map for an input image, wherein the machine learning model includes a plurality of parameters to be updated during the training, the operations comprising:
 receiving data representing training samples that comprise a plurality of image pairs, each image pair comprises a target image capturing a particular scene from a first orientation by a first image sensor and a reference image capturing the particular scene from a second orientation by a second image sensor; 
 for each of the plurality of image pairs:
 generating a compressed cost volume for the image pair, wherein the compressed cost volume comprises one or more cost values selected for each pixel of a reference image in the image pair; 
 providing the compressed cost volume as an input to the machine learning model; 
 generating, using the machine learning model, output data representing a predicted disparity map for the compressed cost volume; and 
 generating a total loss using the predicted disparity map for the compressed cost volume, the total loss comprising a boundary loss, an occlusion loss, and a transfer loss; and 
 
 updating the plurality of parameters of the machine learning model by minimizing the total losses. 
 
     
     
       9. The system of  claim 8 , wherein generating the compressed cost volume for the image pair comprises:
 generating a cost volume for the image pair, wherein the cost volume includes a plurality of cost values for each pixel of the reference image of the image pair at different disparity levels; 
 for each pixel of the reference image:
 generating a waveform for the pixel along an axis of disparity level based on the plurality of cost values for the pixel; and 
 determining a plurality of local minima of the plurality of cost values in the waveform; and 
 
 generating the compressed cost volume for the reference image by selecting one or more local minima of the plurality of local minima and their corresponding disparity levels. 
 
     
     
       10. The system of  claim 8 , wherein the operations further comprise generating the boundary loss for the image pair, the generating comprising:
 generating a boundary map for the image pair using a trained machine learning model, wherein the boundary map includes a probability for each pixel of the reference image indicating a likelihood of the pixel being a boundary pixel; and 
 generating a boundary loss for the image pair based on the predicted disparity map and the boundary map, comprising:
 generating depth gradients based on the predicted disparity map; 
 generating an enhanced boundary confidence map based on the boundary map; and 
 generating the boundary loss by multiplying the depth gradients and the enhanced boundary confidence map. 
 
 
     
     
       11. The system of  claim 8 , wherein the operations further comprise generating the occlusion loss for the image pair, the generating comprising:
 exchanging the target image and the reference image of the image pair to generate a swapped image pair; 
 generating, using the machine learning model, data representing a predicted second disparity map for the swapped image pair; and 
 generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map. 
 
     
     
       12. The system of  claim 11 , wherein generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map, further comprises:
 for each pixel in the reference image of the image pair:
 determining a disparity discrepancy based on disparity values for the pixel from the predicted disparity map and the predicted second disparity map; 
 determining whether the disparity discrepancy is less than or equal to a threshold value; 
 in response to determining that the disparity discrepancy is less than or equal to the threshold value, calculating a pixel-wise loss for the pixel to be aggregated to the occlusion loss; or 
 in response to determining that the disparity discrepancy is greater than the threshold value, setting the pixel-wise loss to be zero. 
 
 
     
     
       13. The system of  claim 8 , wherein the operations further comprise generating the transfer loss for the image pair, the generating comprising:
 processing the predicted disparity map for the image pair using a neural network to generate a first feature map; 
 processing a ground truth disparity map for another image pair using the neural network to generate a second feature map; 
 generating data representing a first Gram matrix for the first feature map; 
 generating data representing a second Gram matrix for the second feature map; and 
 generating the transfer loss for the image pair based on a sum of squared distances between the first Gram matrix and the second Gram matrix. 
 
     
     
       14. The system of  claim 8 , wherein the operations further comprise:
 deploying the trained machine learning model with the updated plurality of parameters on an edge device; and 
 generating a predicted disparity map for an input image pair by performing inference operations of the trained machine learning model for processing the input image pair. 
 
     
     
       15. One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform respective operations for training a machine learning model configured to generate a prediction of a depth map for an input image, wherein the machine learning model includes a plurality of parameters to be updated during the training, the operations comprising:
 receiving data representing training samples that comprise a plurality of image pairs, each image pair comprises a target image capturing a particular scene from a first orientation by a first image sensor and a reference image capturing the particular scene from a second orientation by a second image sensor; 
 for each of the plurality of image pairs:
 generating a compressed cost volume for the image pair, wherein the compressed cost volume comprises one or more cost values selected for each pixel of a reference image in the image pair; 
 providing the compressed cost volume as an input to the machine learning model; 
 generating, using the machine learning model, output data representing a predicted disparity map for the compressed cost volume; and 
 generating a total loss using the predicted disparity map for the compressed cost volume, the total loss comprising a boundary loss, an occlusion loss, and a transfer loss; and 
 
 updating the plurality of parameters of the machine learning model by minimizing the total losses. 
 
     
     
       16. The one or more non-transitory computer-readable storage media of  claim 15 , wherein generating the compressed cost volume for the image pair comprises:
 generating a cost volume for the image pair, wherein the cost volume includes a plurality of cost values for each pixel of the reference image of the image pair at different disparity levels; 
 for each pixel of the reference image:
 generating a waveform for the pixel along an axis of disparity level based on the plurality of cost values for the pixel; and 
 determining a plurality of local minima of the plurality of cost values in the waveform; and 
 
 generating the compressed cost volume for the reference image by selecting one or more local minima of the plurality of local minima and their corresponding disparity levels. 
 
     
     
       17. The one or more non-transitory computer-readable storage media of  claim 15 , wherein the operations further comprise generating the boundary loss for the image pair, the generating comprising:
 generating a boundary map for the image pair using a trained machine learning model, wherein the boundary map includes a probability for each pixel of the reference image indicating a likelihood of the pixel being a boundary pixel; and 
 generating a boundary loss for the image pair based on the predicted disparity map and the boundary map, comprising:
 generating depth gradients based on the predicted disparity map; 
 generating an enhanced boundary confidence map based on the boundary map; and 
 generating the boundary loss by multiplying the depth gradients and the enhanced boundary confidence map. 
 
 
     
     
       18. The one or more non-transitory computer-readable storage media of  claim 15 , wherein the operations further comprise generating the occlusion loss for the image pair, the generating comprising:
 exchanging the target image and the reference image of the image pair to generate a swapped image pair; 
 generating, using the machine learning model, data representing a predicted second disparity map for the swapped image pair; and 
 generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map. 
 
     
     
       19. The one or more non-transitory computer-readable storage media of  claim 18 , wherein generating the occlusion loss for the image pair based on the predicted disparity map and the predicted second disparity map, further comprises:
 for each pixel in the reference image of the image pair:
 determining a disparity discrepancy based on disparity values for the pixel from the predicted disparity map and the predicted second disparity map; 
 determining whether the disparity discrepancy is less than or equal to a threshold value; 
 in response to determining that the disparity discrepancy is less than or equal to the threshold value, calculating a pixel-wise loss for the pixel to be aggregated to the occlusion loss; or 
 in response to determining that the disparity discrepancy is greater than the threshold value, setting the pixel-wise loss to be zero. 
 
 
     
     
       20. The one or more non-transitory computer-readable storage media of  claim 15 , wherein the operations further comprise generating the transfer loss for the image pair, the generating comprising:
 processing the predicted disparity map for the image pair using a neural network to generate a first feature map; 
 processing a ground truth disparity map for another image pair using the neural network to generate a second feature map; 
 generating data representing a first Gram matrix for the first feature map; 
 generating data representing a second Gram matrix for the second feature map; and 
 generating the transfer loss for the image pair based on a sum of squared distances between the first Gram matrix and the second Gram matrix.