### Setup Data Directories and Load Metadata (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This script initializes and creates necessary directories for storing captured data, including RGB images, masks, and camera ground truth. It also loads metadata from JSON files for Kubric captures. Dependencies include `json`, `numpy`, `imageio`, `pdb`, `pathlib`, `cv2`, and `PIL`.

```python
import json
import numpy as np
import imageio
import pdb
from pathlib import Path
import cv2
from PIL import Image

# @markdown The base directory for all captures. This can be anything if you're running this notebook on your own Jupyter runtime.
save_dir = './data/'  
capture_name = 'kubric_single_car_rand_v2_np'
# The root directory for this capture.
root_dir = Path(save_dir, capture_name)
# Where to save RGB images.
rgb_dir = root_dir / 'rgb'
rgb_raw_dir = root_dir / 'rgb-raw'
static_rgb_dir = root_dir / 'static-rgb'
static_rgb_raw_dir = root_dir / 'static-rgb-raw'
# Where to save mask images.
mask_dir = root_dir / 'mask'
mask_raw_dir = root_dir / 'mask-raw'
camera_gt_dir = root_dir / 'camera-gt'
camera_test_gt_dir = root_dir / 'camera-paths-gt' / 'orbit-mild'

camera_gt_dir.mkdir(exist_ok=True, parents=True)
camera_test_gt_dir.mkdir(exist_ok=True, parents=True)
if static_rgb_raw_dir.exists():
    static_rgb_dir.mkdir(exist_ok=True, parents=True)

# make dirs for freeze test if available
freeze_test_dir = root_dir / 'freeze-test'
if freeze_test_dir.exists():
    rgb_raw_freeze_test_dir = freeze_test_dir / 'rgb-raw'
    rgb_freeze_test_dir = freeze_test_dir / 'rgb'
    cam_freeze_test_dir = freeze_test_dir / 'camera-gt'

    rgb_freeze_test_dir.mkdir(exist_ok=True, parents=True)
    cam_freeze_test_dir.mkdir(exist_ok=True, parents=True)

    with open(freeze_test_dir / 'metadata_kubric.json','r') as f:
            k_freeze_meta = json.load(f)

with open(root_dir / 'metadata_kubric.json','r') as f:
    k_meta = json.load(f)

tmp_rgb_raw_dir = rgb_raw_dir

rgb_list = list(sorted(rgb_raw_dir.glob( '*.png')))
image_size = imageio.imread(rgb_list[0]).shape[:2]
rgb_list[:3]

```

--------------------------------

### SE3Field Warp Module Initialization and Usage (Python)

Source: https://context7.com/chikayan/d2nerf/llms.txt

Demonstrates the initialization and example usage of the SE3Field module, which predicts rigid body transformations for point deformation. It requires JAX and Flax libraries. The module can optionally return the Jacobian for regularization purposes.

```python
from hypernerf import warping
import jax.numpy as jnp
from flax import linen as nn

warp_field = warping.SE3Field(
    min_deg=0,
    max_deg=8,
    trunk_depth=6,
    trunk_width=128,
    rotation_depth=0,
    rotation_width=128,
    translation_depth=0,
    translation_width=128,
)

points = jnp.array([[0.1, 0.2, 0.3]])
metadata_embed = jnp.array([[0.1] * 8])
extra_params = {'warp_alpha': 1.0}

# The warp field transforms points and optionally returns Jacobian
# warp_output = warp_field(points, metadata_embed, extra_params, return_jacobian=True)
# warped_points = warp_output['warped_points']
# jacobian = warp_output['jacobian']
```

--------------------------------

### Evaluate D²NeRF Model using Bash Script

Source: https://context7.com/chikayan/d2nerf/llms.txt

This snippet shows how to evaluate a trained D²NeRF model using the evaluation script. It covers running evaluation alongside training, in a single pass mode, and customizing the number of examples for evaluation using Gin bindings.

```bash
# Run evaluation alongside training
python eval.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_configs $CONFIG_PATH

# Run single evaluation (eval_once mode)
python eval.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_bindings="EvalConfig.eval_once=True" \
    --gin_configs $CONFIG_PATH

# Evaluate with custom number of examples
python eval.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_bindings="EvalConfig.num_train_eval=20" \
    --gin_bindings="EvalConfig.num_val_eval=20" \
    --gin_configs $CONFIG_PATH
```

--------------------------------

### Training Step Initialization and Execution (Python)

Source: https://context7.com/chikayan/d2nerf/llms.txt

This snippet details the initialization of the training state, including optimizer and scalar parameters, and demonstrates how to set up and execute a parallelized training step using JAX's pmap. It covers various loss weights and regularization options.

