System and method for segment-aware semantic segmentation

A computer-implemented system comprising:

• at least one memory;
• at least one processor communicatively coupled to the memory;
• and
• a segment-aware semantic segmentation unit, operated by the at least one processor, configured to:
• receive an image, a ground truth annotation mask and a segment ID mask, wherein the image comprises a plurality of pixels, the ground truth annotation mask comprises for each pixel a label of a correct class and the segment ID mask comprises, for each pixel, a segment identifier identifying a segment to which the pixel belongs,
• compute a segment-aware value using the image, the ground truth annotation mask and the segment ID mask, and
• use the segment-aware value to perform any of:
• train a machine learning model for semantic segmentation, validate the machine learning model for semantic segmentation and test the machine learning model for semantic segmentation,
• wherein the segment-aware semantic segmentation unit further comprises
• a segment-aware function unit configured to compute a segment-aware value,
• a machine learning model for semantic segmentation,
• a training unit configured to train the machine learning model for semantic segmentation using the segment-aware value,
• a validation unit configured to provide a metric indicative of the quality of the machine learning model for semantic segmentation using the segment-aware value, and
• an inference unit configured to use the machine learning model for semantic segmentation and provide a segment-aware segmentation result containing for each pixel a predicted class, and
• wherein the segment-aware function unit further comprises a segment-aware loss module configured to use the segment ID map to compute a segment-aware loss value indicative of an ability of the machine learning model for semantic segmentation to correctly classify all segments in the image and a segment-aware metric module configured to use the segment ID map to compute a segment-aware metric value indicative of a quality of the machine learning model for semantic segmentation, and
• wherein the segment-aware loss module is configured to
• compute for each one of the plurality of pixels a per-pixel loss value between the ground truth annotation mask, and a per class probability map containing for each pixel a probability of belonging to each class, and
• compute a per-segment loss value for each one of the plurality of segments in the segment ID map using a plurality of per-pixel loss values of pixels associated with the one of the plurality of segments.

A computer-implemented method comprising:

• receiving an image, a ground truth annotation mask and a segment ID mask, wherein the image comprises a plurality of pixels, the ground truth annotation mask comprises for each pixel a label of a correct class, and the segment ID mask comprises for each pixel a segment identifier identifying a segment to which the pixel belongs;
• feeding the image to a machine learning model for semantic segmentation and computing an outcome;
• computing a per-segment value for each one of the plurality of segments from the information in the outcome and in the ground truth annotation mask related to pixels associated with the one of the plurality of segments in the segment ID mask;
• calculating a segment-aware value from a plurality of per-segment values;
• and
• using the segment-aware value for validating the machine learning model for semantic segmentation and testing the machine learning model for semantic segmentation,
• wherein the step of computing a per-segment value further comprises
• computing for each one of the plurality of pixels a per-pixel metric value using the ground truth annotation mask and the outcome, wherein the outcome is a segmentation result containing for each pixel a label of a predicted class, and
• computing a per-segment value for each one of the plurality of segments in the segment ID mask using a plurality of per-pixel metric values of pixels associated with the one of the plurality of segments thereby creating a per-segment metric,
• wherein computing a per-segment metric further comprises counting for each segment a number of correctly classified pixels,
• wherein a correctly classified pixel is a pixel having a label of a predicted class in the segmentation result identical to the label in the ground truth annotation mask, and
• wherein the computing a segment-aware metric further comprises computing a ratio between the number of segments for which at least a predetermined number of pixels were classified correctly, and a total number of segments present in the ground truth annotation mask.