Systems and methods for training an image colorization model

A computing system comprising:

• one or more processors;
• a colorization model configured to receive a training input image and, in response to receiving the training input image, output a predicted color map that describes a predicted colorization for the training input image;
• one or more non-transitory computer-readable media that collectively store instructions that, when executed by the one or more processors, cause the computing system to perform operations, the operations comprising:
• inputting the training input image into the colorization model;
• receiving the predicted color map as an output of the colorization model;
• calculating a first color distance between a first pixel included in the predicted color map and a second pixel included in the predicted color map;
• calculating a second color distance between a third pixel included in a ground truth color map and a fourth pixel included in the ground truth color map, wherein the third pixel and fourth pixels included in the ground truth color map spatially correspond, respectively, with the first pixel and second pixel included in the predicted color map;
• evaluating a neighborhood color loss function that evaluates a difference between the first color distance and the second color distance;
• and
• adjusting parameters associated with the colorization model based on the neighborhood color loss function;
• and
• iteratively generating respective neighborhood color loss maps for respective additional first pixels of the predicted color map.

A method for training an image colorization model, the method comprising:

• inputting, by one or more computing devices, a training input image into a colorization model configured to receive the training input image and process the training input image to output a predicted color map that describes a predicted colorization for the training input image;
• receiving, by the one or more computing devices, the predicted color map as an output of the colorization model;
• calculating, by the one or more computing devices, a first color distance between a first pixel included in the predicted color map and a second pixel included in the predicted color map;
• calculating, by the one or more computing devices, a second color distance between a third pixel included in a ground truth color map and a fourth pixel included in the ground truth color map, wherein the third pixel and fourth pixel included in the ground truth color map spatially correspond, respectively, with the first pixel and second pixel included in the predicted color map;
• evaluating, by the one or more computing devices, a neighborhood color loss function that evaluates a difference between the first color distance and the second color distance;
• and
• adjusting, by the one or more computing devices, parameters associated with the colorization model based on the neighborhood color loss function;
• wherein:
• the colorization model comprises a multinomial prediction model and a refinement model, wherein the multinomial prediction model is configured to receive the training input image, and in response to receiving the training input image, output a multinomial color distribution that describes a plurality of colorizations for the training input image in a discretized color space, and wherein the refinement model is configured to receive the multinomial color distribution, and in response to receiving the multinomial color distribution, output the predicted color map;
• inputting the training input image into the colorization model comprises inputting the training input image into the multinomial prediction model;
• and
• the method further comprises:
• receiving the multinomial color distribution as an output of the multinomial prediction model;
• and
• inputting the multinomial color distribution into the refinement model.