```python
from hypernerf import training
from hypernerf import model_utils
from hypernerf import models
from flax import optim
import jax

optimizer_def = optim.Adam(learning_rate=0.001)
params = {'model': model_params}
optimizer = optimizer_def.create(params)

state = model_utils.TrainState(
    optimizer=optimizer,
    nerf_alpha=0.0,
    warp_alpha=0.0,
    hyper_alpha=0.0,
    hyper_sheet_alpha=0.0,
)

scalar_params = training.ScalarParams(
    learning_rate=0.001,
    elastic_loss_weight=0.001,
    warp_reg_loss_weight=0.001,
    background_loss_weight=0.001,
    blendw_loss_weight=0.01,
    blendw_loss_skewness=1.75,
    shadow_r_loss_weight=0.1,
)

ptrain_step = jax.pmap(
    lambda *args: training.train_step(
        model,
        *args,
        use_elastic_loss=True,
        elastic_reduce_method='weight',
        elastic_loss_type='log_svals',
        use_background_loss=True,
        use_warp_reg_loss=True,
    ),
    axis_name='batch',
    in_axes=(0, 0, 0, None),
    donate_argnums=(2,)
)

rng_keys = jax.random.split(jax.random.PRNGKey(0), jax.local_device_count())
state, stats, keys, model_out = ptrain_step(rng_keys, state, batch, scalar_params)

psnr = stats['fine']['metric/psnr']
rgb_loss = stats['fine']['loss/rgb']
total_loss = stats['fine']['loss/total']
```

--------------------------------

### Gin Configuration Management

Source: https://context7.com/chikayan/d2nerf/llms.txt

Demonstrates how to define model architecture, training hyperparameters, and evaluation settings using Gin configuration files.

```python
include 'configs/decompose/train.gin'

image_scale = 4
max_steps = 100000

DecomposeNerfModel.use_viewdirs = True
DecomposeNerfModel.use_shadow_model = True
TrainConfig.batch_size = 2048
TrainConfig.lr_schedule = {
    'type': 'exponential',
    'initial_value': 0.001,
    'final_value': 0.0001,
    'num_steps': %max_steps,
}
```

--------------------------------

### Load Image with PIL

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Demonstrates how to list image files from a directory and open the first image using the PIL library.

```python
image_paths = list((rgb_dir / '1x').iterdir())
Image.open(image_paths[0])
```

--------------------------------

### Initialize and Run DecomposeNerfModel in Python

Source: https://context7.com/chikayan/d2nerf/llms.txt

This Python snippet demonstrates the initialization and forward pass of the DecomposeNerfModel. It shows how to set up the model with specific embeddings and parameters, prepare input rays, and perform a forward pass using JAX.

```python
import jax
from jax import random
import jax.numpy as jnp
from hypernerf import models

# Initialize the decomposed model
key = random.PRNGKey(0)
batch_size = 1024
embeddings_dict = {
    'warp': list(range(100)),      # Frame indices for warp embeddings
    'appearance': list(range(100)), # Frame indices for appearance
    'camera': list(range(1)),       # Camera IDs
}

# Construct the model and initialize parameters
model, params = models.construct_decompose_nerf(
    key=key,
    batch_size=batch_size,
    embeddings_dict=embeddings_dict,
    near=0.1,  # Near plane distance
    far=10.0   # Far plane distance
)

# Prepare input rays dictionary
rays_dict = {
    'origins': jnp.ones((batch_size, 3), jnp.float32),
    'directions': jnp.ones((batch_size, 3), jnp.float32),
    'metadata': {
        'warp': jnp.ones((batch_size, 1), jnp.uint32),
        'camera': jnp.ones((batch_size, 1), jnp.uint32),
        'appearance': jnp.ones((batch_size, 1), jnp.uint32),
        'time': jnp.ones((batch_size, 1), jnp.float32),
    }
}

# Extra parameters for warp field
extra_params = {
    'nerf_alpha': 1.0,           # Positional encoding alpha for NeRF
    'warp_alpha': 1.0,           # Positional encoding alpha for warp
    'hyper_alpha': 1.0,          # Alpha for hyper coordinates
    'hyper_sheet_alpha': 1.0,    # Alpha for hyper sheet MLP
}

# Forward pass
key1, key2 = random.split(key)
output = model.apply(
    {'params': params},
    rays_dict,
    extra_params=extra_params,
    rngs={'coarse': key1, 'fine': key2}
)

```

--------------------------------

### Train D²NeRF Model using Bash Script

Source: https://context7.com/chikayan/d2nerf/llms.txt

This snippet demonstrates how to train a D²NeRF model using the provided bash script. It covers setting environment variables for dataset and experiment paths, and running training with basic, debug, or custom configurations via Gin bindings.

```bash
# Set environment variables for paths
export DATASET_PATH=/path/to/dataset
export EXPERIMENT_PATH=/path/to/save/experiment
export CONFIG_PATH=configs/rl/001.gin

# Run training with basic configuration
python train.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_configs $CONFIG_PATH

# Run training with debug mode (JIT disabled for debugging)
python train.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_configs $CONFIG_PATH \
    --debug

# Run training with custom parameters via gin bindings
python train.py \
    --base_folder $EXPERIMENT_PATH \
    --gin_bindings="data_dir='$DATASET_PATH'" \
    --gin_bindings="TrainConfig.max_steps=200000" \
    --gin_bindings="TrainConfig.batch_size=2048" \
    --gin_configs $CONFIG_PATH
```

--------------------------------

### Camera and Ray Conversion Utilities

Source: https://context7.com/chikayan/d2nerf/llms.txt

