text: import os, sys
  import tensorflow as tf
  import numpy as np
  import imageio 
  import json
  
  
  # Misc utils
  
  img2mse = lambda x, y : tf.reduce_mean(tf.square(x - y))
  mse2psnr = lambda x : -10.*tf.log(x)/tf.log(10.)
  to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
  
      
  # Positional encoding
  
  class Embedder:
      
      def __init__(self, **kwargs):
          
          self.kwargs = kwargs
          self.create_embedding_fn()
          
          
      def create_embedding_fn(self):
          
          embed_fns = []
          d = self.kwargs['input_dims']
          out_dim = 0
          if self.kwargs['include_input']:
              embed_fns.append(lambda x : x)
              out_dim += d
              
          max_freq = self.kwargs['max_freq_log2']
          N_freqs = self.kwargs['num_freqs']
          
          if self.kwargs['log_sampling']:
              freq_bands = 2.**tf.linspace(0., max_freq, N_freqs)
          else:
              freq_bands = tf.linspace(2.**0., 2.**max_freq, N_freqs)
              
          for freq in freq_bands:
              for p_fn in self.kwargs['periodic_fns']:
                  embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
                  out_dim += d
                      
          self.embed_fns = embed_fns
          self.out_dim = out_dim
          
      
      def embed(self, inputs):
          return tf.concat([fn(inputs) for fn in self.embed_fns], -1)
              
          
          
  
  def get_embedder(multires, i=0):
      
      if i == -1:
          return tf.identity, 3
      
      embed_kwargs = {
                  'include_input' : True,
                  'input_dims' : 3,
                  'max_freq_log2' : multires-1,
                  'num_freqs' : multires,
                  'log_sampling' : True,
                  'periodic_fns' : [tf.math.sin, tf.math.cos],
      }
      
      embedder_obj = Embedder(**embed_kwargs)
      embed = lambda x, eo=embedder_obj : eo.embed(x)
      return embed, embedder_obj.out_dim
  
  
  
  # Model architecture
  
  def init_nerf_model(D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False):
      
      relu = tf.keras.layers.ReLU()    
      dense = lambda W, act=relu : tf.keras.layers.Dense(W, activation=act)
      
      print('MODEL', input_ch, input_ch_views, type(input_ch), type(input_ch_views), use_viewdirs)
      input_ch = int(input_ch)
      input_ch_views = int(input_ch_views)
      
      inputs = tf.keras.Input(shape=(input_ch + input_ch_views)) 
      inputs_pts, inputs_views = tf.split(inputs, [input_ch, input_ch_views], -1)
      inputs_pts.set_shape([None, input_ch])
      inputs_views.set_shape([None, input_ch_views])
      
      print(inputs.shape, inputs_pts.shape, inputs_views.shape)
      outputs = inputs_pts
      for i in range(D):
          outputs = dense(W)(outputs)
          if i in skips:
              outputs = tf.concat([inputs_pts, outputs], -1)
                          
      if use_viewdirs:
          alpha_out = dense(1, act=None)(outputs)
          bottleneck = dense(256, act=None)(outputs)
          inputs_viewdirs = tf.concat([bottleneck, inputs_views], -1)  # concat viewdirs
          for i in range(D//2):
              outputs = dense(W//2)(inputs_viewdirs)
          outputs = dense(3, act=None)(outputs)
          outputs = tf.concat([outputs, alpha_out], -1)
      else:
          outputs = dense(output_ch, act=None)(outputs)
      
      model = tf.keras.Model(inputs=inputs, outputs=outputs)
      return model
  
  
  # Ray helpers
  
  def get_rays(H, W, focal, c2w):
      i, j = tf.meshgrid(tf.range(W, dtype=tf.float32), tf.range(H, dtype=tf.float32), indexing='xy')
      dirs = tf.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -tf.ones_like(i)], -1)
      rays_d = tf.reduce_sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)
      rays_o = tf.broadcast_to(c2w[:3,-1], tf.shape(rays_d))
      return rays_o, rays_d
  
  
  def get_rays_np(H, W, focal, c2w):
      i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32), indexing='xy')
      dirs = np.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -np.ones_like(i)], -1)
      rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)
      rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))
      return rays_o, rays_d
  
  
  
  def ndc_rays(H, W, focal, near, rays_o, rays_d):
      
      # Shift ray origins to near plane
      t = -(near + rays_o[...,2]) / rays_d[...,2]
      rays_o = rays_o + t[...,None] * rays_d
      
      # Projection
      o0 = -1./(W/(2.*focal)) * rays_o[...,0] / rays_o[...,2]
      o1 = -1./(H/(2.*focal)) * rays_o[...,1] / rays_o[...,2]
      o2 = 1. + 2. * near / rays_o[...,2]
  
      d0 = -1./(W/(2.*focal)) * (rays_d[...,0]/rays_d[...,2] - rays_o[...,0]/rays_o[...,2])
      d1 = -1./(H/(2.*focal)) * (rays_d[...,1]/rays_d[...,2] - rays_o[...,1]/rays_o[...,2])
      d2 = -2. * near / rays_o[...,2]
      
      rays_o = tf.stack([o0,o1,o2], -1)
      rays_d = tf.stack([d0,d1,d2], -1)
      
      return rays_o, rays_d
      
      
  # Hierarchical sampling helper    
  
  def sample_pdf(bins, weights, N_samples, det=False, pytest=False):
      # Get pdf
      weights += 1e-5 # prevent nans
      pdf = weights / tf.reduce_sum(weights, -1, keepdims=True)
      cdf = tf.cumsum(pdf, -1)
      cdf = tf.concat([tf.zeros_like(cdf[...,:1]), cdf], -1)
  
      # Take uniform samples
      if det:
          u = tf.linspace(0., 1., N_samples)
          u = tf.broadcast_to(u, list(cdf.shape[:-1]) + [N_samples])
      else:
          u = tf.random.uniform(list(cdf.shape[:-1]) + [N_samples])
  
      # Pytest, overwrite u with numpy's fixed random numbers
      if pytest:
          np.random.seed(0)
          new_shape = cdf.get_shape().as_list()[:-1] + [N_samples]
          if det:
              u = np.linspace(0., 1., N_samples)
              u = np.broadcast_to(u, new_shape)
          else:
              u = np.random.rand(*new_shape)
          u = tf.cast(u, tf.float32)
  
      # Invert CDF
      inds = tf.searchsorted(cdf, u, side='right')
      below = tf.maximum(0, inds-1)
      above = tf.minimum(cdf.shape[-1]-1, inds)
      inds_g = tf.stack([below, above], -1)
  
      cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)
      bins_g = tf.gather(bins, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2)
  
      denom = (cdf_g[...,1]-cdf_g[...,0])
      denom = tf.where(denom<1e-5, tf.ones_like(denom), denom)
      t = (u-cdf_g[...,0])/denom
      samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])
  
      return samples