google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Utility functions for light."""
17
import tensorflow as tf
18
from osf import run_nerf_helpers
19

20

21
def compute_indirect_radiance(intersect, results, light_rgb, white_bkgd):
22
  """Computes per-primary sample radiance from per-object secondary rays.
23

24
  Args:
25
    intersect:
26
    results: [RS, SO, K]
27
    light_rgb:
28
    white_bkgd:
29

30
  Returns:
31
    indirect_radiance: [R, S, 3] tf.float32. Radiance along each indirect ray.
32
  """
33
  num_rays = tf.shape(intersect['normalized_rgb'])[0]
34

35
  # Compose the combined object outputs into the final rendered result.
36
  outputs = run_nerf_helpers.compose_outputs(  # [RS, K]
37
      results=results,  # [RS, SO, K]
38
      light_rgb=light_rgb,
39
      white_bkgd=white_bkgd)
40

41
  # Reshape the rgb values.
42
  indirect_radiance = outputs['rgb_map']  # [RS, 3]
43
  indirect_radiance = tf.reshape(  # [R, S, 3]
44
      indirect_radiance, [num_rays, -1, 3])
45
  return indirect_radiance
46