Self-supervised learning model for road feature detection in driving environments
Techniques are described for self-supervised training of machine-learned models for detecting road features based on images of a driving environment. The model may be implemented as a neural network having two output heads: a first semantic segmentation head trained to output a likelihood of a particular road feature at various locations in the image, and a second geometric shape trained to output the parameters of a predicted road feature at the various locations. During training, the segmentation head of the model may be trained using labeled data associated with road features (e.g., driveways, lane geometry, etc.). The geometric shape head may be trained based on training data masked by the output of the segmentation head. In some examples, a Voronoi diagram may be generated to expand the labeled training data, and the segmentation output mask may be applied to the Voronoi diagram as the masked training data for the geometric shape head.
1 . A system for training a self-supervised machine-learned model comprising:
one or more processors; and
one or more non-transitory computer-readable media storing computer-executable instructions that, when executed, cause the one or more processors to perform operations comprising:
receiving top-down image data associated with a driving environment representing the driving environment from a top-down perspective;
receiving training data associated with a road feature type, the training data associated with the top-down image data;
providing the top-down image data as input to a machine-learned model comprising:
a first subset of output layers configured to output a segmentation output indicating a likelihood of a road feature of the road feature type being present at a location in the top-down image data; and
a second subset of output layers configured to output a geometric shape output indicating dimensions of a predicted road feature of the road feature type at the location in the top-down image data;
determining, based at least in part on the segmentation output, a first prediction of a road feature being present at a first location in the top-down image data;
training the machine-learned model based at least in part on a comparison between the geometric shape output and the training data at the first location;
determining, based at least in part on the segmentation output, a second prediction of an absence of a road feature at a second location in the top-down image data; and
refraining, based at least in part on the second prediction, from training the machine-learned model based on the geometric shape output at the second location.
2 . The system of claim 1 , the operations further comprising:
receiving labeled image data of the driving environment;
generating a Voronoi diagram based on the labeled image data; and
generating the training data, based at least in part on the Voronoi diagram.
3 . The system of claim 2 , wherein generating the training data comprises:
determining a mask based at least in part on the segmentation output of the machine-learned model; and
applying the mask to the Voronoi diagram.
4 . The system of claim 1 , wherein the road feature type comprises at least one of:
a driveway road feature type;
a parking lot road feature type;
an alleyway road feature type; or
a driving lane road feature type.
5 . The system of claim 1 , the operations further comprising:
generating map data based at least in part on the machine-learned model; and
transmitting the map data to an autonomous vehicle configured to use the map data to control operation of the autonomous vehicle in the driving environment.
6 . A method comprising:
receiving labeled image data of a driving environment;
generating, based on the labeled image data, a data structure representative of a Voronoi diagram;
generating training data, based at least in part on the data structure;
receiving data associated with the driving environment;
providing the data as input to a machine-learned model;
receiving, from a first portion of the machine-learned model, a segmentation output indicating a likelihood of a road feature of a road feature type being present at a first location;
receiving, from a second portion of the machine-learned model, a geometric shape output representing a predicted road feature of the road feature type associated with the first location;
modifying a set of parameters associated with the machine-learned model based at least in part on a comparison between the segmentation output and the training data;
based at least in part on a likelihood indicated by the segmentation output meeting or exceeding a threshold likelihood modifying the set of parameters associated with the machine-learned model based at least in part on a comparison between the geometric shape output and the training data.
7 . The method of claim 6 , further comprising:
refraining from modifying the machine-learned model based at least in part on a portion of the segmentation output having no indication of being associated with the road feature type.
8 . The method of claim 6 , further comprising:
determining a mask based at least in part on the segmentation output of the machine-learned model; and
applying the mask to the Voronoi diagram to generate masked data,
wherein modifying the machine-learned model is based at least in part on the masked data.
9 . The method of claim 6 , wherein the machine-learned model comprises at least one of a convolutional neural network (CNN) or a vision transformer model, and wherein the machine-learned model is configured to output, for each particular location of a plurality of locations in the data:
a first subset of output layers corresponding to the segmentation output at the particular location; and
a second subset of output layers representing dimensions of a road feature of the road feature type predicted at the particular location.
10 . The method of claim 6 , wherein the road feature type comprises at least one of:
a driveway road feature type,
a parking lot road feature type,
an alleyway road feature type, or
a driving lane road feature type.
11 . The method of claim 6 , further comprising generating the data, wherein generating the data comprises:
receiving vehicle sensor data captured by a vehicle traversing the driving environment; and
generating, based at least in part on the vehicle sensor data, a top-down image of the driving environment.
12 . The method of claim 6 , further comprising:
generating map data based at least in part on the machine-learned model; and
transmitting the map data to an autonomous vehicle configured to use the map data to control operations of the autonomous vehicle in the driving environment.
13 . One or more non-transitory computer-readable media storing instructions executable by one or more processors, wherein the instructions, when executed, cause the one or more processors to perform operations comprising:
receiving data associated with a driving environment;
providing the data as input to a machine-learned model;
receiving, from a first portion of the machine-learned model, a segmentation output indicating a likelihood of a road feature of a road feature type being present at a first location;
receiving, from a second portion of the machine-learned model, a geometric shape output representing a predicted road feature of the road feature type associated with the first location; and
modifying a set of parameters associated with the machine-learned model based at least in part on a comparison between the segmentation output and training data, wherein modifying the set of parameters comprises:
when the segmentation output meets or exceeds a threshold likelihood, modifying the set of parameters associated with the machine-learned model based at least in part on a comparison between the geometric shape output and the training data; and
when the segmentation output does not meet or exceed the threshold likelihood, refraining from modifying the set of parameters.
14 . The one or more non-transitory computer-readable media of claim 13 , the operations further comprising:
receiving labeled image data of the driving environment;
generating a Voronoi diagram based on the labeled image data; and
generating the training data, based at least in part on the Voronoi diagram.
15 . The one or more non-transitory computer-readable media of claim 14 , the operations further comprising:
determining a mask based at least in part on the segmentation output of the machine-learned model; and
applying the mask to the Voronoi diagram to generate masked data,
wherein modifying the machine-learned model is based at least in part on the masked data.
16 . The one or more non-transitory computer-readable media of claim 13 , wherein the machine-learned model comprises at least one of a convolutional neural network (CNN) or a vision transformer model, and wherein the machine-learned model is configured to output, for each particular location of a plurality of locations in the data:
a first subset of output layers corresponding to the segmentation output at the particular location; and
a second subset of output layers representing dimensions of a road feature of the road feature type predicted at the particular location.
17 . The one or more non-transitory computer-readable media of claim 13 , wherein the road feature type comprises at least one of:
a driveway road feature type,
a parking lot road feature type,
an alleyway road feature type, or
a driving lane road feature type.
18 . The one or more non-transitory computer-readable media of claim 13 , the operations further comprising generating the data, wherein generating the data comprises:
receiving vehicle sensor data captured by a vehicle traversing the driving environment; and generating, based at least in part on the vehicle sensor data, a top-down image of the driving environment.