Provides methods to load camera parameters from JSON, scale resolutions, and convert camera coordinates into ray origins and directions for rendering.

```python
from hypernerf import camera as cam
from hypernerf.datasets import core as datasets_core

camera = cam.Camera.from_json('/path/to/camera.json')
scaled_camera = camera.scale(0.25)
rays = datasets_core.camera_to_rays(camera)
origins = rays['origins']
directions = rays['directions']
```

--------------------------------

### List Raw RGB Images (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This snippet lists the first three raw RGB image file paths from a specified directory. It uses `pathlib.Path.glob` to find PNG files and `imageio.imread` to determine the image size. This is useful for verifying data loading and understanding image dimensions.

```python
rgb_list = list(sorted(rgb_raw_dir.glob( '*.png')))
image_size = imageio.imread(rgb_list[0]).shape[:2]
print(rgb_list[:3])
```

--------------------------------

### NerfiesDataSource for Dataset Loading (Python)

Source: https://context7.com/chikayan/d2nerf/llms.txt

Illustrates how to use the NerfiesDataSource class to load camera parameters, images, and metadata from datasets in the Nerfies format. It supports various configurations like image scaling, random seeding, and enabling/disabling embeddings. The class also provides methods to create training iterators and load individual items.

```python
from hypernerf import datasets
from hypernerf import gpath
import jax

datasource = datasets.NerfiesDataSource(
    data_dir=gpath.GPath('/path/to/dataset'),
    image_scale=4,
    random_seed=12345,
    use_warp_id=True,
    use_appearance_id=False,
    use_camera_id=False,
    use_time=False,
    mask_interest_region=False,
)

train_ids = datasource.train_ids
val_ids = datasource.val_ids
near = datasource.near
far = datasource.far
embeddings_dict = datasource.embeddings_dict

train_iter = datasource.create_iterator(
    item_ids=datasource.train_ids,
    flatten=True,
    shuffle=True,
    batch_size=2048,
    prefetch_size=3,
    shuffle_buffer_size=5000000,
    devices=jax.local_devices(),
)

item = datasource.get_item(train_ids[0])
rgb = item['rgb']
camera_params = item['camera_params']
metadata = item['metadata']
```

--------------------------------

### DataSource and Dataset Loading

Source: https://context7.com/chikayan/d2nerf/llms.txt

The NerfiesDataSource class handles loading camera parameters, images, and metadata from datasets in the Nerfies format.

```APIDOC
## DataSource and Dataset Loading

### Description
Loads camera parameters, images, and metadata from datasets in the Nerfies format using the `NerfiesDataSource` class.

### Method
N/A (Class instantiation and method calls)

### Endpoint
N/A

### Parameters
#### Class Instantiation Parameters (`NerfiesDataSource`)
- **data_dir** (gpath.GPath) - Path to the dataset directory.
- **image_scale** (int, optional) - Factor to downscale images. Defaults to 1.
- **random_seed** (int, optional) - Seed for random operations. Defaults to None.
- **use_warp_id** (bool, optional) - Whether to use warp IDs. Defaults to False.
- **use_appearance_id** (bool, optional) - Whether to use appearance IDs. Defaults to False.
- **use_camera_id** (bool, optional) - Whether to use camera IDs. Defaults to False.
- **use_time** (bool, optional) - Whether to use time information. Defaults to False.
- **mask_interest_region** (bool, optional) - Whether to mask the region of interest. Defaults to False.

#### Method Parameters (`create_iterator`)
- **item_ids** (list) - List of item IDs to include in the iterator.
- **flatten** (bool, optional) - Whether to flatten the data to rays. Defaults to False.
- **shuffle** (bool, optional) - Whether to shuffle the data. Defaults to False.
- **batch_size** (int, optional) - The batch size for the iterator. Defaults to 1.
- **prefetch_size** (int, optional) - Number of batches to prefetch. Defaults to 0.
- **shuffle_buffer_size** (int, optional) - Size of the shuffle buffer. Defaults to 0.
- **devices** (list, optional) - JAX devices to use for prefetching. Defaults to None.

### Request Example
```python
from hypernerf import datasets
from hypernerf import gpath
import jax

datasource = datasets.NerfiesDataSource(
    data_dir=gpath.GPath('/path/to/dataset'),
    image_scale=4,
    random_seed=12345,
    use_warp_id=True,
)

train_iter = datasource.create_iterator(
    item_ids=datasource.train_ids,
    flatten=True,
    shuffle=True,
    batch_size=2048,
    devices=jax.local_devices(),
)

item = datasource.get_item(datasource.train_ids[0])
rgb = item['rgb']
camera_params = item['camera_params']
metadata = item['metadata']
```

### Response
#### Success Response (200)
Provides access to dataset properties and iterators.

#### Response Example
```json
{
  "train_ids": ["frame_0001", "frame_0002", ...],
  "val_ids": ["frame_0050", "frame_0051", ...],
  "near": 0.1,
  "far": 10.0,
  "embeddings_dict": { ... },
  "iterator": "[Data iterator object]"
}
```
```

--------------------------------

### Training Step Function

Source: https://context7.com/chikayan/d2nerf/llms.txt

