# Path: ibrnet/render_ray.py
# def render_rays(ray_batch,
#                 model,
#                 featmaps,
#                 projector,
#                 N_samples,
#                 inv_uniform=False,
#                 N_importance=0,
#                 det=False,
#                 white_bkgd=False):
#     '''
#     :param ray_batch: {'ray_o': [N_rays, 3] , 'ray_d': [N_rays, 3], 'view_dir': [N_rays, 2]}
#     :param model:  {'net_coarse':  , 'net_fine': }
#     :param N_samples: samples along each ray (for both coarse and fine model)
#     :param inv_uniform: if True, uniformly sample inverse depth for coarse model
#     :param N_importance: additional samples along each ray produced by importance sampling (for fine model)
#     :param det: if True, will deterministicly sample depths
#     :return: {'outputs_coarse': {}, 'outputs_fine': {}}
#     '''
# 
#     ret = {'outputs_coarse': None,
#            'outputs_fine': None}
# 
#     # pts: [N_rays, N_samples, 3]
#     # z_vals: [N_rays, N_samples]
#     pts, z_vals = sample_along_camera_ray(ray_o=ray_batch['ray_o'],
#                                           ray_d=ray_batch['ray_d'],
#                                           depth_range=ray_batch['depth_range'],
#                                           N_samples=N_samples, inv_uniform=inv_uniform, det=det)
#     N_rays, N_samples = pts.shape[:2]
# 
#     rgb_feat, ray_diff, mask = projector.compute(pts, ray_batch['camera'],
#                                                  ray_batch['src_rgbs'],
#                                                  ray_batch['src_cameras'],
#                                                  featmaps=featmaps[0])  # [N_rays, N_samples, N_views, x]
#     pixel_mask = mask[..., 0].sum(dim=2) > 1   # [N_rays, N_samples], should at least have 2 observations
#     raw_coarse = model.net_coarse(rgb_feat, ray_diff, mask)   # [N_rays, N_samples, 4]
#     outputs_coarse = raw2outputs(raw_coarse, z_vals, pixel_mask,
#                                  white_bkgd=white_bkgd)
#     ret['outputs_coarse'] = outputs_coarse
# 
#     if N_importance > 0:
#         assert model.net_fine is not None
#         # detach since we would like to decouple the coarse and fine networks
#         weights = outputs_coarse['weights'].clone().detach()            # [N_rays, N_samples]
#         if inv_uniform:
#             inv_z_vals = 1. / z_vals
#             inv_z_vals_mid = .5 * (inv_z_vals[:, 1:] + inv_z_vals[:, :-1])   # [N_rays, N_samples-1]
#             weights = weights[:, 1:-1]      # [N_rays, N_samples-2]
#             inv_z_vals = sample_pdf(bins=torch.flip(inv_z_vals_mid, dims=[1]),
#                                     weights=torch.flip(weights, dims=[1]),
#                                     N_samples=N_importance, det=det)  # [N_rays, N_importance]
#             z_samples = 1. / inv_z_vals
#         else:
#             # take mid-points of depth samples
#             z_vals_mid = .5 * (z_vals[:, 1:] + z_vals[:, :-1])   # [N_rays, N_samples-1]
#             weights = weights[:, 1:-1]      # [N_rays, N_samples-2]
#             z_samples = sample_pdf(bins=z_vals_mid, weights=weights,
#                                    N_samples=N_importance, det=det)  # [N_rays, N_importance]
# 
#         z_vals = torch.cat((z_vals, z_samples), dim=-1)  # [N_rays, N_samples + N_importance]
# 
#         # samples are sorted with increasing depth
#         z_vals, _ = torch.sort(z_vals, dim=-1)
#         N_total_samples = N_samples + N_importance
# 
#         viewdirs = ray_batch['ray_d'].unsqueeze(1).repeat(1, N_total_samples, 1)
#         ray_o = ray_batch['ray_o'].unsqueeze(1).repeat(1, N_total_samples, 1)
#         pts = z_vals.unsqueeze(2) * viewdirs + ray_o  # [N_rays, N_samples + N_importance, 3]
# 
#         rgb_feat_sampled, ray_diff, mask = projector.compute(pts, ray_batch['camera'],
#                                                              ray_batch['src_rgbs'],
#                                                              ray_batch['src_cameras'],
#                                                              featmaps=featmaps[1])
# 
#         pixel_mask = mask[..., 0].sum(dim=2) > 1  # [N_rays, N_samples]. should at least have 2 observations
#         raw_fine = model.net_fine(rgb_feat_sampled, ray_diff, mask)
#         outputs_fine = raw2outputs(raw_fine, z_vals, pixel_mask,
#                                    white_bkgd=white_bkgd)
#         ret['outputs_fine'] = outputs_fine
# 
#     return ret


import torch
from collections import OrderedDict
from ibrnet.render_ray import render_rays
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.




