Scene re-lighting using direct shading control

1 . A method comprising:

obtaining an input image and a lighting direction indicator that describes a lighting direction;

generating, using a direct shading model, a direct shading map based on the input image and the lighting direction indicator; and

generating, using an image generation model, a shaded image depicting an object from the input image with shading consistent with the lighting direction based on the direct shading map, wherein generating the shaded image comprises performing a reverse diffusion process using the direct shading map as a control guidance.

2 . The method of claim 1 , wherein generating the direct shading map comprises:

encoding the input image to obtain image features; and

projecting the image features in a depth dimension to obtain a three dimensional (3D) feature grid, wherein the direct shading map is based on the 3D feature grid.

3 . The method of claim 2 , further comprising:

generating a depth map based on the input image, wherein the 3D feature grid is based on the depth map.

4 . The method of claim 2 , further comprising:

generating a density grid based on the 3D feature grid.

5 . The method of claim 4 , further comprising:

generating a shading input based on the density grid, wherein the shading input includes a shadow map, a normal map, an N-dot-L map, or any combination thereof, and wherein the direct shading map is generated based on the shading input.

6 . The method of claim 4 , further comprising:

generating a coarse shading map based on the density grid; and

refining the coarse shading map based on the input image to obtain the direct shading map.

7 . The method of claim 1 , wherein generating the shaded image comprises:

encoding, using a residual control encoder, the direct shading map to obtain lighting control information, wherein the shaded image is generated based on the lighting control information.

8 . The method of claim 7 , wherein generating the shaded image comprises:

obtaining a normal map, wherein the lighting control information is generated based on the normal map.

9 . The method of claim 7 , wherein:

the residual control encoder is trained based on an output of a residual control decoder that reconstructs an input of the residual control encoder.

10 . The method of claim 1 , wherein:

the image generation model is trained to generate images depicting an element of a training image with shading based on a training direct shading map.

11 . A method for training a machine learning model, comprising:

obtaining a training set including a training image and a direct shading map for the training image;

encoding, using a residual control encoder, the direct shading map to obtain an encoder output;

decoding, using a residual control decoder, the encoder output to obtain a reconstructed shading map; and

training, using the training set, an image generation model to generate images depicting an object from the training image with shading based on the direct shading map, wherein the image generation model is separate from the residual control encoder and the residual control decoder.

12 . The method of claim 11 , wherein obtaining the training set comprises:

generating, using a direct shading model, the direct shading map based on the training image.

13 . The method of claim 12 , further comprising:

training the direct shading model to generate the direct shading map based on the training image.

14 . The method of claim 12 , further comprising:

extracting a lighting direction from the training image, wherein the direct shading map is based on the lighting direction.

15 . The method of claim 11 , wherein training the image generation model comprises:

computing a reconstruction loss based on the reconstructed shading map and the direct shading map, wherein the image generation model is updated based on the reconstruction loss.

16 . An apparatus comprising:

at least one processor;

at least one memory storing instructions executable by the at least one processor;

a direct shading model comprising parameters stored in the at least one memory, wherein the direct shading model is configured to generate a direct shading map based on an input image and a lighting direction indicator that describes a lighting direction; and

an image generation model comprising parameters stored in the at least one memory, wherein the image generation model is configured to generate a shaded image depicting an object from the input image with shading consistent with the lighting direction based on the direct shading map by performing a reverse diffusion process using the direct shading map as control guidance.

17 . The apparatus of claim 16 , wherein:

the image generation model comprises a residual control encoder trained to generate lighting control information for an image generator.

18 . The apparatus of claim 16 , wherein:

the direct shading model comprises a density network trained to generate a density grid based on 3D features of the input image.

19 . The apparatus of claim 18 , wherein:

the direct shading model comprises a shading network trained to generate the direct shading map based on the density grid.