The core training step computes losses (RGB reconstruction, elastic regularization, blending weight entropy) and updates model parameters.

```APIDOC
## Training Step Function

### Description
Defines and executes a single training step, calculating various losses and updating model parameters. Supports parallel processing across devices.

### Method
N/A (Function definition and execution)

### Endpoint
N/A

### Parameters
#### `ScalarParams` Structure
- **learning_rate** (float) - The learning rate for the optimizer.
- **elastic_loss_weight** (float) - Weight for the elastic regularization loss.
- **warp_reg_loss_weight** (float) - Weight for the warp regularization loss.
- **background_loss_weight** (float) - Weight for the background reconstruction loss.
- **blendw_loss_weight** (float) - Weight for the blending weight entropy loss.
- **blendw_loss_skewness** (float) - Skewness parameter for the blending weight loss.
- **shadow_r_loss_weight** (float) - Weight for the shadow reconstruction loss.

#### `train_step` Function Arguments (within `pmap`)
- **model**: The NeRF model.
- **state**: The current training state (optimizer, alphas, etc.).
- **batch**: A dictionary containing the training data batch.
- **scalar_params**: An instance of `ScalarParams`.
- **use_elastic_loss** (bool) - Whether to use elastic loss.
- **elastic_reduce_method** (str) - Method for reducing elastic loss ('weight').
- **elastic_loss_type** (str) - Type of elastic loss ('log_svals').
- **use_background_loss** (bool) - Whether to use background loss.
- **use_warp_reg_loss** (bool) - Whether to use warp regularization loss.

### Request Example
```python
from hypernerf import training
from hypernerf import model_utils
from flax import optim
import jax

# Assume model_params, model, batch are defined
optimizer_def = optim.Adam(learning_rate=0.001)
params = {'model': model_params}
optimizer = optimizer_def.create(params)

state = model_utils.TrainState(
    optimizer=optimizer,
    nerf_alpha=0.0,
    warp_alpha=0.0,
    hyper_alpha=0.0,
    hyper_sheet_alpha=0.0,
)

scalar_params = training.ScalarParams(
    learning_rate=0.001,
    elastic_loss_weight=0.001,
    warp_reg_loss_weight=0.001,
    background_loss_weight=0.001,
    blendw_loss_weight=0.01,
    blendw_loss_skewness=1.75,
    shadow_r_loss_weight=0.1,
)

ptrain_step = jax.pmap(
    lambda *args: training.train_step(
        model,
        *args,
        use_elastic_loss=True,
        elastic_reduce_method='weight',
        elastic_loss_type='log_svals',
        use_background_loss=True,
        use_warp_reg_loss=True,
    ),
    axis_name='batch',
    in_axes=(0, 0, 0, None),
    donate_argnums=(2,)
)

rng_keys = jax.random.split(jax.random.PRNGKey(0), jax.local_device_count())
state, stats, keys, model_out = ptrain_step(rng_keys, state, batch, scalar_params)

psnr = stats['fine']['metric/psnr']
rgb_loss = stats['fine']['loss/rgb']
total_loss = stats['fine']['loss/total']
```

### Response
#### Success Response (200)
Returns updated training state, statistics, random keys, and model output.

#### Response Example
```json
{
  "state": {
    "optimizer": {"params": {...}},
    "nerf_alpha": 0.01,
    "warp_alpha": 0.005,
    "hyper_alpha": 0.0,
    "hyper_sheet_alpha": 0.0
  },
  "stats": {
    "fine": {
      "metric/psnr": 25.5,
      "loss/rgb": 0.01,
      "loss/total": 0.02
    }
  },
  "keys": "[Updated PRNG keys]",
  "model_out": { ... }
}
```
```

--------------------------------

### Downsample and Save Images

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Iterates through raw image directories, resizes images to be divisible by a factor, downsamples them according to defined scales, and saves the processed images to disk. This ensures consistent input dimensions for the NeRF model.

```python
if static_rgb_raw_dir.exists():
    for rgb_path in Path(static_rgb_raw_dir).glob('*.png'):
      image = make_divisible(imageio.imread(rgb_path), max(image_scales))
      for scale in image_scales:
        save_image(
            static_rgb_dir / f"{scale}x/{rgb_path.stem}.png",
            image_to_uint8(downsample_image(image, scale)))
            
if freeze_test_dir.exists() and (freeze_test_dir / 'static-rgb-raw').exists():
  (freeze_test_dir / 'static-rgb').mkdir(exist_ok=True, parents=True)
  for rgb_path in Path(freeze_test_dir / 'static-rgb-raw').glob('*.png'):
    image = make_divisible(imageio.imread(rgb_path), max(image_scales))
    for scale in image_scales:
      save_image(
          (freeze_test_dir / 'static-rgb') / f"{scale}x/{rgb_path.stem}.png",
          image_to_uint8(downsample_image(image, scale)))
```

--------------------------------

### Accessing Rendered Outputs

Source: https://context7.com/chikayan/d2nerf/llms.txt

This snippet shows how to access various rendered output components from a dictionary, such as RGB colors, depth maps, accumulated weights, and blending weights. These outputs are typically generated after a rendering process.