def render_single_image(ray_sampler,
                        ray_batch,
                        model,
                        projector,
                        chunk_size,
                        N_samples,
                        inv_uniform=False,
                        N_importance=0,
                        det=False,
                        white_bkgd=False,
                        render_stride=1,

# Path: ibrnet/render_ray.py
# def render_rays(ray_batch,
#                 model,
#                 featmaps,
#                 projector,
#                 N_samples,
#                 inv_uniform=False,
#                 N_importance=0,
#                 det=False,
#                 white_bkgd=False):
#     '''
#     :param ray_batch: {'ray_o': [N_rays, 3] , 'ray_d': [N_rays, 3], 'view_dir': [N_rays, 2]}
#     :param model:  {'net_coarse':  , 'net_fine': }
#     :param N_samples: samples along each ray (for both coarse and fine model)
#     :param inv_uniform: if True, uniformly sample inverse depth for coarse model
#     :param N_importance: additional samples along each ray produced by importance sampling (for fine model)
#     :param det: if True, will deterministicly sample depths
#     :return: {'outputs_coarse': {}, 'outputs_fine': {}}
#     '''
# 
#     ret = {'outputs_coarse': None,
#            'outputs_fine': None}
# 
#     # pts: [N_rays, N_samples, 3]
#     # z_vals: [N_rays, N_samples]
#     pts, z_vals = sample_along_camera_ray(ray_o=ray_batch['ray_o'],
#                                           ray_d=ray_batch['ray_d'],
#                                           depth_range=ray_batch['depth_range'],
#                                           N_samples=N_samples, inv_uniform=inv_uniform, det=det)
#     N_rays, N_samples = pts.shape[:2]
# 
#     rgb_feat, ray_diff, mask = projector.compute(pts, ray_batch['camera'],
#                                                  ray_batch['src_rgbs'],
#                                                  ray_batch['src_cameras'],
#                                                  featmaps=featmaps[0])  # [N_rays, N_samples, N_views, x]
#     pixel_mask = mask[..., 0].sum(dim=2) > 1   # [N_rays, N_samples], should at least have 2 observations
#     raw_coarse = model.net_coarse(rgb_feat, ray_diff, mask)   # [N_rays, N_samples, 4]
#     outputs_coarse = raw2outputs(raw_coarse, z_vals, pixel_mask,
#                                  white_bkgd=white_bkgd)
#     ret['outputs_coarse'] = outputs_coarse
# 
#     if N_importance > 0:
#         assert model.net_fine is not None
#         # detach since we would like to decouple the coarse and fine networks
#         weights = outputs_coarse['weights'].clone().detach()            # [N_rays, N_samples]
#         if inv_uniform:
#             inv_z_vals = 1. / z_vals
#             inv_z_vals_mid = .5 * (inv_z_vals[:, 1:] + inv_z_vals[:, :-1])   # [N_rays, N_samples-1]
#             weights = weights[:, 1:-1]      # [N_rays, N_samples-2]
#             inv_z_vals = sample_pdf(bins=torch.flip(inv_z_vals_mid, dims=[1]),
#                                     weights=torch.flip(weights, dims=[1]),
#                                     N_samples=N_importance, det=det)  # [N_rays, N_importance]
#             z_samples = 1. / inv_z_vals
#         else:
#             # take mid-points of depth samples
#             z_vals_mid = .5 * (z_vals[:, 1:] + z_vals[:, :-1])   # [N_rays, N_samples-1]
#             weights = weights[:, 1:-1]      # [N_rays, N_samples-2]
#             z_samples = sample_pdf(bins=z_vals_mid, weights=weights,
#                                    N_samples=N_importance, det=det)  # [N_rays, N_importance]
# 
#         z_vals = torch.cat((z_vals, z_samples), dim=-1)  # [N_rays, N_samples + N_importance]
# 
#         # samples are sorted with increasing depth
#         z_vals, _ = torch.sort(z_vals, dim=-1)
#         N_total_samples = N_samples + N_importance
# 
#         viewdirs = ray_batch['ray_d'].unsqueeze(1).repeat(1, N_total_samples, 1)
#         ray_o = ray_batch['ray_o'].unsqueeze(1).repeat(1, N_total_samples, 1)
#         pts = z_vals.unsqueeze(2) * viewdirs + ray_o  # [N_rays, N_samples + N_importance, 3]
# 
#         rgb_feat_sampled, ray_diff, mask = projector.compute(pts, ray_batch['camera'],
#                                                              ray_batch['src_rgbs'],
#                                                              ray_batch['src_cameras'],
#                                                              featmaps=featmaps[1])
# 
#         pixel_mask = mask[..., 0].sum(dim=2) > 1  # [N_rays, N_samples]. should at least have 2 observations
#         raw_fine = model.net_fine(rgb_feat_sampled, ray_diff, mask)
#         outputs_fine = raw2outputs(raw_fine, z_vals, pixel_mask,
#                                    white_bkgd=white_bkgd)
#         ret['outputs_fine'] = outputs_fine
# 
#     return ret


# Path: ibrnet/render_image.py
import torch
from collections import OrderedDict
from ibrnet.render_ray import render_rays

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.




def render_single_image(ray_sampler,
                        ray_batch,
                        model,
                        projector,
                        chunk_size,
                        N_samples,
                        inv_uniform=False,
                        N_importance=0,
                        det=False,
                        white_bkgd=False,
                        render_stride=1,

                        featmaps=None):