```python
rgb = output['fine']['rgb']
 depth = output['fine']['depth']
 acc = output['fine']['acc']
 blendw = output['fine']['blendw']
 rgb_d = output['fine']['rgb_d']
 rgb_s = output['fine']['rgb_s']
```

--------------------------------

### Save Metadata Information to metadata.json

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This Python script generates and saves metadata for training and validation image IDs. It creates a metadata dictionary mapping image IDs to warp, appearance, and camera IDs. It handles both training and validation sets, ensuring correct indexing for validation IDs. Finally, it saves this metadata to a JSON file named metadata.json in the specified root directory. It also includes logic to create a separate metadata file for freeze tests, copying frame information from the main metadata.

```python
# @title Save metadata information to `metadata.json`.
import bisect

metadata_json = {}
for i, image_id in enumerate(train_ids):
  metadata_json[image_id] = {
      'warp_id': i,
      'appearance_id': i,
      'camera_id': 0,
  }
for i, image_id in enumerate(val_ids):
  i = bisect.bisect_left(train_ids, image_id)
  metadata_json[image_id] = {
      'warp_id': i,
      'appearance_id': i,
      'camera_id': 0,
  }

metadata_json_path = root_dir / 'metadata.json'
with metadata_json_path.open('w') as f:
    json.dump(metadata_json, f, indent=2)

print(f'Saved metadata information to {metadata_json_path}')


# write meta info for freeze test if exists
if freeze_test_dir.exists():
  freeze_frame = metadata_json[list(sorted(rgb_raw_freeze_test_dir.glob('*.png')))[k_freeze_meta['test_freeze_frame'] - 1].stem]
  # freeze_frame = metadata_json[list(sorted(rgb_raw_freeze_test_dir.glob('*.png')))[k_freeze_meta['test_freeze_frame']].stem]

  metadata_freeze_json = {}
  for i, image_path in enumerate(list(sorted(rgb_raw_freeze_test_dir.glob('*.png')))):
    image_id = image_path.stem
    metadata_freeze_json[image_id] = freeze_frame

  metadata_json_path = freeze_test_dir / 'metadata.json'
  with metadata_json_path.open('w') as f:
      json.dump(metadata_freeze_json, f, indent=2)

  print(f'Saved freeze metadata information to {metadata_json_path}')
```

--------------------------------

### Accessing Model Outputs

Source: https://context7.com/chikayan/d2nerf/llms.txt

This section details how to access various rendered outputs from the model, such as RGB colors, depth, accumulated weights, and component-wise RGB values.

```APIDOC
## Accessing Model Outputs

### Description
Accesses rendered RGB colors, depth, accumulated weights, blending weights, and component-wise RGB values from the model output.

### Parameters
None

### Request Example
```python
output = {
    'fine': {
        'rgb': [ ... ],
        'depth': [ ... ],
        'acc': [ ... ],
        'blendw': [ ... ],
        'rgb_d': [ ... ],
        'rgb_s': [ ... ]
    }
}
rgb = output['fine']['rgb']
depth = output['fine']['depth']
acc = output['fine']['acc']
blendw = output['fine']['blendw']
rgb_d = output['fine']['rgb_d']
rgb_s = output['fine']['rgb_s']
```

### Response
#### Success Response (200)
Returns a dictionary containing various output components.

#### Response Example
```json
{
  "fine": {
    "rgb": "[Rendered RGB colors]",
    "depth": "[Expected depth]",
    "acc": "[Accumulated weights]",
    "blendw": "[Blending weights (dynamic vs static)]",
    "rgb_d": "[Dynamic component RGB]",
    "rgb_s": "[Static component RGB]"
  }
}
```
```

--------------------------------

### Resize Images (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This is a placeholder for a Python function designed to resize images into different scales. The specific implementation details, such as the resizing algorithm and output formats, are not provided in this snippet.

```python
# @title Resize images into different scales.
```

--------------------------------

### Calculate and Save Scene Information

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This script defines a function `get_bbox_corners` to compute the minimum and maximum coordinates of a bounding box from a set of points. It then calculates the scene scale and center based on these bounding box corners. Finally, it saves the scene's scale, center, bounding box, and near/far clipping planes to a JSON file named `scene_gt.json`. This is crucial for defining the spatial properties of the scene for rendering or simulation.

```python
def get_bbox_corners(points):
  lower = points.min(axis=0)
  upper = points.max(axis=0)
  return np.stack([lower, upper])

bbox_corners = get_bbox_corners(np.array(k_meta['instances'][0]['bboxes_3d'])[0])
bbox_corners = np.array([[0,   0,  0],
                         [50, 50, 10]])

# scene_center = np.mean(bbox_corners, axis=0)
scene_scale = 1.0 / np.sqrt(np.sum((bbox_corners[1] - bbox_corners[0]) ** 2))
scene_center = np.array([0., 0., 0.])
# scene_scale = 0.01

print(f'Scene Center: {scene_center}')
print(f'Scene Scale: {scene_scale}')

# @title Save scene information to `scene.json`.

scene_json_path = root_dir /  'scene_gt.json'
with scene_json_path.open('w') as f:
  json.dump({
      'scale': scene_scale,
      'center': scene_center.tolist(),
      'bbox': bbox_corners.tolist(),
      'near': 0.1,
      'far': 1. # bbox_corners[1].max() * scene_scale,
  }, f, indent=2)

print(f'Saved scene information to {scene_json_path}')
```

--------------------------------

### Generate dataset.json for Freeze Test

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

This script checks if a freeze test directory exists. If it does, it collects all image IDs from the freeze test directory, constructs a dataset JSON object containing counts and IDs, and saves it to dataset.json within the freeze test directory. This is useful for organizing and referencing data used in freeze tests.

```python
if freeze_test_dir.exists():
    all_ids_freeze_test = [img.stem for img in list(sorted(rgb_raw_freeze_test_dir.glob('*.png')))]
    train_ids_freeze_test = []
    val_ids_freeze_test = all_ids_freeze_test
    dataset_json = {
        'count': len(all_ids_freeze_test),
        'num_exemplars': len(train_ids_freeze_test),
        'ids': all_ids_freeze_test,
        'train_ids': train_ids_freeze_test,
        'val_ids': val_ids_freeze_test,
    }

    dataset_json_path = freeze_test_dir / 'dataset.json'
    with dataset_json_path.open('w') as f:
        json.dump(dataset_json, f, indent=2)
```

--------------------------------

### Parallel Model Execution with JAX

Source: https://context7.com/chikayan/d2nerf/llms.txt

Defines a parallelized model function using jax.pmap for distributed rendering. It processes ray batches across available devices to generate RGB, depth, and accumulation outputs.

```python
def _model_fn(key_0, key_1, params, rays_dict, extra_params):
    out = model.apply(
        {'params': params},
        rays_dict,
        extra_params=extra_params,
        metadata_encoded=False,
        rngs={'coarse': key_0, 'fine': key_1},
        mutable=False
    )
    return jax.lax.all_gather(out, axis_name='batch')

pmodel_fn = jax.pmap(
    _model_fn,
    in_axes=(0, 0, 0, 0, 0),
    devices=jax.local_devices(),
    axis_name='batch',
)

model_out = evaluation.render_image(
    state=state,
    rays_dict=rays_dict,
    model_fn=pmodel_fn,
    device_count=jax.local_device_count(),
    rng=jax.random.PRNGKey(0),
    chunk=8192,
    normalise_rendering=False,
    use_tsne=False,
)
```

--------------------------------

### Save Image Utility (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Saves a NumPy array as an image to the specified path. It handles directory creation and supports various image formats based on the file extension. Dependencies include Pillow (PIL) and NumPy.

```Python
def save_image(path, image: np.ndarray) -> None:
  print(f'Saving {path}')
  if not path.parent.exists():
    path.parent.mkdir(exist_ok=True, parents=True)
  with path.open('wb') as f:
    image = Image.fromarray(np.asarray(image))
    image.save(f, format=path.suffix.lstrip('.'))
```

--------------------------------

### Generate Circular Camera Trajectory

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Calculates a circular path for test cameras around a central point and saves the resulting camera configurations to JSON files.

```python
from hypernerf import camera as camera_model
positions = np.array([cam['position'] for cam in cameras])
center = np.average(positions, axis=0)
center_axis = (0. - center) / np.linalg.norm(center)
axis_2 = center_axis.copy()
axis_2[0] += 1
axis_a = np.cross(center_axis, axis_2) / np.linalg.norm(np.cross(center_axis, axis_2))
axis_b = np.cross(center_axis, axis_a) / np.linalg.norm(np.cross(center_axis, axis_a))

RADIUS_RATIO = 0.5
N_TEST = 100
angles = np.linspace(0, 2 * np.pi, N_TEST)
for i, ang in enumerate(angles):
    new_pos = center + radius * np.cos(ang) * axis_a + radius * np.sin(ang) * axis_b
    test_cam = camera_model.Camera.from_json(str(camera_gt_dir / rgb_list[0].stem) + '.json')
    test_cam.position = new_pos
    test_cam.look_at(new_pos, np.array([0,0,0]), np.array([1,0,0]))
    with open(str(camera_test_gt_dir / f'{i:06d}.json'), 'w') as f:
        json.dump(test_cam.to_json(), f, indent=2)
```

--------------------------------

### Render Image Function

Source: https://context7.com/chikayan/d2nerf/llms.txt

Render complete images by batching rays and processing them through the model, with support for chunked rendering to handle memory constraints.

```APIDOC
## Render Image Function

### Description
Renders complete images by processing rays through the model, with support for chunked rendering to manage memory usage.

### Method
N/A (Function definition and execution)

### Endpoint
N/A

### Parameters
None explicitly defined in the provided text, but typically requires model, camera parameters, and ray batches.

### Request Example
```python
from hypernerf import evaluation
import jax
import jax.numpy as jnp

# Assume model, state, camera_params, rays are defined
# image = evaluation.render_image(model, state, camera_params, rays, chunk=1024)
```

### Response
#### Success Response (200)
Returns the rendered image as a JAX NumPy array.

#### Response Example
```json
{
  "rendered_image": "[NumPy array representing the image]"
}
```
```

--------------------------------

### Save Dataset Split

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Generates a dataset.json file containing training and validation splits based on image IDs.

```python
VAL_RATIO = 0
all_ids = list(map(lambda x: x.stem, rgb_list))
val_ids = all_ids[::int(1 / VAL_RATIO)] if VAL_RATIO > 0 else []
train_ids = sorted(set(all_ids) - set(val_ids))
dataset_json = {'count': len(all_ids), 'num_exemplars': len(train_ids), 'ids': all_ids, 'train_ids': train_ids, 'val_ids': val_ids}
with (root_dir / 'dataset.json').open('w') as f:
    json.dump(dataset_json, f, indent=2)
```

--------------------------------

### Generate Ground Truth Camera Metadata

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Iterates through camera data to convert quaternions to rotation matrices and save camera parameters into JSON files. This process handles both standard and freeze-test camera configurations.

```python
from scipy.spatial.transform import Rotation as R

cameras = []
for i, rgb in enumerate(rgb_list):
    meta = {}
    quat = k_meta['camera']['quaternions'][i]
    quat = [quat[1],quat[2],quat[3],quat[0]]
    rot = R.from_quat(quat)
    meta['orientation'] = rot.as_matrix().T.tolist()
    meta['position'] = k_meta['camera']['positions'][i]
    meta['focal_length'] = k_meta['camera']['focal_length'] * image_size[0] / k_meta['camera']['sensor_width']
    meta['principal_point'] = [image_size[0]/2, image_size[1]/2]
    meta['skew'] = 0.
    meta['pixel_aspect_ratio'] = 1.
    meta['radial_distortion'] = [0., 0., 0.]
    meta['tangential_distortion'] = [0., 0.]
    meta['image_size'] = image_size
    cameras.append(meta)
    with open(str(camera_gt_dir / rgb.stem) + '.json', 'w') as f:
        json.dump(meta, f, indent=2)
```

--------------------------------

### Render Image Functionality (Python)

Source: https://context7.com/chikayan/d2nerf/llms.txt

This snippet indicates the availability of a render image function within the evaluation module. The function is designed to render complete images by processing rays through the model, with support for chunked rendering to manage memory usage.

```python
from hypernerf import evaluation
import jax
import jax.numpy as jnp
```

--------------------------------

### Batch Image Processing Script (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Processes images from specified directories, downsampling them by various scales and saving the results. It handles both RGB images and masks, applying `make_divisible`, `downsample_image`, `image_to_uint8`, and `save_image` functions. Dependencies include `imageio`, `numpy`, `pathlib`, and `cv2`.

```Python
image_scales = "2"  # @param {type: "string"}
image_scales = [int(x) for x in image_scales.split(',')]


for image_path in Path(tmp_rgb_raw_dir).glob('*.png'):
  image = make_divisible(imageio.imread(image_path), max(image_scales))
  for scale in image_scales:
    save_image(
        rgb_dir / f'{scale}x/{image_path.stem}.png',
        image_to_uint8(downsample_image(image, scale)))

if mask_raw_dir.exists():
    for i, mask_path in enumerate(sorted(list(Path(mask_raw_dir).glob('*.png')))):
      mask = make_divisible(imageio.imread(mask_path), max(image_scales))
      for scale in image_scales:
        save_image(
            mask_dir / f"{scale}x/{rgb_list[i].stem.replace('segmentation','rgba')}.png",
            image_to_uint8(downsample_image(mask, scale)))


if freeze_test_dir.exists():
    for rgb_path in Path(rgb_raw_freeze_test_dir).glob('*.png'):
      image = make_divisible(imageio.imread(rgb_path), max(image_scales))
      for scale in image_scales:
        save_image(
            rgb_freeze_test_dir / f"{scale}x/{rgb_path.stem}.png",
            image_to_uint8(downsample_image(image, scale)))
```

--------------------------------

### Evaluation Metrics Computation

Source: https://context7.com/chikayan/d2nerf/llms.txt

Computes standard image quality metrics including PSNR, SSIM, and Multi-Scale SSIM using TensorFlow and custom utility functions.

```python
from hypernerf import utils
import tensorflow as tf

mse = ((predicted_rgb - target_rgb) ** 2).mean()
psnr = utils.compute_psnr(mse)

def compute_ssim(image1, image2):
    image1 = tf.convert_to_tensor(image1)
    image2 = tf.convert_to_tensor(image2)
    return tf.image.ssim(image1, image2, max_val=1.0)

def compute_multiscale_ssim(image1, image2):
    image1 = tf.convert_to_tensor(image1)
    image2 = tf.convert_to_tensor(image2)
    return tf.image.ssim_multiscale(image1, image2, max_val=1.0, filter_size=4)
```

--------------------------------

### Image to uint8 Conversion (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Converts an input image (NumPy array) to the uint8 data type. It raises a ValueError if the input is not a floating-point type. This is useful for standardizing image formats before saving or further processing.

```Python
def image_to_uint8(image: np.ndarray) -> np.ndarray:
  """Convert the image to a uint8 array."""
  if image.dtype == np.uint8:
    return image
  if not issubclass(image.dtype.type, np.floating):
    raise ValueError(
        f'Input image should be a floating type but is of type {image.dtype!r}')
  return (image * 255).clip(0.0, 255).astype(np.uint8)
```

--------------------------------

### SE3Field Warp Module

Source: https://context7.com/chikayan/d2nerf/llms.txt

The SE(3) warp field predicts rigid body transformations to deform points from observation space to canonical space, enabling handling of object motion.

```APIDOC
## SE3Field Warp Module

### Description
Predicts rigid body transformations to deform points from observation space to canonical space, handling object motion.

### Method
N/A (Class instantiation and method calls)

### Endpoint
N/A

### Parameters
#### Class Instantiation Parameters
- **min_deg** (int) - Minimum degree for the polynomial representation.
- **max_deg** (int) - Maximum degree for the polynomial representation.
- **trunk_depth** (int) - Depth of the trunk network.
- **trunk_width** (int) - Width of the trunk network.
- **rotation_depth** (int) - Depth of the rotation prediction network.
- **rotation_width** (int) - Width of the rotation prediction network.
- **translation_depth** (int) - Depth of the translation prediction network.
- **translation_width** (int) - Width of the translation prediction network.

#### Method Parameters (`__call__`)
- **points** (jax.numpy.ndarray) - Input 3D points.
- **metadata_embed** (jax.numpy.ndarray) - Embedding for warp information.
- **extra_params** (dict) - Additional parameters, e.g., {'warp_alpha': 1.0}.
- **return_jacobian** (bool, optional) - Whether to return the Jacobian matrix. Defaults to False.

### Request Example
```python
from hypernerf import warping
import jax.numpy as jnp

warp_field = warping.SE3Field(
    min_deg=0,
    max_deg=8,
    trunk_depth=6,
    trunk_width=128,
    rotation_depth=0,
    rotation_width=128,
    translation_depth=0,
    translation_width=128,
)

points = jnp.array([[0.1, 0.2, 0.3]])
metadata_embed = jnp.array([[0.1] * 8])
extra_params = {'warp_alpha': 1.0}

warp_output = warp_field(points, metadata_embed, extra_params, return_jacobian=True)
warped_points = warp_output['warped_points']
jacobian = warp_output['jacobian']
```

### Response
#### Success Response (200)
Returns a dictionary containing:
- **warped_points** (jax.numpy.ndarray): The transformed points.
- **jacobian** (jax.numpy.ndarray, optional): The Jacobian matrix if `return_jacobian` is True.

#### Response Example
```json
{
  "warped_points": [[x, y, z]],
  "jacobian": [[...]]
}
```
```

--------------------------------

### Make Image Divisible Utility (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Trims the height and width of an image (NumPy array) so they are divisible by a specified divisor. This function is crucial for ensuring compatibility with models or operations that require specific image dimensions.

```Python
def make_divisible(image: np.ndarray, divisor: int) -> np.ndarray:
  """Trim the image if not divisible by the divisor."""
  height, width = image.shape[:2]
  if height % divisor == 0 and width % divisor == 0:
    return image

  new_height = height - height % divisor
  new_width = width - width % divisor

  return image[:new_height, :new_width]
```

--------------------------------

### Compute LPIPS Perceptual Metric

Source: https://context7.com/chikayan/d2nerf/llms.txt

Calculates the LPIPS score between two images. The function normalizes input images to the [-1, 1] range and transforms them into the (N, C, H, W) format required by the LPIPS model.

```python
import numpy as np
import torch
from lpips import LPIPS

loss_fn_alex = LPIPS(net='alex', version=0.1)

def compute_lpips(image1, image2):
    # Normalize to [-1, 1]
    image1 = (np.array(image1) - 0.5) * 2
    image2 = (np.array(image2) - 0.5) * 2
    # Transpose to (N, C, H, W)
    image1 = np.transpose(image1[None, ...], [0, 3, 1, 2])
    image2 = np.transpose(image2[None, ...], [0, 3, 1, 2])
    return loss_fn_alex(torch.from_numpy(image1), torch.from_numpy(image2)).item()

lpips = compute_lpips(predicted_rgb, target_rgb)
```

--------------------------------

### Downsample Image by Scale (Python)

Source: https://github.com/chikayan/d2nerf/blob/main/kubric_meta_parser.ipynb

Downsamples an image (NumPy array) by an integer scale factor using area interpolation (cv2.INTER_AREA) to prevent moire artifacts. It raises a ValueError if the image dimensions are not divisible by the scale. Dependencies include OpenCV (cv2) and NumPy.

```Python
def downsample_image(image: np.ndarray, scale: int) -> np.ndarray:
  """Downsamples the image by an integer factor to prevent artifacts."""
  if scale == 1:
    return image

  height, width = image.shape[:2]
  if height % scale > 0 or width % scale > 0:
    raise ValueError(f'Image shape ({height},{width}) must be divisible by the'
                     f' scale ({scale}).')
  out_height, out_width = height // scale, width // scale
  resized = cv2.resize(image, (out_width, out_height), cv2.INTER_AREA)
  return resized
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.