from .config_helpers import parse_config

__all__ = (
    "parse_config",
)

magnerf/configs/__init__.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "baby",
    "runname": "test",
    'dset_type': 'handheld',
    "logdir": "./logs/handheld",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/handheld",
    "hold_every": 25,
    "finetune_iters": 1000,

    "bbox": "-4.0,-5.0,1.0,5.0,4.5,8.0",
    "near_plane": 0.2,
    "far_plane": 100.0,
    "use_app_embedding": False,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "dynamic_start": 154,

    # Optimization settings
    "num_steps": 20_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 1.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0,
    'gridchange_loss_weight': 0.5,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 10_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/baby.py
3d-motion-magnification/3d-motion-mag
config = {
  'expname': 'drums_3Hz_0_0005',
  'logdir': './logs/blender',
  'device': 'cuda:0',

  'data_downsample': 1.0,
  'dset_type': 'synthetic',
  'data_dirs': './data/synthetic',
  'contract': False,
  'ndc': False,

  # Optimization settings
  'num_steps': 30001,
  'batch_size': 4096,
  'scheduler_type': 'warmup_cosine',
  'optim_type': 'adam',
  'lr': 0.01,
  'finetune_iters': 10001,
  'finetune_lr': 1e-4,

  # Regularization
  'plane_tv_weight': 0.01,
  'plane_tv_weight_proposal_net': 0.0001,
  'histogram_loss_weight': 1.0,
  'distortion_loss_weight': 0.001,
  "gridchange_loss_weight": 0.0,

  # Training settings
  'save_every': 30000,
  'valid_every': 30000,
  'save_outputs': True,
  'train_fp16': True,

  # Raymarching settings
  'single_jitter': False,
  'num_samples': 48,
  # proposal sampling
  'num_proposal_samples': [256, 128],
  'num_proposal_iterations': 2,
  'use_same_proposal_network': False,
  'use_proposal_weight_anneal': True,
  'proposal_net_args_list': [
    {'num_input_coords': 3, 'num_output_coords': 8, 'resolution': [64, 64, 64]},
    {'num_input_coords': 3, 'num_output_coords': 8, 'resolution': [128, 128, 128]}
  ],

  # Model settings
  'multiscale_res': [1],
  'density_activation': 'trunc_exp',
  'concat_features_across_scales': True,
  'linear_decoder': False,
  'grid_config': [{
    'grid_dimensions': 2,
    'input_coordinate_dim': 3,
    'output_coordinate_dim': 32,
    'resolution': [512, 512, 512]
  }],
}

magnerf/configs/blender.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "breathing",
    "runname": "",
    'dset_type': 'tripod',

    "logdir": "./logs/tripod",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/tripod",
    "hold_every": 25,

    "bbox": "-5.0,-4.0,2.0,5.0,4.0,9.5",

    "use_app_embedding": True,
    "near_plane": 1.0,
    "far_plane": 15.0,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "frame_offset": 1,

    # Optimization settings
    "num_steps": 30_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 5.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0,
    'gridchange_loss_weight': 1,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 5_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/breathing.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "fork",
    "runname": "",
    'dset_type': 'tripod',

    "logdir": "./logs/tripod",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/tripod",
    "hold_every": 25,

    "bbox": "-5.8,-3.0,3.5,4.5,3.0,11.0",

    "use_app_embedding": True,
    "near_plane": 0.5,
    "far_plane": 12.0,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "frame_offset": 1,

    # Optimization settings
    "num_steps": 30_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 5.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0,
    'gridchange_loss_weight': 1,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 5_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/fork.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "jackson",
    "runname": "test",
    'dset_type': 'handheld',
    "logdir": "./logs/handheld",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/handheld",
    "hold_every": 25,
    "finetune_iters": 1000,

    "bbox": "-6.0,-4.5,1.0,6.0,5.0,23.0",
    "near_plane": 1.0,
    "far_plane": 200.0,
    "use_app_embedding": False,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "dynamic_start": 20,

    # Optimization settings
    "num_steps": 20_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 2.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0,
    'gridchange_loss_weight': 0.4,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 10_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/jackson.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "ronaldo",
    "runname": "test",
    'dset_type': 'handheld',
    "logdir": "./logs/handheld",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/handheld",
    "hold_every": 25,
    "finetune_iters": 1000,

    "bbox": "-10.5,-9.5,6.0,7.0,6.5,20.0",
    "near_plane": 2.0,
    "far_plane": 150.0,
    "use_app_embedding": False,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "dynamic_start": 100,

    # Optimization settings
    "num_steps": 20_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 1.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0.0,
    'gridchange_loss_weight': 1.5,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 10_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/ronaldo.py
3d-motion-magnification/3d-motion-mag
config = {
    "expname": "standing",
    "runname": "",
    'dset_type': 'tripod',

    "logdir": "./logs/tripod",
    "device": "cuda:0",

    # Data settings
    "data_downsample": 1,
    "data_dirs": "./data/tripod",
    "hold_every": 25,

    "bbox": "-3.0,-3.0,2.65,4.0,2.5,6.5",

    "use_app_embedding": True,
    "near_plane": 2.5,
    "far_plane": 20.0,
    "opaque_background": False,
    "contract": False,
    "bg_color_aug": "random",
    "frame_offset": 1,

    # Optimization settings
    "num_steps": 30_001,
    "batch_size": 8192,
    "eval_batch_size": 8192,
    "scheduler_type": "warmup_cosine",
    "optim_type": "adam",
    "lr": 2e-2,

    # Regularization
    'plane_tv_weight': 2e-4,
    'plane_tv_weight_proposal_net': 2e-4,
    'distortion_loss_weight': 5.0,
    'histogram_loss_weight': 1.0,
    "l1_proposal_net_weight": 0,
    'gridchange_loss_weight': 1,

    # Training settings
    "train_fp16": True,
    "save_every": 10_000,
    "valid_every": 5_000,
    "save_outputs": True,

    # Raymarching settings
    "num_samples": 64,
    "single_jitter": False,
    # proposal sampling
    "num_proposal_samples": [256, 128],
    "num_proposal_iterations": 2,
    "use_same_proposal_network": False,
    "use_proposal_weight_anneal": True,
    "proposal_net_args_list": [
        {"resolution": [128, 128, 128], "num_input_coords": 3, "num_output_coords": 8},
        {"resolution": [256, 256, 256], "num_input_coords": 3, "num_output_coords": 8},
    ],

    # Model settings
    'appearance_embedding_dim': 8,
    'multiscale_res': [1],
    'density_activation': 'trunc_exp',
    'concat_features_across_scales': True,
    'linear_decoder': False,
    'grid_config': [{
        'grid_dimensions': 2,
        'input_coordinate_dim': 3,
        'output_coordinate_dim': 32,
        'resolution': [512, 512, 512]
    }],
}

magnerf/configs/standing.py
3d-motion-magnification/3d-motion-mag
from .synthetic_nerf_dataset import SyntheticNerfDataset
from .nerf_360_v2 import HandheldDataset
from .tripod import TripodDataset

__all__ = (
    "SyntheticNerfDataset",
    "HandheldDataset",
    "TripodDataset",
)

magnerf/datasets/__init__.py
3d-motion-magnification/3d-motion-mag
from abc import ABC
import os
from typing import Optional, List, Union

import torch
from torch.utils.data import Dataset

from .intrinsics import Intrinsics


class BaseDataset(Dataset, ABC):
    def __init__(self,
                 datadir: str,
                 scene_bbox: torch.Tensor,
                 split: str,
                 is_ndc: bool,
                 is_contracted: bool,
                 rays_o: Optional[torch.Tensor],
                 rays_d: Optional[torch.Tensor],
                 intrinsics: Union[Intrinsics, List[Intrinsics]],
                 batch_size: Optional[int] = None,
                 imgs: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
                 sampling_weights: Optional[torch.Tensor] = None,
                 weights_subsampled: int = 1,
                 ):
        self.datadir = datadir
        self.name = os.path.basename(self.datadir)
        self.scene_bbox = scene_bbox
        self.split = split
        self.is_ndc = is_ndc
        self.is_contracted = is_contracted
        self.weights_subsampled = weights_subsampled
        self.batch_size = batch_size
        if self.split == 'train':
            assert self.batch_size is not None
        self.rays_o = rays_o
        self.rays_d = rays_d
        self.imgs = imgs
        if self.imgs is not None:
            self.num_samples = len(self.imgs)
        elif self.rays_o is not None:
            self.num_samples = len(self.rays_o)
        else:
            self.num_samples = None
            #raise RuntimeError("Can't figure out num_samples.")
        self.intrinsics = intrinsics
        self.sampling_weights = sampling_weights
        if self.sampling_weights is not None:
            assert len(self.sampling_weights) == self.num_samples, (
                f"Expected {self.num_samples} sampling weights but given {len(self.sampling_weights)}."
            )
        self.sampling_batch_size = 2_000_000  # Increase this?
        if self.num_samples is not None:
            self.use_permutation = self.num_samples < 100_000_000  # 64M is static
        else:
            self.use_permutation = True
        self.perm = None

    @property
    def img_h(self) -> Union[int, List[int]]:
        if isinstance(self.intrinsics, list):
            return [i.height for i in self.intrinsics]
        return self.intrinsics.height

    @property
    def img_w(self) -> Union[int, List[int]]:
        if isinstance(self.intrinsics, list):
            return [i.width for i in self.intrinsics]
        return self.intrinsics.width

    def reset_iter(self):
        if self.sampling_weights is None and self.use_permutation:
            self.perm = torch.randperm(self.num_samples)
        else:
            del self.perm
            self.perm = None

    def get_rand_ids(self, index):
        assert self.batch_size is not None, "Can't get rand_ids for test split"
        if self.sampling_weights is not None:
            batch_size = self.batch_size // (self.weights_subsampled ** 2)
            num_weights = len(self.sampling_weights)
            if num_weights > self.sampling_batch_size:
                # Take a uniform random sample first, then according to the weights
                subset = torch.randint(
                    0, num_weights, size=(self.sampling_batch_size,),
                    dtype=torch.int64, device=self.sampling_weights.device)
                samples = torch.multinomial(
                    input=self.sampling_weights[subset], num_samples=batch_size)
                return subset[samples]
            return torch.multinomial(
                input=self.sampling_weights, num_samples=batch_size)
        else:
            batch_size = self.batch_size
            if self.use_permutation:
                return self.perm[index * batch_size: (index + 1) * batch_size]
            else:
                return torch.randint(0, self.num_samples, size=(batch_size, ))

    def __len__(self):
        if self.split == 'train':
            return (self.num_samples + self.batch_size - 1) // self.batch_size
        else:
            return self.num_samples

    def __getitem__(self, index, return_idxs: bool = False):
        if self.split == 'train':
            index = self.get_rand_ids(index)
        out = {}
        if self.rays_o is not None:
            out["rays_o"] = self.rays_o[index]
        if self.rays_d is not None:
            out["rays_d"] = self.rays_d[index]
        if self.imgs is not None:
            out["imgs"] = self.imgs[index]
        else:
            out["imgs"] = None
        if return_idxs:
            return out, index
        return out

magnerf/datasets/base_dataset.py
3d-motion-magnification/3d-motion-mag
import collections
import struct

import numpy as np


Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"])
CAMERA_MODELS = {
    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
    CameraModel(model_id=7, model_name="FOV", num_params=5),
    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS])


class Image(BaseImage):
    def qvec2rotmat(self):
        return qvec2rotmat(self.qvec)


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
    """Read and unpack the next bytes from a binary file.
    :param fid:
    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
    :param endian_character: Any of {@, =, <, >, !}
    :return: Tuple of read and unpacked values.
    """
    data = fid.read(num_bytes)
    return struct.unpack(endian_character + format_char_sequence, data)


def read_cameras_text(path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasText(const std::string& path)
        void Reconstruction::ReadCamerasText(const std::string& path)
    """
    cameras = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                camera_id = int(elems[0])
                model = elems[1]
                width = int(elems[2])
                height = int(elems[3])
                params = np.array(tuple(map(float, elems[4:])))
                cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params)
    return cameras


def read_cameras_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasBinary(const std::string& path)
        void Reconstruction::ReadCamerasBinary(const std::string& path)
    """
    cameras = {}
    with open(path_to_model_file, "rb") as fid:
        num_cameras = read_next_bytes(fid, 8, "Q")[0]
        for camera_line_index in range(num_cameras):
            camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ")
            camera_id = camera_properties[0]
            model_id = camera_properties[1]
            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name  # only pinhole (model_id = 1) supported
            width = camera_properties[2]
            height = camera_properties[3]
            num_params = CAMERA_MODEL_IDS[model_id].num_params
            params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params)
            cameras[camera_id] = Camera(
                id=camera_id, model=model_name, width=width, height=height, params=np.array(params)
            )
        assert len(cameras) == num_cameras
    return cameras


def read_images_text(path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesText(const std::string& path)
        void Reconstruction::WriteImagesText(const std::string& path)
    """
    images = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                image_id = int(elems[0])
                qvec = np.array(tuple(map(float, elems[1:5])))
                tvec = np.array(tuple(map(float, elems[5:8])))
                camera_id = int(elems[8])
                image_name = elems[9]
                elems = fid.readline().split()
                xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))])
                point3D_ids = np.array(tuple(map(int, elems[2::3])))
                images[image_id] = Image(
                    id=image_id,
                    qvec=qvec,
                    tvec=tvec,
                    camera_id=camera_id,
                    name=image_name,
                    xys=xys,
                    point3D_ids=point3D_ids,
                )
    return images


def read_images_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesBinary(const std::string& path)
        void Reconstruction::WriteImagesBinary(const std::string& path)
    """
    images = {}
    with open(path_to_model_file, "rb") as fid:
        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
        for image_index in range(num_reg_images):
            binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi")
            image_id = binary_image_properties[0]
            qvec = np.array(binary_image_properties[1:5])
            tvec = np.array(binary_image_properties[5:8])
            camera_id = binary_image_properties[8]
            image_name = ""
            current_char = read_next_bytes(fid, 1, "c")[0]
            while current_char != b"\x00":  # look for the ASCII 0 entry
                image_name += current_char.decode("utf-8")
                current_char = read_next_bytes(fid, 1, "c")[0]
            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
            x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D)
            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))])
            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
            images[image_id] = Image(
                id=image_id,
                qvec=qvec,
                tvec=tvec,
                camera_id=camera_id,
                name=image_name,
                xys=xys,
                point3D_ids=point3D_ids,
            )
    return images


def qvec2rotmat(qvec):
    return np.array(
        [
            [
                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
            ],
            [
                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
            ],
            [
                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
            ],
        ]
    )


def rotmat2qvec(R):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
    K = (
        np.array(
            [
                [Rxx - Ryy - Rzz, 0, 0, 0],
                [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
                [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
                [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
            ]
        )
        / 3.0
    )
    eigvals, eigvecs = np.linalg.eigh(K)
    qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
    if qvec[0] < 0:
        qvec *= -1
    return

magnerf/datasets/colmap_utils.py
3d-motion-magnification/3d-motion-mag
from typing import Tuple, Optional, Dict, Any, List
import logging as log
import os
import resource

import torch
from torch.multiprocessing import Pool
import torchvision.transforms
from PIL import Image
import imageio.v3 as iio

from magnerf.utils.my_tqdm import tqdm

pil2tensor = torchvision.transforms.ToTensor()
# increase ulimit -n (number of open files) otherwise parallel loading might fail
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (16192, rlimit[1]))


def _load_phototourism_image(idx: int,
                             paths: List[str],
                             out_h: List[int],
                             out_w: List[int]) -> torch.Tensor:
    f_path = paths[idx]
    img = Image.open(f_path).convert('RGB')
    img.resize((out_w[idx], out_h[idx]), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]
    return img


def _parallel_loader_phototourism_image(args):
    torch.set_num_threads(1)
    return _load_phototourism_image(**args)


def _load_llff_image(idx: int,
                     paths: List[str],
                     data_dir: str,
                     out_h: int,
                     out_w: int,
                     ) -> torch.Tensor:
    f_path = os.path.join(data_dir, paths[idx])
    img = Image.open(f_path).convert('RGB')

    img = img.resize((out_w, out_h), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]
    return img


def _parallel_loader_llff_image(args):
    torch.set_num_threads(1)
    return _load_llff_image(**args)


def _load_nerf_image_pose(idx: int,
                          frames: List[Dict[str, Any]],
                          data_dir: str,
                          out_h: Optional[int],
                          out_w: Optional[int],
                          downsample: float,
                          ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
    # Fix file-path
    f_path = os.path.join(data_dir, frames[idx]['file_path'])
    if '.' not in os.path.basename(f_path):
        f_path += '.png'  # so silly...
    if not os.path.exists(f_path):  # there are non-exist paths in fox...
        return None
    img = Image.open(f_path)
    if out_h is None:
        out_h = int(img.size[0] / downsample)
    if out_w is None:
        out_w = int(img.size[1] / downsample)
    # Now we should downsample to out_h, out_w and low-pass filter to resolution * 2.
    # We only do the low-pass filtering if resolution * 2 is lower-res than out_h, out_w
    if out_h != out_w:
        log.warning("360 non-square")
    img = img.resize((out_w, out_h), Image.LANCZOS)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]

    pose = torch.tensor(frames[idx]['transform_matrix'], dtype=torch.float32)

    return (img, pose)

def _load_humman_image_pose(idx: int,
                          frames: List[Dict[str, Any]],
                          data_dir: str,
                          out_h: Optional[int],
                          out_w: Optional[int],
                          downsample: float,
                          ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
    # Fix file-path
    f_path = os.path.join(data_dir, frames[idx]['kinect_color'])
    mask_path = os.path.join(data_dir, frames[idx]['kinect_mask'])
    print(f_path, mask_path)
    print(os.path.exists(f_path))
    if '.' not in os.path.basename(f_path):
        f_path += '.png'  # so silly...
    if not os.path.exists(f_path):  # there are non-exist paths in fox...
        return None
    img = Image.open(f_path)
    mask = Image.open(mask_path)
    if out_h is None:
        out_h = int(img.size[1] / downsample)
    if out_w is None:
        out_w = int(img.size[0] / downsample)

    # Now we should downsample to out_h, out_w and low-pass filter to resolution * 2.
    # We only do the low-pass filtering if resolution * 2 is lower-res than out_h, out_w
    if out_h != out_w:
        log.warning("360 non-square")
    img = img.resize((out_w, out_h), Image.BILINEAR)
    img = pil2tensor(img)  # [C, H, W]
    img = img.permute(1, 2, 0)  # [H, W, C]
    # pretend like mask is an additional channel
    mask = mask.resize((out_w, out_h), Image.BILINEAR)
    mask = pil2tensor(mask) # [C, H, W]
    mask = mask.permute(1, 2, 0)[:,:,-1:]  # [H, W, C]

    img = torch.cat([img, mask], -1)

    pose = torch.tensor(frames[idx]['transform_matrix'])

    # blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
    # c2w = np.array(frames[idx]['transform_matrix']) @ blender2opencv
    # print(c2w.shape)
    # c2w = np.concatenate([c2w[1:2, :], -c2w[0:1, :], c2w[2:, :]], 0)
    # c2w = np.moveaxis(c2w, -1, 0).astype(np.float32)
    # c2w[0:3, 2] *= -1  # flip the y and z axis
    # c2w[0:3, 1] *= -1
    # c2w = c2w[[1, 0, 2, 3], :]  # swap y and z
    # c2w[2, :] *= -1  # flip whole world upside down

    # pose = torch.tensor(c2w, dtype=torch.float32)

    return (img, pose)


def _parallel_loader_nerf_image_pose(args):
    torch.set_num_threads(1)
    return _load_nerf_image_pose(**args)

def _parallel_loader_humman_image_pose(args):
    torch.set_num_threads(1)
    return _load_humman_image_pose(**args)


def _load_video_1cam(idx: int,
                     paths: List[str],
                     poses: torch.Tensor,
                     out_h: int,
                     out_w: int,
                     load_every: int = 1
                     ):  # -> Tuple[List[torch.Tensor], torch.Tensor, List[int]]:
    filters = [
        ("scale", f"w={out_w}:h={out_h}")
    ]
    all_frames = iio.imread(
        paths[idx], plugin='pyav', format='rgb24', constant_framerate=True, thread_count=2,
        filter_sequence=filters,)
    imgs, timestamps = [], []
    for frame_idx, frame in enumerate(all_frames):
        if frame_idx % load_every != 0:
            continue
        if frame_idx >= 300:  # Only look at the first 10 seconds
            break
        # Frame is np.ndarray in uint8 dtype (H, W, C)
        imgs.append(
            torch.from_numpy(frame)
        )
        timestamps.append(frame_idx)
    imgs = torch.stack(imgs, 0)
    med_img, _ = torch.median(imgs, dim=0)  # [h, w, 3]
    return (imgs,
            poses[idx].expand(len(timestamps), -1, -1),
            med_img,
            torch.tensor(timestamps, dtype=torch.int32))


def _parallel_loader_video(args):
    torch.set_num_threads(1)
    return _load_video_1cam(**args)


def parallel_load_images(tqdm_title,
                         dset_type: str,
                         num_images: int,
                         **kwargs) -> List[Any]:
    max_threads = 10
    if dset_type == 'llff':
        fn = _parallel_loader_llff_image
    elif dset_type == 'synthetic':
        fn = _parallel_loader_nerf_image_pose
    elif dset_type == 'phototourism':
        fn = _parallel_loader_phototourism_image
    elif dset_type == 'video':
        fn = _parallel_loader_video
        # giac: Can increase to e.g. 10 if loading 4x subsampled images. Otherwise OOM.
        max_threads = 8
    elif dset_type == 'humman':
        fn = _parallel_loader_humman_image_pose
    else:
        raise ValueError(dset_type)
    p = Pool(min(max_threads, num_images))

    iterator = p.imap(fn, [{"idx": i, **kwargs} for i in range(num_images)])
    outputs = []
    for _ in tqdm(range(num_images), desc=tqdm_title):
        out = next(iterator)
        if out is not None:
            outputs.append(out)

    return outputs

magnerf/datasets/data_loading.py
3d-motion-magnification/3d-motion-mag
from dataclasses import dataclass


@dataclass
class Intrinsics:
    width: int
    height: int
    focal_x: float
    focal_y: float
    center_x: float
    center_y: float

    def scale(self, factor: float):
        nw = round(self.width * factor)
        nh = round(self.height * factor)
        sw = nw / self.width
        sh = nh / self.height
        self.focal_x *= sw
        self.focal_y *= sh
        self.center_x *= sw
        self.center_y *= sh
        self.width = int(nw)
        self.height = int(nh)

    def __repr__(self):
        return (f"Intrinsics(width={self.width}, height={self.height}, "
                f"focal_x={self.focal_x}, focal_y={self.focal_y}, "
                f"center_x={self.center_x}, center_y={self.center_y})")

magnerf/datasets/intrinsics.py
3d-motion-magnification/3d-motion-mag
import os
import sys
from typing import Optional
import numpy as np
from PIL import Image
import torch

from .ray_utils import (
    generate_interpolated_path,
    create_meshgrid,
    stack_camera_dirs,
)
from .base_dataset import BaseDataset
from .intrinsics import Intrinsics

_PATH = os.path.abspath(__file__)
sys.path.insert(0, os.path.join(os.path.dirname(_PATH), ".", "pycolmap", "pycolmap"))
from scene_manager import SceneManager

def similarity_from_cameras(c2w, strict_scaling):
    """
    reference: nerf-factory
    Get a similarity transform to normalize dataset
    from c2w (OpenCV convention) cameras
    :param c2w: (N, 4)
    :return T (4,4) , scale (float)
    """
    t = c2w[:, :3, 3]
    R = c2w[:, :3, :3]

    # (1) Rotate the world so that z+ is the up axis
    # we estimate the up axis by averaging the camera up axes
    ups = np.sum(R * np.array([0, -1.0, 0]), axis=-1)
    world_up = np.mean(ups, axis=0)
    world_up /= np.linalg.norm(world_up)

    up_camspace = np.array([0.0, -1.0, 0.0])
    c = (up_camspace * world_up).sum()
    cross = np.cross(world_up, up_camspace)
    skew = np.array(
        [
            [0.0, -cross[2], cross[1]],
            [cross[2], 0.0, -cross[0]],
            [-cross[1], cross[0], 0.0],
        ]
    )
    if c > -1:
        R_align = np.eye(3) + skew + (skew @ skew) * 1 / (1 + c)
    else:
        # In the unlikely case the original data has y+ up axis,
        # rotate 180-deg about x axis
        R_align = np.array([[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])

    #  R_align = np.eye(3) # DEBUG
    R = R_align @ R
    fwds = np.sum(R * np.array([0, 0.0, 1.0]), axis=-1)
    t = (R_align @ t[..., None])[..., 0]

    # (2) Recenter the scene using camera center rays
    # find the closest point to the origin for each camera's center ray
    nearest = t + (fwds * -t).sum(-1)[:, None] * fwds

    # median for more robustness
    translate = -np.median(nearest, axis=0)

    #  translate = -np.mean(t, axis=0)  # DEBUG

    transform = np.eye(4)
    transform[:3, 3] = translate
    transform[:3, :3] = R_align

    # (3) Rescale the scene using camera distances
    scale_fn = np.max if strict_scaling else np.median
    scale = 1.0 / scale_fn(np.linalg.norm(t + translate, axis=-1))

    return transform, scale


def handheld_load_colmap_temporal(
    cutoff: int,
    data_dir: str,
    factor: int = 1,
    hold_every: int = 8,
    timestep: str = "001",
):
    colmap_dir = os.path.join(data_dir, "sparse/0/")

    manager = SceneManager(colmap_dir)
    manager.load_cameras()
    manager.load_images()
    manager.load_points3D()
    pts_ret = manager.get_filtered_points3D(return_colors=True)

    # Assume shared intrinsics between all cameras.
    cam = manager.cameras[1]
    fx, fy, cx, cy = cam.fx, cam.fy, cam.cx, cam.cy
    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    K[:2, :] /= factor

    intrinsics = Intrinsics(
        width=int(cam.width / factor),
        height=int(cam.height / factor),
        focal_x=fx / factor,
        focal_y=fy / factor,
        center_x=cx / factor,
        center_y=cy / factor,
    )

    # Extract extrinsic matrices in world-to-camera format.
    imdata = manager.images
    w2c_mats = []
    bottom = np.array([0, 0, 0, 1]).reshape(1, 4)
    for k in imdata:
        im = imdata[k]
        rot = im.R()
        trans = im.tvec.reshape(3, 1)
        w2c = np.concatenate([np.concatenate([rot, trans], 1), bottom], axis=0)
        w2c_mats.append(w2c)
    w2c_mats = np.stack(w2c_mats, axis=0)

    # Convert extrinsics to camera-to-world.
    camtoworlds = np.linalg.inv(w2c_mats)

    # Image names from COLMAP. No need for permuting the poses according to
    # image names anymore.
    image_names = [imdata[k].name for k in imdata]
    # Get distortion parameters.
    type_ = cam.camera_type

    if type_ == 0 or type_ == "SIMPLE_PINHOLE":
        params = None

    elif type_ == 1 or type_ == "PINHOLE":
        params = None

    if type_ == 2 or type_ == "SIMPLE_RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1

    elif type_ == 3 or type_ == "RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2

    elif type_ == 4 or type_ == "OPENCV":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["p1"] = cam.p1
        params["p2"] = cam.p2

    elif type_ == 5 or type_ == "OPENCV_FISHEYE":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "k4"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["k3"] = cam.k3
        params["k4"] = cam.k4

    # Previous Nerf results were generated with images sorted by filename,
    # ensure metrics are reported on the same test set.
    inds = np.argsort(image_names)
    image_names = [image_names[i] for i in inds]
    camtoworlds = camtoworlds[inds]

    colmap_image_dir = os.path.join(data_dir, "images")
    image_dir = os.path.join(data_dir, "images")
    for d in [image_dir, colmap_image_dir]:
        if not os.path.exists(d):
            raise ValueError(f"Image folder {d} does not exist.")
    # Downsampled images may have different names vs images used for COLMAP,
    # so we need to map between the two sorted lists of files.
    colmap_files = sorted(os.listdir(colmap_image_dir))
    image_files = sorted(os.listdir(image_dir))
    colmap_to_image = dict(zip(colmap_files, image_files))
    image_paths = [os.path.join(image_dir, colmap_to_image[f]) for f in image_names]
    images = []
    for x in image_paths:
        pil_image = Image.open(x).convert("RGB")
        if factor > 1:
            width, height = pil_image.size
            newsize = (int(width / factor), int(height / factor))
            pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
        images.append(np.array(pil_image, dtype="uint8"))
    images = np.stack(images, axis=0)

    if cutoff == -1:
        cutoff = 1

    all_indices = np.array([cutoff - 1 + int(timestep)])
    split_indices = {
        "test": all_indices,
        "train": all_indices,
    }
    return images, camtoworlds, intrinsics, split_indices, pts_ret


def handheld_load_colmap(
    cutoff: int, data_dir: str, factor: int = 1, hold_every: int = 8
):
    assert factor in [1, 2, 4, 8]

    colmap_dir = os.path.join(data_dir, "sparse/0/")

    manager = SceneManager(colmap_dir)
    manager.load_cameras()
    manager.load_images()
    manager.load_points3D()
    pts_ret = manager.get_filtered_points3D(return_colors=True)

    # Assume shared intrinsics between all cameras.
    cam = manager.cameras[1]
    fx, fy, cx, cy = cam.fx, cam.fy, cam.cx, cam.cy
    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    K[:2, :] /= factor

    intrinsics = Intrinsics(
        width=int(cam.width / factor),
        height=int(cam.height / factor),
        focal_x=fx / factor,
        focal_y=fy / factor,
        center_x=cx / factor,
        center_y=cy / factor,
    )

    # Extract extrinsic matrices in world-to-camera format.
    imdata = manager.images
    w2c_mats = []
    bottom = np.array([0, 0, 0, 1]).reshape(1, 4)
    for k in imdata:
        im = imdata[k]
        rot = im.R()
        trans = im.tvec.reshape(3, 1)
        w2c = np.concatenate([np.concatenate([rot, trans], 1), bottom], axis=0)
        w2c_mats.append(w2c)
    w2c_mats = np.stack(w2c_mats, axis=0)

    # Convert extrinsics to camera-to-world.
    camtoworlds = np.linalg.inv(w2c_mats)

    # Image names from COLMAP. No need for permuting the poses according to
    # image names anymore.
    image_names = [imdata[k].name for k in imdata]

    # # Switch from COLMAP (right, down, fwd) to Nerf (right, up, back) frame.
    # poses = poses @ np.diag([1, -1, -1, 1])

    # Get distortion parameters.
    type_ = cam.camera_type

    if type_ == 0 or type_ == "SIMPLE_PINHOLE":
        params = None
        camtype = "perspective"

    elif type_ == 1 or type_ == "PINHOLE":
        params = None
        camtype = "perspective"

    if type_ == 2 or type_ == "SIMPLE_RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1

    elif type_ == 3 or type_ == "RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2

    elif type_ == 4 or type_ == "OPENCV":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["p1"] = cam.p1
        params["p2"] = cam.p2

    elif type_ == 5 or type_ == "OPENCV_FISHEYE":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "k4"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["k3"] = cam.k3
        params["k4"] = cam.k4

    # Previous Nerf results were generated with images sorted by filename,
    # ensure metrics are reported on the same test set.
    inds = np.argsort(image_names)
    if cutoff == -1:
        cutoff = len(inds)
    inds = inds[:cutoff]

    image_names = [image_names[i] for i in inds]
    camtoworlds = camtoworlds[inds]

    colmap_image_dir = os.path.join(data_dir, "images")
    image_dir = os.path.join(data_dir, "images")
    for d in [image_dir, colmap_image_dir]:
        if not os.path.exists(d):
            raise ValueError(f"Image folder {d} does not exist.")
    # Downsampled images may have different names vs images used for COLMAP,
    # so we need to map between the two sorted lists of files.
    colmap_files = sorted(os.listdir(colmap_image_dir))
    image_files = sorted(os.listdir(image_dir))
    colmap_to_image = dict(zip(colmap_files, image_files))
    image_paths = [os.path.join(image_dir, colmap_to_image[f]) for f in image_names]

    images = []
    for x in image_paths:
        pil_image = Image.open(x).convert("RGB")
        if factor > 1:
            width, height = pil_image.size
            newsize = (int(width / factor), int(height / factor))
            pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
        images.append(np.array(pil_image, dtype="uint8"))

    images = np.stack(images, axis=0)

    # Select the split.
    all_indices = np.arange(images.shape[0])
    split_indices = {
        "test": all_indices[all_indices % hold_every == 0],
        "train": all_indices,
    }
    return images, camtoworlds, intrinsics, split_indices, pts_ret


class HandheldDataset(BaseDataset):
    def __init__(
        self,
        datadir,
        split: str,
        batch_size: Optional[int] = None,
        downsample: int = 4,
        hold_every: int = 8,
        contraction: bool = False,
        train_cutoff: int = -1,
        dynamic_start: int = 0,
        near_plane_offset: float = 0.0,
        start_interp_index: int = 5,
        aabb=[-1.0, -1.0, -1.0, 1.0, 1.0, 1.0],
        use_app: bool = False,
        near_plane: float = 0.2,
        far_plane: float = 30.0,
        bg_color_aug: str = "random",
    ):
        self.downsample = downsample
        self.hold_every = hold_every
        self.bg_color_aug = bg_color_aug
        if not isinstance(contraction, bool):
            contraction = eval(contraction)

        data_dir = "/".join(datadir.split("/")[:-1])
        self.timestep = datadir.split("/")[-1]
        scene_name = datadir.split("/")[-2]

        self.mask = None
        self.per_view_near_far = False

        if split == "render":
            if self.timestep != "000":
                imgs, poses, intrinsics, split_indices, _ = handheld_load_colmap_temporal(
                    train_cutoff,
                    data_dir,
                    downsample,
                    hold_every,
                    self.timestep,
                )
            else:
                imgs, poses, intrinsics, split_indices, _ = handheld_load_colmap(
                    train_cutoff,
                    data_dir,
                    downsample,
                    hold_every,
                )
            T, sscale = similarity_from_cameras(poses[:train_cutoff], strict_scaling=False)
            poses = np.einsum("nij, ki -> nkj", poses, T)
            poses[:, :3, 3] *= sscale
            poses = torch.from_numpy(poses).float()

            sta_id = start_interp_index
            end_id = len(poses) - 1

            render_poses_0 = poses[sta_id : sta_id + 1, :3].repeat(30, 1, 1)
            render_poses_1 = generate_interpolated_path(
                torch.cat([render_poses_0[0:1], poses[end_id : end_id + 1, :3]]),
                n_interp=30,
                look_at_neg=False,
            )
            render_poses_2 = poses[end_id : end_id + 1, :3].repeat(30, 1, 1)
            render_poses_3 = generate_interpolated_path(
                torch.cat([poses[end_id : end_id + 1, :3], render_poses_0[0:1]]),
                n_interp=30,
                look_at_neg=False,
            )
            self.poses = torch.cat([
                 render_poses_0, render_poses_0,
                 render_poses_1, render_poses_2,
                 render_poses_2, render_poses_3,], dim=0,)

            if scene_name == "ronaldo":
                i_a = 4
                i_b = 92
                pose_a = poses[i_a : i_a + 1, :3]
                pose_b = poses[i_b : i_b + 1, :3]
                a2b = generate_interpolated_path(
                    poses[i_a : i_b + 1, :3],
                    n_points=30,
                    look_at_neg=False,
                    n_interp=-1,
                )
                b2a = torch.flip(a2b, [0])
                self.poses = torch.cat(
                    [pose_b.repeat(150, 1, 1), b2a, pose_a.repeat(150, 1, 1), a2b],
                    dim=0,
                )

            if scene_name == "jackson":
                i_a = 40
                i_b = 70
                pose_a = poses[i_a : i_a + 1, :3]
                pose_b = poses[i_b : i_b + 1, :3]
                a2b = generate_interpolated_path(
                    poses[i_a : i_b + 1, :3],
                    n_points=30,
                    look_at_neg=False,
                    n_interp=-1,
                )
                b2a = torch.flip(a2b, [0])
                self.poses = torch.cat(
                    [pose_a.repeat(120, 1, 1), a2b, pose_b.repeat(120, 1, 1), b2a],
                    dim=0,
                )

            if scene_name == "baby":
                i_a = 55
                i_b = 214
                pose_a = poses[i_a : i_a + 1, :3]
                pose_b = poses[i_b : i_b + 1, :3]
                a2b = generate_interpolated_path(
                    poses[i_a : i_b + 1, :3],
                    n_points=30,
                    look_at_neg=False,
                    n_interp=-1,
                )
                b2a = torch.flip(a2b, [0])
                self.poses = torch.cat(
                    [pose_a.repeat(150, 1, 1), a2b, pose_b.repeat(150, 1, 1), b2a],
                    dim=0,
                )
            imgs = None
        else:
            if self.timestep != "000" and split == "train":
                imgs, poses, intrinsics, split_indices, _ = handheld_load_colmap_temporal(
                    train_cutoff,
                    data_dir,
                    downsample,
                    hold_every,
                    self.timestep,
                )
            else:
                imgs, poses, intrinsics, split_indices, _ = handheld_load_colmap(
                    train_cutoff,
                    data_dir,
                    downsample,
                    hold_every,
                )
            if train_cutoff == -1:
                train_cutoff = len(poses)

            T, sscale = similarity_from_cameras(poses[:train_cutoff], strict_scaling=False)
            poses = np.einsum("nij, ki -> nkj", poses, T)
            poses[:, :3, 3] *= sscale
            self.poses = torch.from_numpy(poses).float()

            indices = split_indices[split]
            if self.timestep != "000":
                indices = [dynamic_start + int(self.timestep)]

            imgs = imgs[indices]
            self.poses = self.poses[indices]
            imgs = torch.from_numpy(imgs).to(torch.uint8)

            if split == "train":
                imgs = imgs.view(-1, imgs.shape[-1])
            else:
                imgs = imgs.view(-1, intrinsics.height * intrinsics.width, imgs.shape[-1])

        self.num_images = train_cutoff
        self.dynamic_start = dynamic_start

        self.near_fars = torch.tensor([near_plane, far_plane])
        if split != "train":
            self.near_fars[0] += near_plane_offset

        if scene_name == "jackson":
            self.per_view_near_far = True
            self.near_fars = self.near_fars.repeat(len(self.poses), 1)
            self.near_fars[:, 0] = torch.tensor([1.0]).view(1, 1) - self.poses[:, 2, 3]

        self.scene_name = scene_name
        self.pose_opt = False

        self.global_translation = torch.tensor([0.0, 0.0, 0.0])
        self.global_scale = aabb
        bbox = torch.tensor([aabb[:3], aabb[3:]])

        if type(use_app) == str:
            use_app = bool(eval(use_app))
        self.use_app_emb = use_app
        self.OPENGL_CAMERA = False

        super().__init__(
            datadir=datadir,
            split=split,
            scene_bbox=bbox,
            batch_size=batch_size,
            imgs=imgs,
            rays_o=None,
            rays_d=None,
            intrinsics=intrinsics,
            is_ndc=False,
            is_contracted=contraction,
        )

    def __getitem__(self, index):
        h = self.intrinsics.height
        w = self.intrinsics.width
        dev = "cpu"
        out = {}

        if self.split == "train":
            index = self.get_rand_ids(index)
            if self.mask is not None:
                mask = self.mask[index]
                index = index[mask]
                if len(index) == 0:
                    index = self.get_rand_ids(index)

            image_id = torch.div(index, h * w, rounding_mode="floor").long()
            y = torch.remainder(index, h * w).div(w, rounding_mode="floor")
            x = torch.remainder(index, h * w).remainder(w)
            x = x + 0.5
            y = y + 0.5

            if self.bg_color_aug == "random":
                out["bg_color"] = torch.rand_like(torch.tensor([[0.0, 0.0, 0.0]]))
            elif self.bg_color_aug == "white":
                out["bg_color"] = torch.tensor([[1.0, 1.0, 1.0]])
            else:
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])

            if self.per_view_near_far:
                out["near_fars"] = self.near_fars[image_id].view(-1, 2)
            else:
                out["near_fars"] = self.near_fars.view(-1, 2)

        else:
            image_id = torch.LongTensor([index])
            x, y = create_meshgrid(height=h, width=w, dev=dev, add_half=True, flat=True)
            if self.bg_color_aug == "random":
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])
            elif self.bg_color_aug == "white":
                out["bg_color"] = torch.tensor([[1.0, 1.0, 1.0]])
            else:
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])

            if self.per_view_near_far:
                out["near_fars"] = self.near_fars[image_id].view(-1, 2)
            else:
                out["near_fars"] = self.near_fars.view(-1, 2)

        if self.imgs is not None:
            out["imgs"] = (
                self.imgs[index] / 255.0
            )  # (num_rays, 3)   this converts to f32
        else:
            out["imgs"] = None

        c2w = self.poses[image_id]  # (num_rays, 3, 4)
        camera_dirs = stack_camera_dirs(
            x, y, self.intrinsics, self.OPENGL_CAMERA
        )  # [num_rays, 3]
        directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
        origins = torch.broadcast_to(c2w[:, :3, -1], directions.shape)
        viewdirs = directions / torch.linalg.norm(directions, dim=-1, keepdims=True)

        if self.pose_opt and self.split != "render":
            out["indices"] = (
                image_id
                if self.timestep == "000"
                else self.dynamic_start * torch.ones_like(image_id)
            )
            out["camera_dirs"] = camera_dirs
            out["c2w"] = c2w

        out["rays_o"] = origins
        out["rays_d"] = viewdirs

        if self.use_app_emb:
            out["timestamps"] = (image_id + 1) % self.num_images
            if self.timestep != "000":
                out["timestamps"] = (
                    self.dynamic_start * torch.ones_like(image_id) + 1
                ) % self.num_images

        return out

magnerf/datasets/nerf_360_v2.py
3d-motion-magnification/3d-motion-mag
from typing import Tuple, Optional

import numpy as np
import torch
import scipy

from .intrinsics import Intrinsics

__all__ = (
    "create_meshgrid",
    "get_ray_directions",
    "stack_camera_dirs",
    "get_rays",
    "ndc_rays_blender",
    "center_poses",
    "generate_spiral_path",
    "generate_hemispherical_orbit",
    "generate_spherical_poses",
    "normalize",
    "average_poses",
    "viewmatrix",
)


def create_meshgrid(height: int,
                    width: int,
                    dev: str = 'cpu',
                    add_half: bool = True,
                    flat: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
    xs = torch.arange(width, dtype=torch.float32, device=dev)
    ys = torch.arange(height, dtype=torch.float32, device=dev)
    if add_half:
        xs += 0.5
        ys += 0.5
    # generate grid by stacking coordinates
    yy, xx = torch.meshgrid([ys, xs], indexing="ij")  # both HxW
    if flat:
        return xx.flatten(), yy.flatten()
    return xx, yy


def stack_camera_dirs(x: torch.Tensor, y: torch.Tensor, intrinsics: Intrinsics, opengl_camera: bool):
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    x = x.float()
    y = y.float()
    return torch.stack([
        (x - intrinsics.center_x) / intrinsics.focal_x,
        (y - intrinsics.center_y) / intrinsics.focal_y
        * (-1.0 if opengl_camera else 1.0),
        torch.full_like(x, fill_value=-1.0 if opengl_camera else 1.0)
    ], -1)  # (H, W, 3)


def get_ray_directions(intrinsics: Intrinsics, opengl_camera: bool, add_half: bool = True) -> torch.Tensor:
    """
    Get ray directions for all pixels in camera coordinate.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems
    Inputs:
        height, width, focal: image height, width and focal length
    Outputs:
        directions: (height, width, 3), the direction of the rays in camera coordinate
    """
    xx, yy = create_meshgrid(intrinsics.height, intrinsics.width, add_half=add_half)

    return stack_camera_dirs(xx, yy, intrinsics, opengl_camera)


def get_rays(directions: torch.Tensor,
             c2w: torch.Tensor,
             ndc: bool,
             ndc_near: float = 1.0,
             intrinsics: Optional[Intrinsics] = None,
             normalize_rd: bool = True):
    """Get ray origin and normalized directions in world coordinate for all pixels in one image.
    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Args:
        directions:
        c2w:
        ndc:
        ndc_near:
        intrinsics:
        normalize_rd:

    Returns:

    """
    directions = directions.view(-1, 3)  # [n_rays, 3]
    if len(c2w.shape) == 2:
        c2w = c2w[None, ...]
    rd = (directions[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
    ro = torch.broadcast_to(c2w[:, :3, 3], directions.shape)
    if ndc:
        assert intrinsics is not None, "intrinsics must not be None when NDC active."
        ro, rd = ndc_rays_blender(
            intrinsics=intrinsics, near=ndc_near, rays_o=ro, rays_d=rd)
    if normalize_rd:
        rd /= torch.linalg.norm(rd, dim=-1, keepdim=True)
    return ro, rd


def ndc_rays_blender(intrinsics: Intrinsics, near: float, rays_o: torch.Tensor,
                     rays_d: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # Shift ray origins to near plane
    t = -(near + rays_o[..., 2]) / rays_d[..., 2]
    rays_o = rays_o + t[..., None] * rays_d

    # Projection
    ndc_coef_x = - (2 * intrinsics.focal_x) / intrinsics.width
    ndc_coef_y = - (2 * intrinsics.focal_y) / intrinsics.height
    o0 = ndc_coef_x * rays_o[..., 0] / rays_o[..., 2]
    o1 = ndc_coef_y * rays_o[..., 1] / rays_o[..., 2]
    o2 = 1. + 2. * near / rays_o[..., 2]

    d0 = ndc_coef_x * (rays_d[..., 0] / rays_d[..., 2] - rays_o[..., 0] / rays_o[..., 2])
    d1 = ndc_coef_y * (rays_d[..., 1] / rays_d[..., 2] - rays_o[..., 1] / rays_o[..., 2])
    d2 = -2. * near / rays_o[..., 2]

    rays_o = torch.stack([o0, o1, o2], -1)
    rays_d = torch.stack([d0, d1, d2], -1)

    return rays_o, rays_d


def normalize(v: np.ndarray) -> np.ndarray:
    """Normalize a vector."""
    return v / np.linalg.norm(v)


def average_poses(poses: np.ndarray) -> np.ndarray:
    """
    Calculate the average pose, which is then used to center all poses
    using @center_poses. Its computation is as follows:
    1. Compute the center: the average of pose centers.
    2. Compute the z axis: the normalized average z axis.
    3. Compute axis y': the average y axis.
    4. Compute x' = y' cross product z, then normalize it as the x axis.
    5. Compute the y axis: z cross product x.

    Note that at step 3, we cannot directly use y' as y axis since it's
    not necessarily orthogonal to z axis. We need to pass from x to y.
    Inputs:
        poses: (N_images, 3, 4)
    Outputs:
        pose_avg: (3, 4) the average pose
    """
    # 1. Compute the center
    center = poses[..., 3].mean(0)  # (3)
    # 2. Compute the z axis
    z = normalize(poses[..., 2].mean(0))  # (3)
    # 3. Compute axis y' (no need to normalize as it's not the final output)
    y_ = poses[..., 1].mean(0)  # (3)
    # 4. Compute the x axis
    x = normalize(np.cross(z, y_))  # (3)
    # 5. Compute the y axis (as z and x are normalized, y is already of norm 1)
    y = np.cross(x, z)  # (3)

    pose_avg = np.stack([x, y, z, center], 1)  # (3, 4)

    return pose_avg


def center_poses(poses):
    """
    Center the poses so that we can use NDC.
    See https://github.com/bmild/nerf/issues/34
    Inputs:
        poses: (N_images, 3, 4)
    Outputs:
        poses_centered: (N_images, 3, 4) the centered poses
        pose_avg: (3, 4) the average pose
    """
    pose_avg = average_poses(poses)  # (3, 4)
    pose_avg_homo = np.eye(4)
    pose_avg_homo[:3] = pose_avg  # convert to homogeneous coordinate for faster computation
    pose_avg_homo = pose_avg_homo
    # by simply adding 0, 0, 0, 1 as the last row
    last_row = np.tile(np.array([0, 0, 0, 1]), (len(poses), 1, 1))  # (N_images, 1, 4)
    poses_homo = \
        np.concatenate([poses, last_row], 1)  # (N_images, 4, 4) homogeneous coordinate

    poses_centered = np.linalg.inv(pose_avg_homo) @ poses_homo  # (N_images, 4, 4)
    poses_centered = poses_centered[:, :3]  # (N_images, 3, 4)

    return poses_centered, pose_avg_homo


def viewmatrix(z, up, pos):
    vec2 = normalize(z)
    vec1_avg = up
    vec0 = normalize(np.cross(vec1_avg, vec2))
    vec1 = normalize(np.cross(vec2, vec0))
    return np.stack([-vec0, vec1, vec2, pos], axis=1)


def generate_spiral_path(poses: np.ndarray,
                         near_fars: np.ndarray,
                         n_frames=120,
                         n_rots=2,
                         zrate=.5,
                         dt=0.75,
                         percentile=70,
                         center=None) -> np.ndarray:
    """Calculates a forward facing spiral path for rendering.

    From https://github.com/google-research/google-research/blob/342bfc150ef1155c5254c1e6bd0c912893273e8d/regnerf/internal/datasets.py
    and https://github.com/apchenstu/TensoRF/blob/main/dataLoader/llff.py

    :param poses: [N, 3, 4]
    :param near_fars:
    :param n_frames:
    :param n_rots:
    :param zrate:
    :param dt:
    :return:
    """
    # center pose
    if center is not None:
        c2w = center
    else:
        c2w = average_poses(poses)  # [3, 4]

    # Get average pose
    up = normalize(poses[:, :3, 1].sum(0))

    # Find a reasonable "focus depth" for this dataset as a weighted average
    # of near and far bounds in disparity space.
    close_depth, inf_depth = np.min(near_fars) * 1.0, np.max(near_fars) * 5.0
    focal = 1.0 / (((1.0 - dt) / close_depth + dt / inf_depth))

    # Get radii for spiral path using 90th percentile of camera positions.
    positions = poses[:, :3, 3]
    radii = np.percentile(np.abs(positions), percentile, 0)
    radii = np.concatenate([radii, [1.]])

    # Generate poses for spiral path.
    render_poses = []
    for theta in np.linspace(0., 2. * np.pi * n_rots, n_frames, endpoint=False):
        t = radii * [np.cos(theta), -np.sin(theta), -np.sin(theta * zrate), 1.]
        position = c2w @ t
        lookat = c2w @ np.array([0, 0, -focal, 1.0])
        z_axis = normalize(position - lookat)
        render_poses.append(viewmatrix(z_axis, up, position))
    return np.stack(render_poses, axis=0)


def generate_hemispherical_orbit(poses: torch.Tensor, n_frames=120, opengl_cam=True):
    """Calculates a render path which orbits around the z-axis.
    Based on https://github.com/google-research/google-research/blob/342bfc150ef1155c5254c1e6bd0c912893273e8d/regnerf/internal/datasets.py
    """
    origins = poses[:, :3, 3]
    radius = torch.sqrt(torch.mean(torch.sum(origins ** 2, dim=-1)))

    # Assume that z-axis points up towards approximate camera hemisphere
    sin_phi = torch.mean(origins[:, 2], dim=0) / radius + 0.4
    cos_phi = torch.sqrt(1 - sin_phi ** 2)
    render_poses = []

    up = torch.tensor([0., 0., 1.])
    for theta in np.linspace(0., 2. * np.pi, n_frames, endpoint=False):
        camorigin = radius * torch.tensor(
            [cos_phi * np.cos(theta), cos_phi * np.sin(theta), sin_phi])
        render_poses.append(torch.from_numpy(viewmatrix(camorigin, up, camorigin)))

    render_poses = torch.stack(render_poses, dim=0)

    return render_poses


trans_t = lambda t: torch.Tensor([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, t],
    [0, 0, 0, 1]]).float()

rot_phi = lambda phi: torch.Tensor([
    [1, 0, 0, 0],
    [0, np.cos(phi), -np.sin(phi), 0],
    [0, np.sin(phi), np.cos(phi), 0],
    [0, 0, 0, 1]]).float()

rot_theta = lambda th: torch.Tensor([
    [np.cos(th), 0, -np.sin(th), 0],
    [0, 1, 0, 0],
    [np.sin(th), 0, np.cos(th), 0],
    [0, 0, 0, 1]]).float()


def generate_spherical_poses(theta, phi, radius):
    c2w = trans_t(radius)
    c2w = rot_phi(phi / 180. * np.pi) @ c2w
    c2w = rot_theta(theta / 180. * np.pi) @ c2w
    c2w = torch.Tensor(np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])) @ c2w
    return c2w


def focus_point_fn(poses: np.ndarray) -> np.ndarray:
  """Calculate nearest point to all focal axes in poses."""
  directions, origins = poses[:, :3, 2:3], poses[:, :3, 3:4]
  m = np.eye(3) - directions * np.transpose(directions, [0, 2, 1])
  mt_m = np.transpose(m, [0, 2, 1]) @ m
  focus_pt = np.linalg.inv(mt_m.mean(0)) @ (mt_m @ origins).mean(0)[:, 0]
  return focus_pt


def generate_ellipse_path(poses: np.ndarray,
                          n_frames: int = 120,
                          z_variation: float = 0.,
                          z_phase: float = 0.) -> np.ndarray:
  """Generate an elliptical render path based on the given poses."""
  # Calculate the focal point for the path (cameras point toward this).
  center = focus_point_fn(poses)
  # Path height sits at z=0 (in middle of zero-mean capture pattern).
  offset = np.array([center[0], center[1], 0])

  # Calculate scaling for ellipse axes based on input camera positions.
  sc = np.percentile(np.abs(poses[:, :3, 3] - offset), 90, axis=0)
  # Use ellipse that is symmetric about the focal point in xy.
  low = -sc + offset
  high = sc + offset
  # Optional height variation need not be symmetric
  z_low = np.percentile((poses[:, :3, 3]), 10, axis=0)
  z_high = np.percentile((poses[:, :3, 3]), 90, axis=0)

  def get_positions(theta):
    # Interpolate between bounds with trig functions to get ellipse in x-y.
    # Optionally also interpolate in z to change camera height along path.
    return np.stack([
        low[0] + (high - low)[0] * (np.cos(theta) * .5 + .5),
        low[1] + (high - low)[1] * (np.sin(theta) * .5 + .5),
        z_variation * (z_low[2] + (z_high - z_low)[2] *
                       (np.cos(theta + 2 * np.pi * z_phase) * .5 + .5)),
    ], -1)

  theta = np.linspace(0, 2. * np.pi, n_frames + 1, endpoint=True)
  positions = get_positions(theta)

  # Throw away duplicated last position.
  positions = positions[:-1]

  # Set path's up vector to axis closest to average of input pose up vectors.
  avg_up = poses[:, :3, 1].mean(0)
  avg_up = avg_up / np.linalg.norm(avg_up)
  ind_up = np.argmax(np.abs(avg_up))
  up = np.eye(3)[ind_up] * np.sign(avg_up[ind_up])

  return np.stack([viewmatrix(p - center, up, p) for p in positions])

def viewmatrix(lookdir: np.ndarray, up: np.ndarray,
                position: np.ndarray) -> np.ndarray:
    """Construct lookat view matrix."""
    vec2 = normalize(lookdir)
    vec0 = normalize(np.cross(up, vec2))
    vec1 = normalize(np.cross(vec2, vec0))
    m = np.stack([vec0, vec1, vec2, position], axis=1)
    return m

def poses_to_points(poses, dist, look_at_neg):
    """Converts from pose matrices to (position, lookat, up) format."""
    pos = poses[:, :3, -1]
    if look_at_neg:
        lookat = poses[:, :3, -1] + dist * poses[:, :3, 2]
        up = poses[:, :3, -1] - dist * poses[:, :3, 1]
    else:
        lookat = poses[:, :3, -1] - dist * poses[:, :3, 2]
        up = poses[:, :3, -1] + dist * poses[:, :3, 1]
    return np.stack([pos, lookat, up], 1)


def points_to_poses(points):
    """Converts from (position, lookat, up) format to pose matrices."""
    pose = np.array([viewmatrix(p - l, u - p, p) for p, l, u in points])
    return torch.from_numpy(pose).float()


def interp(points, n, k, s):
    """Runs multidimensional B-spline interpolation on the input points."""
    sh = points.shape
    pts = np.reshape(points, (sh[0], -1))
    k = min(k, sh[0] - 1)
    tck, _ = scipy.interpolate.splprep(pts.T, k=k, s=s)
    u = np.linspace(0, 1, n, endpoint=False)
    new_points = np.array(scipy.interpolate.splev(u, tck))
    new_points = np.reshape(new_points.T, (n, sh[1], sh[2]))
    return new_points


def rotx(θ):
    """ Returns a rotation matrix for the given angle about the X axis """
    return np.array(((1,0,0,0),
                    (0,np.cos(θ),-np.sin(θ),0),
                    (0,np.sin(θ),np.cos(θ),0),
                    (0,0,0,1)))

def roty(θ):
    """ Returns a rotation matrix for the given angle about the X axis """
    return np.array(((np.cos(θ),0,np.sin(θ),0),
                    (0,1,0,0),
                    (-np.sin(θ),0,np.cos(θ),0),
                    (0,0,0,1)))

def rotz(θ):
    """ Returns a rotation matrix for the given angle about the X axis """
    return np.array(((np.cos(θ), -np.sin(θ),0,0),
                    (np.sin(θ), np.cos(θ),0,0),
                    (0,0,1,0),
                    (0,0,0,1)))

def trans(x=0, y=0, z=0):
    """ Returns a translation matrix for the given offset """
    return np.array(((1,0,0,x),
                    (0,1,0,y),
                    (0,0,1,z),
                    (0,0,0,1)))


def generate_interpolated_path(poses: np.ndarray,
                               n_interp: int,
                               spline_degree: int = 5,
                               smoothness: float = .03,
                               rot_weight: float = .1,
                               look_at_neg: bool = False,
                               n_points: int = None):
  """Creates a smooth spline path between input keyframe camera poses.
  Spline is calculated with poses in format (position, lookat-point, up-point).
  Args:
    poses: (n, 3, 4) array of input pose keyframes.
    n_interp: returned path will have n_interp * (n - 1) total poses.
    spline_degree: polynomial degree of B-spline.
    smoothness: parameter for spline smoothing, 0 forces exact interpolation.
    rot_weight: relative weighting of rotation/translation in spline solve.
  Returns:
    Array of new camera poses with shape (n_interp * (n - 1), 3, 4).
  """
  if n_points is None:
    n_points = n_interp * (poses.shape[0] - 1)
  points = poses_to_points(poses, dist=rot_weight, look_at_neg=look_at_neg)
  new_points = interp(points,
                      n_points,
                      k=spline_degree,
                      s=smoothness)
  return points_to_poses(new_points)



import struct
def write_pointcloud(filename,xyz_points,rgb_points=None):
    assert xyz_points.shape[1] == 3,'Input XYZ points should be Nx3 float array'
    if rgb_points is None:
        rgb_points = np.ones(xyz_points.shape).astype(np.uint8)*255
    assert xyz_points.shape == rgb_points.shape,'Input RGB colors should be Nx3 float array and have same size as input XYZ points'

    # Write header of .ply file
    fid = open(filename,'wb')
    fid.write(bytes('ply\n', 'utf-8'))
    fid.write(bytes('format binary_little_endian 1.0\n', 'utf-8'))
    fid.write(bytes('element vertex %d\n'%xyz_points.shape[0], 'utf-8'))
    fid.write(bytes('property float x\n', 'utf-8'))
    fid.write(bytes('property float y\n', 'utf-8'))
    fid.write(bytes('property float z\n', 'utf-8'))
    fid.write(bytes('property uchar red\n', 'utf-8'))
    fid.write(bytes('property uchar green\n', 'utf-8'))
    fid.write(bytes('property uchar blue\n', 'utf-8'))
    fid.write(bytes('end_header\n', 'utf-8'))

    # Write 3D points to .ply file
    for i in range(xyz_points.shape[0]):
        fid.write(bytearray(struct.pack("fffccc",xyz_points[i,0],xyz_points[i,1],xyz_points[i,2], rgb_points[i,0].tobytes(),rgb_points[i,1].tobytes(), rgb_points[i,2].tobytes())))
    fid.close()


magnerf/datasets/ray_utils.py
3d-motion-magnification/3d-motion-mag
import json
import logging as log
import os
from typing import Tuple, Optional, Any

import numpy as np
import torch

from .data_loading import parallel_load_images
from .ray_utils import get_ray_directions, generate_hemispherical_orbit, get_rays
from .intrinsics import Intrinsics
from .base_dataset import BaseDataset


class SyntheticNerfDataset(BaseDataset):
    def __init__(self,
                 datadir,
                 split: str,
                 batch_size: Optional[int] = None,
                 downsample: float = 1.0,
                 max_frames: Optional[int] = None,
                 radius: float = 1.3,
                 near_far: Optional[list] = None,):
        self.downsample = downsample
        self.max_frames = max_frames
        if near_far is not None:
            self.near_far = near_far
        else:
            self.near_far = [2.0, 6.0]
        if split == 'render':
            frames, transform = load_360_frames(datadir, 'test', self.max_frames)
            imgs, poses = load_360_images(frames, datadir, 'test', self.downsample)
            render_poses = generate_hemispherical_orbit(poses, n_frames=120)
            self.poses = torch.cat([render_poses[30:31].repeat(30, 1, 1), render_poses[30:90], render_poses[90:91].repeat(30, 1, 1)])
            intrinsics = load_360_intrinsics(
                transform, img_h=imgs[0].shape[0], img_w=imgs[0].shape[1],
                downsample=self.downsample)
            imgs = None
        else:
            frames, transform = load_360_frames(datadir, split, self.max_frames)
            imgs, poses = load_360_images(frames, datadir, split, self.downsample)
            self.poses = poses
            intrinsics = load_360_intrinsics(
                transform, img_h=imgs[0].shape[0], img_w=imgs[0].shape[1],
                downsample=self.downsample)
        rays_o, rays_d, imgs = create_360_rays(
            imgs, self.poses, merge_all=split == 'train', intrinsics=intrinsics)

        super().__init__(
            datadir=datadir,
            split=split,
            scene_bbox=get_360_bbox(datadir, is_contracted=False, radius=radius),
            is_ndc=False,
            is_contracted=False,
            batch_size=batch_size,
            imgs=imgs,
            rays_o=rays_o,
            rays_d=rays_d,
            intrinsics=intrinsics,
        )
        log.info(f"SyntheticNerfDataset. Loaded {split} set from {datadir}."
                 f"{len(poses)} images of shape {self.img_h}x{self.img_w}. "
                 f"Images loaded: {imgs is not None}. "
                 f"Sampling without replacement={self.use_permutation}. {intrinsics}")

    def __getitem__(self, index):
        out = super().__getitem__(index)
        pixels = out["imgs"]
        if self.split == 'train':
            bg_color = torch.rand((1, 3), dtype=pixels.dtype, device=pixels.device)
        else:
            if pixels is None:
                bg_color = torch.ones((1, 3), dtype=torch.float32, device='cuda:0')
            else:
                bg_color = torch.ones((1, 3), dtype=pixels.dtype, device=pixels.device)
        # Alpha compositing
        if pixels is not None and pixels.shape[-1] == 4:
            pixels = pixels[:, :3] * pixels[:, 3:] + bg_color * (1.0 - pixels[:, 3:])
        out["imgs"] = pixels
        out["bg_color"] = bg_color
        out["near_fars"] = torch.tensor([[2.0, 6.0]])
        return out


def get_360_bbox(datadir, is_contracted=False, radius=1.3):
    if is_contracted:
        radius = 2
    elif "ship" in datadir:
        radius = 1.5
    else:
        radius = radius
    return torch.tensor([[-radius, -radius, -radius], [radius, radius, radius]])


def create_360_rays(
              imgs: Optional[torch.Tensor],
              poses: torch.Tensor,
              merge_all: bool,
              intrinsics: Intrinsics) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    directions = get_ray_directions(intrinsics, opengl_camera=True)  # [H, W, 3]
    num_frames = poses.shape[0]

    all_rays_o, all_rays_d = [], []
    for i in range(num_frames):
        rays_o, rays_d = get_rays(directions, poses[i], ndc=False, normalize_rd=True)  # h*w, 3
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)

    all_rays_o = torch.cat(all_rays_o, 0).to(dtype=torch.float32)  # [n_frames * h * w, 3]
    all_rays_d = torch.cat(all_rays_d, 0).to(dtype=torch.float32)  # [n_frames * h * w, 3]
    if imgs is not None:
        imgs = imgs.view(-1, imgs.shape[-1]).to(dtype=torch.float32)   # [N*H*W, 3/4]
    if not merge_all:
        num_pixels = intrinsics.height * intrinsics.width
        if imgs is not None:
            imgs = imgs.view(num_frames, num_pixels, -1)  # [N, H*W, 3/4]
        all_rays_o = all_rays_o.view(num_frames, num_pixels, 3)  # [N, H*W, 3]
        all_rays_d = all_rays_d.view(num_frames, num_pixels, 3)  # [N, H*W, 3]
    return all_rays_o, all_rays_d, imgs


def load_360_frames(datadir, split, max_frames: int) -> Tuple[Any, Any]:
    with open(os.path.join(datadir, f"transforms_{split}.json"), 'r') as f:
        meta = json.load(f)
        frames = meta['frames']

        # Subsample frames
        tot_frames = len(frames)
        num_frames = min(tot_frames, max_frames or tot_frames)
        if split == 'train' or split == 'test':
            subsample = int(round(tot_frames / num_frames))
            frame_ids = np.arange(tot_frames)[::subsample]
            if subsample > 1:
                log.info(f"Subsampling {split} set to 1 every {subsample} images.")
        else:
            frame_ids = np.arange(num_frames)
        frames = np.take(frames, frame_ids).tolist()
    return frames, meta


def load_360_images(frames, datadir, split, downsample) -> Tuple[torch.Tensor, torch.Tensor]:
    img_poses = parallel_load_images(
        dset_type="synthetic",
        tqdm_title=f'Loading {split} data',
        num_images=len(frames),
        frames=frames,
        data_dir=datadir,
        out_h=None,
        out_w=None,
        downsample=downsample,
    )
    imgs, poses = zip(*img_poses)
    imgs = torch.stack(imgs, 0)  # [N, H, W, 3/4]
    poses = torch.stack(poses, 0)  # [N, ????]
    return imgs, poses


def load_360_intrinsics(transform, img_h, img_w, downsample) -> Intrinsics:
    height = img_h
    width = img_w
    # load intrinsics
    if 'fl_x' in transform or 'fl_y' in transform:
        fl_x = (transform['fl_x'] if 'fl_x' in transform else transform['fl_y']) / downsample
        fl_y = (transform['fl_y'] if 'fl_y' in transform else transform['fl_x']) / downsample
    elif 'camera_angle_x' in transform or 'camera_angle_y' in transform:
        # blender, assert in radians. already downscaled since we use H/W
        fl_x = width / (2 * np.tan(transform['camera_angle_x'] / 2)) if 'camera_angle_x' in transform else None
        fl_y = height / (2 * np.tan(transform['camera_angle_y'] / 2)) if 'camera_angle_y' in transform else None
        if fl_x is None:
            fl_x = fl_y
        if fl_y is None:
            fl_y = fl_x
    else:
        raise RuntimeError('Failed to load focal length, please check the transforms.json!')

    cx = (transform['cx'] / downsample) if 'cx' in transform else (width / 2)
    cy = (transform['cy'] / downsample) if 'cy' in transform else (height / 2)
    return Intrinsics(height=height, width=width, focal_x=fl_x, focal_y=fl_y, center_x=cx, center_y=cy)

magnerf/datasets/synthetic_nerf_dataset.py
3d-motion-magnification/3d-motion-mag
import os
import sys
import logging as log
from typing import Optional

import numpy as np
import torch

from PIL import Image

from .ray_utils import (
    generate_interpolated_path,
    create_meshgrid,
    stack_camera_dirs,
    roty,
    rotx,
    rotz,
    trans,
)
from .base_dataset import BaseDataset
from .intrinsics import Intrinsics

_PATH = os.path.abspath(__file__)

sys.path.insert(0, os.path.join(os.path.dirname(_PATH), ".", "pycolmap", "pycolmap"))
from scene_manager import SceneManager


def _load_colmap_temporal(
    data_dir: str,
    factor: int = 1,
    hold_every: int = 8,
    timestep: str = "001",
    frame_offset: int = 2,
):
    colmap_dir = os.path.join(data_dir, "sparse/0/")

    manager = SceneManager(colmap_dir)
    manager.load_cameras()
    manager.load_images()
    manager.load_points3D()
    pts_ret = manager.get_filtered_points3D(return_colors=True)

    # Assume shared intrinsics between all cameras.
    cam = manager.cameras[1]
    fx, fy, cx, cy = cam.fx, cam.fy, cam.cx, cam.cy
    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    K[:2, :] /= factor

    intrinsics = Intrinsics(
        width=int(cam.width / factor),
        height=int(cam.height / factor),
        focal_x=fx / factor,
        focal_y=fy / factor,
        center_x=cx / factor,
        center_y=cy / factor,
    )

    # Extract extrinsic matrices in world-to-camera format.
    imdata = manager.images
    w2c_mats = []
    bottom = np.array([0, 0, 0, 1]).reshape(1, 4)
    for k in imdata:
        im = imdata[k]
        rot = im.R()
        trans = im.tvec.reshape(3, 1)
        w2c = np.concatenate([np.concatenate([rot, trans], 1), bottom], axis=0)
        w2c_mats.append(w2c)
    w2c_mats = np.stack(w2c_mats, axis=0)

    # Convert extrinsics to camera-to-world.
    camtoworlds = np.linalg.inv(w2c_mats)

    # Image names from COLMAP. No need for permuting the poses according to
    # image names anymore.
    image_names = [imdata[k].name for k in imdata]

    # Get distortion parameters.
    type_ = cam.camera_type

    if type_ == 0 or type_ == "SIMPLE_PINHOLE":
        params = None
        camtype = "perspective"

    elif type_ == 1 or type_ == "PINHOLE":
        params = None
        camtype = "perspective"

    if type_ == 2 or type_ == "SIMPLE_RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        camtype = "perspective"

    elif type_ == 3 or type_ == "RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        camtype = "perspective"

    elif type_ == 4 or type_ == "OPENCV":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["p1"] = cam.p1
        params["p2"] = cam.p2
        camtype = "perspective"

    elif type_ == 5 or type_ == "OPENCV_FISHEYE":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "k4"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["k3"] = cam.k3
        params["k4"] = cam.k4
        camtype = "fisheye"

    # assert params is None, f"Only support pinhole camera model. Params: {params}"

    # Previous Nerf results were generated with images sorted by filename,
    # ensure metrics are reported on the same test set.
    inds = np.argsort(image_names)
    image_names = [image_names[i] for i in inds]
    camtoworlds = camtoworlds[inds]

    # Load images.
    image_dir = os.path.join(data_dir, f"video_frames_{frame_offset}")
    print(camtoworlds.shape)

    names = sorted(os.listdir(image_dir))

    # image_paths = [
    #    os.path.join(image_dir, f'image_{int(timestep):04d}.png')
    # ]

    image_paths = [os.path.join(image_dir, names[int(timestep) - 1])]

    images = []
    for x in image_paths:
        pil_image = Image.open(x).convert("RGB")
        if factor > 1:
            width, height = pil_image.size
            newsize = (int(width / factor), int(height / factor))
            pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
        images.append(np.array(pil_image, dtype="uint8"))

    images = np.stack(images, axis=0)
    images = np.repeat(images, len(camtoworlds), axis=0)

    # Select the split.
    all_indices = np.array([len(camtoworlds) - frame_offset])
    split_indices = {
        "test": all_indices,
        "train": all_indices,
    }
    return images, camtoworlds, intrinsics, split_indices


def _load_colmap(data_dir: str, factor: int = 1, hold_every: int = 8):
    assert factor in [1, 2, 4, 8]

    colmap_dir = os.path.join(data_dir, "sparse/0/")

    manager = SceneManager(colmap_dir)
    manager.load_cameras()
    manager.load_images()
    manager.load_points3D()
    pts_ret = manager.get_filtered_points3D(return_colors=True)

    # Assume shared intrinsics between all cameras.
    cam = manager.cameras[1]
    fx, fy, cx, cy = cam.fx, cam.fy, cam.cx, cam.cy
    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    K[:2, :] /= factor

    intrinsics = Intrinsics(
        width=int(cam.width / factor),
        height=int(cam.height / factor),
        focal_x=fx / factor,
        focal_y=fy / factor,
        center_x=cx / factor,
        center_y=cy / factor,
    )

    # Extract extrinsic matrices in world-to-camera format.
    imdata = manager.images
    w2c_mats = []
    bottom = np.array([0, 0, 0, 1]).reshape(1, 4)
    for k in imdata:
        im = imdata[k]
        rot = im.R()
        trans = im.tvec.reshape(3, 1)
        w2c = np.concatenate([np.concatenate([rot, trans], 1), bottom], axis=0)
        w2c_mats.append(w2c)
    w2c_mats = np.stack(w2c_mats, axis=0)

    # Convert extrinsics to camera-to-world.
    camtoworlds = np.linalg.inv(w2c_mats)

    # Image names from COLMAP. No need for permuting the poses according to
    # image names anymore.
    image_names = [imdata[k].name for k in imdata]

    # # Switch from COLMAP (right, down, fwd) to Nerf (right, up, back) frame.
    # poses = poses @ np.diag([1, -1, -1, 1])

    # Get distortion parameters.
    type_ = cam.camera_type

    if type_ == 0 or type_ == "SIMPLE_PINHOLE":
        params = None
        camtype = "perspective"

    elif type_ == 1 or type_ == "PINHOLE":
        params = None
        camtype = "perspective"

    if type_ == 2 or type_ == "SIMPLE_RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        camtype = "perspective"

    elif type_ == 3 or type_ == "RADIAL":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        camtype = "perspective"

    elif type_ == 4 or type_ == "OPENCV":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "p1", "p2"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["p1"] = cam.p1
        params["p2"] = cam.p2
        camtype = "perspective"

    elif type_ == 5 or type_ == "OPENCV_FISHEYE":
        params = {k: 0.0 for k in ["k1", "k2", "k3", "k4"]}
        params["k1"] = cam.k1
        params["k2"] = cam.k2
        params["k3"] = cam.k3
        params["k4"] = cam.k4
        camtype = "fisheye"

    # assert params is None, f"Only support pinhole camera model. Params: {params}"

    # Previous Nerf results were generated with images sorted by filename,
    # ensure metrics are reported on the same test set.
    inds = np.argsort(image_names)
    image_names = [image_names[i] for i in inds]
    camtoworlds = camtoworlds[inds]

    # Load images.
    # if factor > 1:
    #    image_dir_suffix = f"_{int(factor)}"
    # else:
    #    image_dir_suffix = ""

    colmap_image_dir = os.path.join(data_dir, "images")
    image_dir = os.path.join(data_dir, "images")
    for d in [image_dir, colmap_image_dir]:
        if not os.path.exists(d):
            raise ValueError(f"Image folder {d} does not exist.")
    # Downsampled images may have different names vs images used for COLMAP,
    # so we need to map between the two sorted lists of files.
    colmap_files = sorted(os.listdir(colmap_image_dir))
    image_files = sorted(os.listdir(image_dir))
    colmap_to_image = dict(zip(colmap_files, image_files))
    image_paths = [os.path.join(image_dir, colmap_to_image[f]) for f in image_names]
    print("loading images")
    images = []
    for x in image_paths:
        pil_image = Image.open(x).convert("RGB")
        if factor > 1:
            width, height = pil_image.size
            newsize = (int(width / factor), int(height / factor))
            pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
        images.append(np.array(pil_image, dtype="uint8"))

    # images = [imageio.imread(x) for x in image_paths]

    images = np.stack(images, axis=0)

    # Select the split.
    all_indices = np.arange(images.shape[0])
    split_indices = {
        "test": all_indices[all_indices % hold_every == 0],
        "train": all_indices,
    }
    return images, camtoworlds, intrinsics, split_indices


def similarity_from_cameras(c2w, strict_scaling):
    """
    reference: nerf-factory
    Get a similarity transform to normalize dataset
    from c2w (OpenCV convention) cameras
    :param c2w: (N, 4)
    :return T (4,4) , scale (float)
    """
    t = c2w[:, :3, 3]
    R = c2w[:, :3, :3]

    # (1) Rotate the world so that z+ is the up axis
    # we estimate the up axis by averaging the camera up axes
    ups = np.sum(R * np.array([0, -1.0, 0]), axis=-1)
    world_up = np.mean(ups, axis=0)
    world_up /= np.linalg.norm(world_up)

    up_camspace = np.array([0.0, -1.0, 0.0])
    c = (up_camspace * world_up).sum()
    cross = np.cross(world_up, up_camspace)
    skew = np.array(
        [
            [0.0, -cross[2], cross[1]],
            [cross[2], 0.0, -cross[0]],
            [-cross[1], cross[0], 0.0],
        ]
    )
    if c > -1:
        R_align = np.eye(3) + skew + (skew @ skew) * 1 / (1 + c)
    else:
        # In the unlikely case the original data has y+ up axis,
        # rotate 180-deg about x axis
        R_align = np.array([[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])

    #  R_align = np.eye(3) # DEBUG
    R = R_align @ R
    fwds = np.sum(R * np.array([0, 0.0, 1.0]), axis=-1)
    t = (R_align @ t[..., None])[..., 0]

    # (2) Recenter the scene using camera center rays
    # find the closest point to the origin for each camera's center ray
    nearest = t + (fwds * -t).sum(-1)[:, None] * fwds

    # median for more robustness
    translate = -np.median(nearest, axis=0)

    #  translate = -np.mean(t, axis=0)  # DEBUG

    transform = np.eye(4)
    transform[:3, 3] = translate
    transform[:3, :3] = R_align

    # (3) Rescale the scene using camera distances
    scale_fn = np.max if strict_scaling else np.median
    scale = 1.0 / scale_fn(np.linalg.norm(t + translate, axis=-1))

    return transform, scale


class TripodDataset(BaseDataset):
    def __init__(
        self,
        datadir,
        split: str,
        batch_size: Optional[int] = None,
        downsample: int = 4,
        hold_every: int = 8,
        contraction: bool = False,
        frame_offset: int = 1,
        near_plane_offset: float = 0.0,
        start_interp_index: int = 5,
        aabb=[-1.0, -1.0, -1.0, 1.0, 1.0, 1.0],
        use_app: bool = False,
        near_plane: float = 0.2,
        far_plane: float = 30.0,
        bg_color_aug: str = "black",
    ):
        self.downsample = downsample
        self.hold_every = hold_every
        self.bg_color_aug = bg_color_aug
        if not isinstance(contraction, bool):
            contraction = eval(contraction)

        data_dir = "/".join(datadir.split("/")[:-1])
        self.timestep = datadir.split("/")[-1]
        scene_name = datadir.split("/")[-2]

        if split == "render":
            if self.timestep != "000":
                imgs, poses, intrinsics, split_indices = _load_colmap_temporal(
                    data_dir, downsample, hold_every, self.timestep, frame_offset
                )
            else:
                imgs, poses, intrinsics, split_indices = _load_colmap(
                    data_dir, downsample, hold_every
                )
            T, sscale = similarity_from_cameras(poses, strict_scaling=False)
            poses = np.einsum("nij, ki -> nkj", poses, T)
            poses[:, :3, 3] *= sscale
            poses = torch.from_numpy(poses).float()

            sta_id = start_interp_index
            end_id = len(poses) - frame_offset

            render_poses_0 = poses[sta_id:sta_id+1, :3].repeat(30, 1, 1)
            render_poses_1 = generate_interpolated_path(
                torch.cat([render_poses_0[0:1], poses[end_id : end_id + 1, :3]]),
                n_interp=30,
                look_at_neg=False,
            )
            render_poses_2 = poses[end_id : end_id + 1, :3].repeat(30, 1, 1)
            render_poses_3 = generate_interpolated_path(
                torch.cat([poses[end_id : end_id + 1, :3], render_poses_0[0:1]]),
                n_interp=30,
                look_at_neg=False,
            )
            self.poses = torch.cat(
                [
                    render_poses_0,
                    render_poses_0,
                    render_poses_1,
                    render_poses_2,
                    render_poses_2,
                    render_poses_3,
                ],
                dim=0,
            )
            if scene_name == "standing":
                rot = (
                    torch.from_numpy(rotx(-1 * np.pi / 36)).float()
                    @ torch.from_numpy(roty(-1 * np.pi / 30)).float()
                )
                rot = rot @ torch.from_numpy(rotz(-1 * np.pi / 30)).float()

                rot = rot @ torch.from_numpy(trans(x=0.6, z=-0.0, y=-0.5)).float()
                pose_a = poses[end_id : end_id + 1, :3]
                pose_b = poses[end_id : end_id + 1, :3] @ rot
                a2b = generate_interpolated_path(
                    torch.cat([pose_a, pose_b]), n_interp=30, look_at_neg=False
                )
                rot = (
                    torch.from_numpy(roty(1 * np.pi / 36)).float()
                    @ torch.from_numpy(rotz(1 * np.pi / 30)).float()
                )
                rot = rot @ torch.from_numpy(trans(x=-0.5, z=0.2, y=0.2)).float()
                pose_c = poses[end_id : end_id + 1, :3] @ rot
                c2a = generate_interpolated_path(
                    torch.cat([pose_c, pose_a]), n_interp=30, look_at_neg=False
                )
                b2c = generate_interpolated_path(
                    torch.cat([pose_b, pose_c]), n_interp=30, look_at_neg=False
                )
                self.poses = torch.cat([pose_a.repeat(60, 1, 1), a2b, b2c, c2a], dim=0)

            if scene_name == "breathing":
                rot = (
                    torch.from_numpy(rotx(-1 * np.pi / 18)).float()
                    @ torch.from_numpy(roty(-np.pi / 15)).float()
                )
                rot = rot @ torch.from_numpy(trans(x=2.0, z=-1.0, y=-1.5)).float()
                pose_a = poses[end_id : end_id + 1, :3]
                pose_b = poses[end_id : end_id + 1, :3] @ rot
                a2b = generate_interpolated_path(
                    torch.cat([pose_a, pose_b]), n_interp=30, look_at_neg=False
                )
                rot = (
                    torch.from_numpy(roty(np.pi / 30)).float()
                    @ torch.from_numpy(rotz(np.pi / 15)).float()
                )
                rot = rot @ torch.from_numpy(trans(x=-0.0, z=1.0, y=0.5)).float()
                pose_c = poses[end_id : end_id + 1, :3] @ rot
                c2a = generate_interpolated_path(
                    torch.cat([pose_c, pose_a]), n_interp=30, look_at_neg=False
                )
                b2c = generate_interpolated_path(
                    torch.cat([pose_b, pose_c]), n_interp=30, look_at_neg=False
                )
                self.poses = torch.cat(
                    [
                        pose_a.repeat(60, 1, 1),
                        a2b,
                        pose_b.repeat(60, 1, 1),
                        b2c,
                        pose_c.repeat(60, 1, 1),
                        c2a,
                    ],
                    dim=0,
                )

            if scene_name == "fork":
                pose_a = poses[21 : 21 + 1, :3]
                pose_b = poses[1 : 1 + 1, :3]
                pose_c = poses[4 : 4 + 1, :3]
                a2b = generate_interpolated_path(
                    torch.cat([pose_a, pose_b]), n_interp=30, look_at_neg=False
                )
                b2c = generate_interpolated_path(
                    torch.cat([pose_b, pose_c]), n_interp=30, look_at_neg=False
                )
                c2a = generate_interpolated_path(
                    torch.cat([pose_c, pose_a]), n_interp=30, look_at_neg=False
                )
                self.poses = torch.cat([pose_a.repeat(30, 1, 1), a2b, b2c, c2a], dim=0)

            imgs = None
        else:
            if self.timestep != "000" and split == "train":
                imgs, poses, intrinsics, split_indices = _load_colmap_temporal(
                    data_dir, downsample, hold_every, self.timestep, frame_offset
                )
            else:
                imgs, poses, intrinsics, split_indices = _load_colmap(
                    data_dir, downsample, hold_every
                )

            T, sscale = similarity_from_cameras(poses, strict_scaling=False)
            poses = np.einsum("nij, ki -> nkj", poses, T)
            poses[:, :3, 3] *= sscale

            self.poses = torch.from_numpy(poses).float()

            indices = split_indices[split]
            imgs = imgs[indices]
            self.poses = self.poses[indices]
            imgs = torch.from_numpy(imgs).to(torch.uint8)

            if split == "train":
                imgs = imgs.view(-1, imgs.shape[-1])
            else:
                imgs = imgs.view(
                    -1, intrinsics.height * intrinsics.width, imgs.shape[-1]
                )

        self.num_images = len(os.listdir(f"{data_dir}/images"))
        self.test_img_id = self.num_images - frame_offset

        self.near_fars = torch.tensor([near_plane, far_plane])
        if split != "train":
            self.near_fars[0] += near_plane_offset

        self.global_translation = torch.tensor([0.0, 0.0, 0.0])
        self.global_scale = aabb
        bbox = torch.tensor([aabb[:3], aabb[3:]])

        if type(use_app) == str:
            use_app = bool(eval(use_app))

        self.use_app_emb = use_app
        self.OPENGL_CAMERA = False

        super().__init__(
            datadir=datadir,
            split=split,
            scene_bbox=bbox,
            batch_size=batch_size,
            imgs=imgs,
            rays_o=None,
            rays_d=None,
            intrinsics=intrinsics,
            is_ndc=False,
            is_contracted=contraction,
        )

        log.info(
            f"TripodDataset. {contraction=}. Loaded {split} set from {datadir}. "
            f"{self.num_images} poses of shape {self.img_h}x{self.img_w}. "
            f"Images loaded: {imgs is not None}. Near-far[:3]: {self.near_fars[:3]}. "
            f"Sampling without replacement={self.use_permutation}. {intrinsics}"
        )

    def __getitem__(self, index):
        h = self.intrinsics.height
        w = self.intrinsics.width
        dev = "cpu"
        out = {"near_fars": self.near_fars.view(-1, 2)}

        if self.split == "train":
            index = self.get_rand_ids(index)
            image_id = torch.div(index, h * w, rounding_mode="floor").long()
            y = torch.remainder(index, h * w).div(w, rounding_mode="floor")
            x = torch.remainder(index, h * w).remainder(w)
            x = x + 0.5
            y = y + 0.5

            if self.bg_color_aug == "random":
                out["bg_color"] = torch.rand_like(torch.tensor([[0.0, 0.0, 0.0]]))
            elif self.bg_color_aug == "white":
                out["bg_color"] = torch.tensor([[1.0, 1.0, 1.0]])
            else:
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])
        else:
            image_id = torch.LongTensor([index])
            x, y = create_meshgrid(height=h, width=w, dev=dev, add_half=True, flat=True)
            if self.bg_color_aug == "random":
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])
            elif self.bg_color_aug == "white":
                out["bg_color"] = torch.tensor([[1.0, 1.0, 1.0]])
            else:
                out["bg_color"] = torch.tensor([[0.0, 0.0, 0.0]])

        if self.imgs is not None:
            out["imgs"] = (
                self.imgs[index] / 255.0
            )  # (num_rays, 3)   this converts to f32
        else:
            out["imgs"] = None

        c2w = self.poses[image_id]  # (num_rays, 3, 4)
        camera_dirs = stack_camera_dirs(
            x, y, self.intrinsics, self.OPENGL_CAMERA
        )  # [num_rays, 3]
        # [num_rays, 3]
        directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
        origins = torch.broadcast_to(c2w[:, :3, -1], directions.shape)
        viewdirs = directions / torch.linalg.norm(directions, dim=-1, keepdims=True)

        out["rays_o"] = origins
        out["rays_d"] = viewdirs

        if self.use_app_emb:
            out["timestamps"] = (image_id + 1) % self.num_images
            if self.timestep != "000":
                out["timestamps"] = (
                    self.test_img_id * torch.ones_like(image_id) + 1
                ) % self.num_images

        return out


magnerf/datasets/tripod.py
3d-motion-magnification/3d-motion-mag
"""
Copyright (c) 2022 Ruilong Li, UC Berkeley.
"""
import numpy as np
import scipy
import torch

import collections

Rays = collections.namedtuple("Rays", ("origins", "viewdirs"))


def namedtuple_map(fn, tup):
    """Apply `fn` to each element of `tup` and cast to `tup`'s namedtuple."""
    return type(tup)(*(None if x is None else fn(x) for x in tup))

def normalize(v: np.ndarray) -> np.ndarray:
    """Normalize a vector."""
    return v / np.linalg.norm(v)

def generate_interpolated_path(poses: np.ndarray,
                               n_interp: int,
                               spline_degree: int = 5,
                               smoothness: float = .03,
                               rot_weight: float = .1,
                               look_at_neg: bool = False):
  """Creates a smooth spline path between input keyframe camera poses.
  Spline is calculated with poses in format (position, lookat-point, up-point).
  Args:
    poses: (n, 3, 4) array of input pose keyframes.
    n_interp: returned path will have n_interp * (n - 1) total poses.
    spline_degree: polynomial degree of B-spline.
    smoothness: parameter for spline smoothing, 0 forces exact interpolation.
    rot_weight: relative weighting of rotation/translation in spline solve.
  Returns:
    Array of new camera poses with shape (n_interp * (n - 1), 3, 4).
  """

  def viewmatrix(lookdir: np.ndarray, up: np.ndarray,
                position: np.ndarray) -> np.ndarray:
    """Construct lookat view matrix."""
    vec2 = normalize(lookdir)
    vec0 = normalize(np.cross(up, vec2))
    vec1 = normalize(np.cross(vec2, vec0))
    m = np.stack([vec0, vec1, vec2, position], axis=1)
    return m

  def poses_to_points(poses, dist, look_at_neg):
    """Converts from pose matrices to (position, lookat, up) format."""
    pos = poses[:, :3, -1]
    if look_at_neg:
        lookat = poses[:, :3, -1] + dist * poses[:, :3, 2]
        up = poses[:, :3, -1] - dist * poses[:, :3, 1]
    else:
        lookat = poses[:, :3, -1] - dist * poses[:, :3, 2]
        up = poses[:, :3, -1] + dist * poses[:, :3, 1]
    return np.stack([pos, lookat, up], 1)

  def points_to_poses(points):
    """Converts from (position, lookat, up) format to pose matrices."""
    pose = np.array([viewmatrix(p - l, u - p, p) for p, l, u in points])
    return torch.from_numpy(pose).float()

  def interp(points, n, k, s):
    """Runs multidimensional B-spline interpolation on the input points."""
    sh = points.shape
    pts = np.reshape(points, (sh[0], -1))
    k = min(k, sh[0] - 1)
    tck, _ = scipy.interpolate.splprep(pts.T, k=k, s=s)
    u = np.linspace(0, 1, n, endpoint=False)
    new_points = np.array(scipy.interpolate.splev(u, tck))
    new_points = np.reshape(new_points.T, (n, sh[1], sh[2]))
    return new_points

  points = poses_to_points(poses, dist=rot_weight, look_at_neg=look_at_neg)
  new_points = interp(points,
                      n_interp * (points.shape[0] - 1),
                      k=spline_degree,
                      s=smoothness)
  return points_to_poses(new_points)
magnerf/datasets/utils.py
3d-motion-magnification/3d-motion-mag
import os
import warnings
warnings.simplefilter("ignore")
os.environ["PYTHONWARNINGS"] = "ignore"

import argparse
import importlib.util
import logging
import pprint
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from typing import List, Dict, Any
import tempfile

import numpy as np
import imageio
import tqdm
from scipy import signal
import subprocess

from plenoptic.simulate import Steerable_Pyramid_Freq


def bandpass_filter(data_3d, mask):
    d_data_3d = data_3d

    data_f = torch.fft.fft(d_data_3d)
    data_f = data_f * mask.to(data_f.device)

    out = torch.fft.ifft(data_f)

    return out.real


def bandpass_filter_np(data_3d, mask):
    d_data_3d = data_3d - data_3d[..., 0:1]
    data_f = np.fft.fft(d_data_3d)
    data_f[~mask] = 0
    filtered = np.fft.ifft(data_f)

    return filtered.real


def create_bandpass_mask(lowcut, highcut, n, fs):
    fl, fh = lowcut / fs, highcut / fs
    fl, fh = fl * 2, fh * 2
    B = signal.firwin(n, cutoff=[fl, fh], window="hamming", pass_zero="bandpass")
    B = torch.FloatTensor(B[:n])

    mask = torch.fft.fft(torch.fft.ifftshift(B))
    mask = mask.unsqueeze(0).unsqueeze(0)

    return mask


def create_bandpass_mask_ideal(lowcut, highcut, n, fs):
    freq = np.arange(n) / n
    freq = freq * fs

    mask = (freq > lowcut) * (freq <= highcut)
    mask = torch.from_numpy(mask)

    return mask


def amplify_one_channel_vid(
    f_vid, alpha, low, high, fs, ideal=False, suppress_others=False, flt_thrs=0.0
):
    if ideal:
        mask = create_bandpass_mask_ideal(
            lowcut=low, highcut=high, n=f_vid.shape[-1], fs=fs
        )
    else:
        mask = create_bandpass_mask(lowcut=low, highcut=high, n=f_vid.shape[-1], fs=fs)

    mask = mask.repeat(f_vid.shape[0], f_vid.shape[1], 1)

    ref_phs = f_vid.angle()[..., 0:1]

    delta = f_vid.angle() - ref_phs
    delta = ((np.pi + delta) % (2 * np.pi)) - np.pi

    flt_phs = bandpass_filter(delta, mask)

    thrs = torch.quantile(flt_phs.abs(), q=flt_thrs)
    flt_phs[flt_phs.abs() < thrs] = 0.0

    amp_phs = alpha * flt_phs
    amp_phs = ((np.pi + amp_phs) % (2 * np.pi)) - np.pi

    if suppress_others:
        amp_vid = f_vid[..., 0:1] * torch.exp(1j * amp_phs)
    else:
        amp_vid = f_vid * torch.exp(1j * amp_phs)

    return amp_vid


def process_feature_images(f_im, alpha, low, high, fs, device, args):
    f_im = torch.cat(f_im, dim=0)

    pyr = Steerable_Pyramid_Freq(
        height="auto",
        image_shape=[f_im.shape[-2], f_im.shape[-1]],
        order=args.orientation,
        is_complex=True,
        downsample=True,
        twidth=0.75,
    )
    pyr.to(device)
    mbsize = 30

    for channel in tqdm.trange(f_im.shape[1], disable=True):
        all_coeffs = {}

        for mb in range(0, len(f_im), mbsize):
            coeff = pyr.forward(
                f_im[mb : mb + mbsize, channel : channel + 1].to(device)
            )
            for key in coeff.keys():
                all_coeffs.setdefault(key, []).append(coeff[key])
            del coeff
        for key in all_coeffs.keys():
            all_coeffs[key] = torch.cat(all_coeffs[key])

        for key in all_coeffs.keys():
            _vid = all_coeffs[key]
            if "residual_highpass" in key or "residual_lowpass" in key:
                continue
            else:
                amp_vid = amplify_one_channel_vid(
                    _vid[:, 0].permute(1, 2, 0),
                    alpha,
                    low,
                    high,
                    fs,
                    ideal=args.ideal,
                    suppress_others=args.suppress_others,
                    flt_thrs=args.flt_thrs,
                )
            all_coeffs[key] = amp_vid.permute(2, 0, 1)[:, None]
            del amp_vid

        for mb in range(0, len(f_im), mbsize):
            mb_coeffs = {}
            for key in all_coeffs.keys():
                mb_coeffs[key] = all_coeffs[key][mb : mb + mbsize]

            out = pyr.recon_pyr(mb_coeffs)
            out = out.cpu()

            f_im[mb : mb + mbsize, channel : channel + 1] = out

            del out, mb_coeffs
        del all_coeffs

    return torch.unbind(f_im.unsqueeze(1), 0)


def get_freer_gpu():
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_fname = os.path.join(tmpdir, "tmp")
        os.system(f'nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >"{tmp_fname}"')
        if os.path.isfile(tmp_fname):
            memory_available = [
                int(x.split()[2]) for x in open(tmp_fname, "r").readlines()
            ]
            if len(memory_available) > 0:
                return np.argmax(memory_available)
    return None  # The grep doesn't work with all GPUs. If it fails we ignore it.


gpu = get_freer_gpu()
if gpu is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
    print(f"CUDA_VISIBLE_DEVICES set to {gpu}")
else:
    print(f"Did not set GPU.")

import torch
import torch.utils.data
from magnerf.runners import static_trainer
from magnerf.utils.parse_args import parse_optfloat


def setup_logging(log_level=logging.INFO):
    handlers = [logging.StreamHandler(sys.stdout)]
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s|%(levelname)8s| %(message)s",
        handlers=handlers,
        force=True,
    )


def load_data(
    model_type: str,
    data_downsample,
    data_dirs,
    validate_only: bool,
    render_only: bool,
    **kwargs,
):
    data_downsample = parse_optfloat(data_downsample, default_val=1.0)

    return static_trainer.load_data(
        data_downsample,
        data_dirs,
        validate_only=validate_only,
        render_only=render_only,
        **kwargs,
    )


def init_trainer(model_type: str, **kwargs):
    from magnerf.runners import static_trainer

    return static_trainer.StaticTrainer(**kwargs)


def vis_grid_features(new_field_feature_imgs, base_logdir):
    planes = torch.cat(new_field_feature_imgs["delta_grids.0.0"]).cpu()
    out = planes.permute(0, 2, 3, 1)[..., :3]
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    imageio.mimwrite(f"{base_logdir}/xy.mp4", out, fps=10, quality=8)

    planes = torch.cat(new_field_feature_imgs["delta_grids.0.1"]).cpu()
    out = planes.permute(0, 2, 3, 1)[..., :3]
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    imageio.mimwrite(f"{base_logdir}/xz.mp4", out, fps=10, quality=8)

    planes = torch.cat(new_field_feature_imgs["delta_grids.0.2"]).cpu()
    out = planes.permute(0, 2, 3, 1)[..., :3]
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    imageio.mimwrite(f"{base_logdir}/yz.mp4", out, fps=10, quality=8)


def main():
    setup_logging()
    logger = logging.getLogger()
    logger.disabled = True

    p = argparse.ArgumentParser(description="")

    p.add_argument("--render-only", action="store_true")
    p.add_argument("--validate-only", action="store_true")
    p.add_argument("--spacetime-only", action="store_true")
    p.add_argument("--config-path", type=str, required=True)
    p.add_argument("--log-dir", type=str, default=None)
    p.add_argument("--timestep", type=str, default="000")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--save_video", action="store_true")

    p.add_argument(
        "--depth", "-n", default=4, type=int, help="Depth of Pyramid. Integer"
    )
    p.add_argument(
        "--orientation", "-k", default=8, type=int, help="Orientation. Integer"
    )

    p.add_argument("--num_timesteps", type=int, default=30)
    p.add_argument("--num_frames", type=int, default=30)
    p.add_argument("--alpha", type=float, default=10.0)
    p.add_argument("--low", type=float, default=2.5)
    p.add_argument("--high", type=float, default=3.5)
    p.add_argument("--fs", type=int, default=30)
    p.add_argument("--view_num", type=int, default=-1)
    p.add_argument("--ideal", action="store_true")
    p.add_argument("--suppress_others", action="store_true")
    p.add_argument("--start_idx", type=int, default=1)
    p.add_argument("--vid_fps", type=int, default=30)
    p.add_argument("--flt_thrs", type=float, default=0.0)

    p.add_argument("override", nargs=argparse.REMAINDER)

    args = p.parse_args()

    num_frames = args.num_frames
    num_timesteps = args.num_timesteps
    alpha = args.alpha
    low, high, fs = args.low, args.high, args.fs

    start_idx = args.start_idx
    device = "cuda"

    # Set random seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Import config
    spec = importlib.util.spec_from_file_location(
        os.path.basename(args.config_path), args.config_path
    )
    cfg = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(cfg)
    config: Dict[str, Any] = cfg.config
    # Process overrides from argparse into config
    # overrides can be passed from the command line as key=value pairs. E.g.
    # python plenoxels/main.py --config-path plenoxels/config/cfg.py max_ts_frames=200
    # note that all values are strings, so code should assume incorrect data-types for anything
    # that's derived from config - and should not a string.
    overrides: List[str] = args.override
    overrides_dict = {ovr.split("=")[0]: ovr.split("=")[1] for ovr in overrides}

    timestep = "000"
    expname = config["expname"]
    data_dir = config["data_dirs"]
    base_logdir = config["logdir"]
    if "runname" in config:
        runname = config["runname"]
    else:
        runname = "."

    if "expname" in overrides_dict:
        expname = overrides_dict["expname"]
    if "runname" in overrides_dict:
        runname = overrides_dict["runname"]
    if "data_dirs" in overrides_dict:
        data_dir = overrides_dict["data_dirs"]
    if "logdir" in overrides_dict:
        base_logdir = overrides_dict["logdir"]

    overrides_dict["data_dirs"] = f"{data_dir}/{expname}/{timestep}"

    expname = f"{expname}/{runname}"

    ckpt_base_logdir = f"{base_logdir}/{expname}/output"

    if "angle" in overrides_dict:
        name2 = overrides_dict["expname2"]
        overrides_dict["data_dirs"] = f"{data_dir}/{name2}/{timestep}"
        base_logdir = f'{base_logdir}/{overrides_dict["angle"]}'
        overrides_dict["logdir"] = base_logdir
        overrides_dict["test_split"] = f"test_angle_{overrides_dict['angle']}"

    config.update(overrides_dict)

    base_logdir = f"{base_logdir}/{expname}/output"
    render_dir = f"{base_logdir}/render/phase"
    os.makedirs(f"{render_dir}", exist_ok=True)

    frame_dir = f"{render_dir}/{low:.1f}_{high:.1f}_x{alpha:.1f}"
    if args.ideal:
        frame_dir = f"{frame_dir}_ideal"
    if args.suppress_others:
        frame_dir = f"{frame_dir}_supp"
    if args.flt_thrs > 0.0:
        frame_dir = f"{frame_dir}_thrs_{args.flt_thrs}"

    frame_dir = f"{frame_dir}/{args.view_num:03d}"

    os.makedirs(f"{frame_dir}", exist_ok=True)

    for f in os.listdir(frame_dir):
        if f.endswith(".png"):
            os.remove(f"{frame_dir}/{f}")

    model_type = "static"
    validate_only = args.validate_only
    render_only = args.render_only
    spacetime_only = args.spacetime_only
    if validate_only and render_only:
        raise ValueError("render_only and validate_only are mutually exclusive.")
    if render_only and spacetime_only:
        raise ValueError("render_only and spacetime_only are mutually exclusive.")
    if validate_only and spacetime_only:
        raise ValueError("validate_only and spacetime_only are mutually exclusive.")

    if "bbox" in config and isinstance(config["bbox"], str):
        config.update({"bbox": list(map(float, config["bbox"].split(",")))})
    if isinstance(config["num_samples"], str):
        config.update({"num_samples": int(config["num_samples"])})

    pprint.pprint(config)

    data = load_data(
        model_type,
        validate_only=validate_only,
        render_only=render_only or spacetime_only,
        **config,
    )
    config.update(data)
    trainer = init_trainer(model_type, **config)

    dset = trainer.test_dataset

    field_feature_imgs = {}
    field_feature_imgs2 = {}

    for i in range(start_idx, num_timesteps + start_idx):
        timestep = f"{i:03d}"
        checkpoint_path = os.path.join(ckpt_base_logdir, f"{timestep}", "model.pth")
        trainer = init_trainer(model_type, **config)
        trainer.load_model(torch.load(checkpoint_path), training_needed=False)

        for n, p in trainer.model.field.named_parameters():
            if "delta" not in n and len(p.shape) == 4:
                field_feature_imgs2.setdefault(n, []).append(p.data.clone().cpu())

        for n, p in trainer.model.field.named_parameters():
            if "delta_grids" in n:
                n2 = n.replace("delta_", "")
                x = field_feature_imgs2[n2][0]
                field_feature_imgs.setdefault(n, []).append(p.data.clone().cpu() + x)

    trainer.model.field.use_delta_features = True
    trainer.model.field.delta_add_back = False

    new_field_feature_imgs = {}
    for n, f_im in field_feature_imgs.items():
        if low < 0 and high < 0:
            new_field_feature_imgs[n] = f_im
        else:
            print(f"Processing {n}")
            processed = process_feature_images(f_im, alpha, low, high, fs, device, args)
            new_field_feature_imgs[n] = processed
    vis_grid_features(new_field_feature_imgs, frame_dir)

    for i in range(num_frames):
        for n, p in trainer.model.field.named_parameters():
            if "delta_grids" in n:
                new_data = new_field_feature_imgs[n][i % num_timesteps].data
                p.data = new_data.to(device)

        img_h, img_w = dset.img_h, dset.img_w

        view = args.view_num if args.view_num > -1 else i % len(dset.poses)
        ts_render = trainer.eval_step(dset[view])
        preds_rgb = (
            ts_render["rgb"]
            .reshape(img_h, img_w, 3)
            .cpu()
            .clamp(0, 1)
            .mul(255.0)
            .byte()
            .numpy()
        )
        imageio.imwrite(f"{frame_dir}/{i:03d}.jpg", preds_rgb)

    if args.save_video:
        subprocess.call(
            [
                "ffmpeg",
                "-framerate",
                f"{args.vid_fps}",
                "-i",
                f"{frame_dir}/%03d.jpg",
                "-vcodec",
                "libx264",
                "-y",
                "-crf",
                "20",
                "-pix_fmt",
                "yuv420p",
                "-vf",
                "pad=ceil(iw/2)*2:ceil(ih/2)*2",
                f"{frame_dir}/../{args.view_num:03d}.mp4",
            ]
        )


if __name__ == "__main__":
    main()

magnerf/mag_render.py
3d-motion-magnification/3d-motion-mag
import argparse
import importlib.util
import logging
import os
import pprint
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from typing import List, Dict, Any
import tempfile

import numpy as np


def get_freer_gpu():
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_fname = os.path.join(tmpdir, "tmp")
        os.system(f'nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >"{tmp_fname}"')
        if os.path.isfile(tmp_fname):
            memory_available = [int(x.split()[2]) for x in open(tmp_fname, 'r').readlines()]
            if len(memory_available) > 0:
                return np.argmax(memory_available)
    return None  # The grep doesn't work with all GPUs. If it fails we ignore it.

gpu = get_freer_gpu()
if gpu is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
    print(f"CUDA_VISIBLE_DEVICES set to {gpu}")
else:
    print(f"Did not set GPU.")

import torch
import torch.utils.data
from magnerf.runners import video_trainer
from magnerf.runners import phototourism_trainer
from magnerf.runners import static_trainer
from magnerf.utils.create_rendering import render_to_path, decompose_space_time
from magnerf.utils.parse_args import parse_optfloat


def setup_logging(log_level=logging.INFO):
    handlers = [logging.StreamHandler(sys.stdout)]
    logging.basicConfig(level=log_level,
                        format='%(asctime)s|%(levelname)8s| %(message)s',
                        handlers=handlers,
                        force=True)


def load_data(model_type: str, data_downsample, data_dirs, validate_only: bool, render_only: bool, **kwargs):
    data_downsample = parse_optfloat(data_downsample, default_val=1.0)

    if model_type == "video":
        return video_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs)
    elif model_type == "phototourism":
        return phototourism_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs
        )
    else:
        return static_trainer.load_data(
            data_downsample, data_dirs, validate_only=validate_only,
            render_only=render_only, **kwargs)


def init_trainer(model_type: str, **kwargs):
    if model_type == "video":
        from magnerf.runners import video_trainer
        return video_trainer.VideoTrainer(**kwargs)
    elif model_type == "phototourism":
        from magnerf.runners import phototourism_trainer
        return phototourism_trainer.PhototourismTrainer(**kwargs)
    else:
        from magnerf.runners import static_trainer
        return static_trainer.StaticTrainer(**kwargs)


def save_config(config):
    log_dir = os.path.join(config['logdir'], config['expname'])
    os.makedirs(log_dir, exist_ok=True)

    with open(os.path.join(log_dir, 'config.py'), 'wt') as out:
        out.write('config = ' + pprint.pformat(config))

    with open(os.path.join(log_dir, 'config.csv'), 'w') as f:
        for key in config.keys():
            f.write("%s\t%s\n" % (key, config[key]))


def main():
    setup_logging()

    p = argparse.ArgumentParser(description="")

    p.add_argument('--render-only', action='store_true')
    p.add_argument('--validate-only', action='store_true')
    p.add_argument('--spacetime-only', action='store_true')
    p.add_argument('--config-path', type=str, required=True)
    p.add_argument('--log-dir', type=str, default=None)
    p.add_argument('--seed', type=int, default=0)
    p.add_argument('override', nargs=argparse.REMAINDER)

    args = p.parse_args()

    # Set random seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Import config
    spec = importlib.util.spec_from_file_location(os.path.basename(args.config_path), args.config_path)
    cfg = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(cfg)
    config: Dict[str, Any] = cfg.config
    # Process overrides from argparse into config
    # overrides can be passed from the command line as key=value pairs. E.g.
    # python plenoxels/main.py --config-path plenoxels/config/cfg.py max_ts_frames=200
    # note that all values are strings, so code should assume incorrect data-types for anything
    # that's derived from config - and should not a string.
    overrides: List[str] = args.override
    overrides_dict = {ovr.split("=")[0]: ovr.split("=")[1] for ovr in overrides}
    config.update(overrides_dict)
    if "keyframes" in config:
        model_type = "video"
    elif "appearance_embedding_dim" in config:
        model_type = "phototourism"
    else:
        model_type = "static"
    validate_only = args.validate_only
    render_only = args.render_only
    spacetime_only = args.spacetime_only
    if validate_only and render_only:
        raise ValueError("render_only and validate_only are mutually exclusive.")
    if render_only and spacetime_only:
        raise ValueError("render_only and spacetime_only are mutually exclusive.")
    if validate_only and spacetime_only:
        raise ValueError("validate_only and spacetime_only are mutually exclusive.")

    pprint.pprint(config)
    if validate_only or render_only:
        assert args.log_dir is not None and os.path.isdir(args.log_dir)
    else:
        save_config(config)

    data = load_data(model_type, validate_only=validate_only, render_only=render_only or spacetime_only, **config)
    config.update(data)
    trainer = init_trainer(model_type, **config)
    if args.log_dir is not None:
        checkpoint_path = os.path.join(args.log_dir, "model.pth")
        training_needed = not (validate_only or render_only or spacetime_only)
        logging.info(f"=> Loaded checkpoint at {checkpoint_path}")
        trainer.load_model(torch.load(checkpoint_path), training_needed=training_needed)

    if validate_only:
        trainer.validate()
    elif render_only:
        render_to_path(trainer, extra_name="")
    elif spacetime_only:
        decompose_space_time(trainer, extra_name="")
    else:
        trainer.train()


if __name__ == "__main__":
    main()

magnerf/main.py
3d-motion-magnification/3d-motion-mag
"""
Density proposal field
"""
from typing import Optional, Callable
import logging as log

import torch
import torch.nn as nn
import tinycudann as tcnn

from magnerf.models.kplane_field import interpolate_ms_features, normalize_aabb, init_grid_param
from magnerf.raymarching.spatial_distortions import SpatialDistortion


class KPlaneDensityField(nn.Module):
    def __init__(self,
                 aabb,
                 resolution,
                 num_input_coords,
                 num_output_coords,
                 density_activation: Callable,
                 spatial_distortion: Optional[SpatialDistortion] = None,
                 linear_decoder: bool = True):
        super().__init__()
        self.aabb = nn.Parameter(aabb, requires_grad=False)
        self.spatial_distortion = spatial_distortion
        self.hexplane = num_input_coords == 4
        self.feature_dim = num_output_coords
        self.density_activation = density_activation
        self.linear_decoder = linear_decoder
        activation = "ReLU"
        if self.linear_decoder:
            activation = "None"

        self.feature_dim = self.feature_dim * 3

        self.grids = init_grid_param(
            grid_nd=2, in_dim=num_input_coords, out_dim=num_output_coords, reso=resolution,
            a=0.1, b=0.15)
        self.sigma_net = tcnn.Network(
            n_input_dims=self.feature_dim,
            n_output_dims=1,
            network_config={
                "otype": "FullyFusedMLP",
                "activation": activation,
                "output_activation": "None",
                "n_neurons": 64,
                "n_hidden_layers": 1,
            },
        )
        log.info(f"Initialized KPlaneDensityField. hexplane={self.hexplane} - "
                 f"resolution={resolution}")
        log.info(f"KPlaneDensityField grids: \n{self.grids}")

    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
        if self.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.aabb)
        n_rays, n_samples = pts.shape[:2]
        if timestamps is not None and self.hexplane:
            timestamps = timestamps[:, None].expand(-1, n_samples)[..., None]  # [n_rays, n_samples, 1]
            pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]

        pts = pts.reshape(-1, pts.shape[-1])
        features = interpolate_ms_features(
            pts, ms_grids=[self.grids], grid_dimensions=2, concat_features=False, num_levels=None)
        density = self.density_activation(
            self.sigma_net(features).to(pts)
            #features.to(pts)
        ).view(n_rays, n_samples, 1)
        return density

    def forward(self, pts: torch.Tensor):
        return self.get_density(pts)

    def get_params(self):
        field_params = {k: v for k, v in self.grids.named_parameters(prefix="grids")}
        nn_params = {k: v for k, v in self.sigma_net.named_parameters(prefix="sigma_net")}
        other_params = {k: v for k, v in self.named_parameters() if (
            k not in nn_params.keys() and k not in field_params.keys()
        )}
        return {
            "nn": list(nn_params.values()),
            "field": list(field_params.values()),
            "other": list(other_params.values()),
        }

magnerf/models/density_fields.py
3d-motion-magnification/3d-motion-mag
import itertools
import logging as log
from typing import Optional, Union, List, Dict, Sequence, Iterable, Collection, Callable

import torch
import torch.nn as nn
import tinycudann as tcnn
import numpy as np

from magnerf.ops.interpolation import grid_sample_wrapper
from magnerf.raymarching.spatial_distortions import SpatialDistortion


def get_normalized_directions(directions):
    """SH encoding must be in the range [0, 1]

    Args:
        directions: batch of directions
    """
    return (directions + 1.0) / 2.0


def normalize_aabb(pts, aabb):
    '''
    x = pts[..., 0:1]
    new_x = torch.where(x > 0.75, 10 * torch.ones_like(x), x)
    new_x = torch.where(x < -0.75, 10 * torch.ones_like(x), new_x)
    y = pts[..., 1:2]
    new_y = torch.where(x > 1.0, 10 * torch.ones_like(x), y)
    new_y = torch.where(x < -1.0, 10 * torch.ones_like(x), new_y)
    z = pts[..., 1:2]
    new_z = torch.where(z > 2, 10 * torch.ones_like(z), z)
    new_z = torch.where(z < 1, 10 * torch.ones_like(z), new_z)
    pts = torch.cat([new_x, new_y, new_z], -1)
    '''
    return (pts - aabb[0]) * (2.0 / (aabb[1] - aabb[0])) - 1.0


def init_grid_param(
        grid_nd: int,
        in_dim: int,
        out_dim: int,
        reso: Sequence[int],
        a: float = 0.1,
        b: float = 0.2):
    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
    has_time_planes = in_dim == 4
    assert grid_nd <= in_dim
    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
    grid_coefs = nn.ParameterList()
    for ci, coo_comb in enumerate(coo_combs):
        new_grid_coef = nn.Parameter(torch.empty(
            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
        ))
        if has_time_planes and 3 in coo_comb:  # Initialize time planes to 1
            nn.init.ones_(new_grid_coef)
        else:
            nn.init.uniform_(new_grid_coef, a=a, b=b)

        grid_coefs.append(new_grid_coef)

    return grid_coefs


def init_delta_grid_param(
        grid_nd: int,
        in_dim: int,
        out_dim: int,
        reso: Sequence[int],
        a: float = 0.1,
        b: float = 0.5):
    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
    assert grid_nd <= in_dim
    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
    grid_coefs = nn.ParameterList()
    for ci, coo_comb in enumerate(coo_combs):
        new_grid_coef = nn.Parameter(torch.empty(
            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
        ))
        nn.init.zeros_(new_grid_coef)
        grid_coefs.append(new_grid_coef)

    return grid_coefs



def init_grid_param_motion(
        grid_nd: int,
        in_dim: int,
        out_dim: int,
        reso: Sequence[int],
        zeros: bool = True):
    assert in_dim == len(reso), "Resolution must have same number of elements as input-dimension"
    assert grid_nd <= in_dim
    coo_combs = list(itertools.combinations(range(in_dim), grid_nd))
    grid_coefs = nn.ParameterList()
    for ci, coo_comb in enumerate(coo_combs):
        new_grid_coef = nn.Parameter(torch.empty(
            [1, out_dim] + [reso[cc] for cc in coo_comb[::-1]]
        ))
        nn.init.uniform_(new_grid_coef, -0.001, 0.001)
        grid_coefs.append(new_grid_coef)

    return grid_coefs


def interpolate_ms_features(pts: torch.Tensor,
                            ms_grids: Collection[Iterable[nn.Module]],
                            grid_dimensions: int,
                            concat_features: bool,
                            num_levels: Optional[int],
                            ) -> torch.Tensor:
    hadamard = False
    coo_combs = list(itertools.combinations(
        range(pts.shape[-1]), grid_dimensions)
    )
    if num_levels is None:
        num_levels = len(ms_grids)
    multi_scale_interp = [] if concat_features else 0.
    grid: nn.ParameterList
    for scale_id, grid in enumerate(ms_grids[:num_levels]):
        interp_space = 1. if hadamard else []
        for ci, coo_comb in enumerate(coo_combs):
            # interpolate in plane
            feature_dim = grid[ci].shape[1]  # shape of grid[ci]: 1, out_dim, *reso
            interp_out_plane = (
                grid_sample_wrapper(grid[ci], pts[..., coo_comb])
                .view(-1, feature_dim)
            )
            # compute product over planes
            interp_space = interp_space * interp_out_plane if hadamard else interp_space + [interp_out_plane]
        if not hadamard:
            interp_space = torch.cat(interp_space, dim = -1)

        # combine over scales
        if concat_features:
            multi_scale_interp.append(interp_space)
        else:
            multi_scale_interp = multi_scale_interp + interp_space

    if concat_features:
        multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
    return multi_scale_interp


def interpolate_ms_features_v2(pts: torch.Tensor,
                            ms_grids: Collection[Iterable[nn.Module]],
                            grid_dimensions: int,
                            concat_features: bool,
                            num_levels: Optional[int],
                            ) -> torch.Tensor:
    coo_combs = list(itertools.combinations(
        range(pts.shape[-1] - 1), grid_dimensions)
    )
    if num_levels is None:
        num_levels = len(ms_grids)
    multi_scale_interp = [] if concat_features else 0.
    grid: nn.ParameterList
    for scale_id, grid in enumerate(ms_grids[:num_levels]):
        interp_space = 1.
        for ci, _coo_comb in enumerate(coo_combs):
            coo_comb = _coo_comb + (3, )
            # interpolate in plane
            feature_dim = grid[ci].shape[1]  # shape of grid[ci]: 1, out_dim, *reso
            interp_out_plane = (
                grid_sample_wrapper(grid[ci], pts[..., coo_comb])
                .view(-1, feature_dim)
            )
            # compute product over planes
            interp_space = interp_space * interp_out_plane

        # combine over scales
        if concat_features:
            multi_scale_interp.append(interp_space)
        else:
            multi_scale_interp = multi_scale_interp + interp_space

    if concat_features:
        multi_scale_interp = torch.cat(multi_scale_interp, dim=-1)
    return multi_scale_interp


class KPlaneField(nn.Module):
    def __init__(
        self,
        aabb,
        grid_config: Union[str, List[Dict]],
        concat_features_across_scales: bool,
        multiscale_res: Optional[Sequence[int]],
        use_appearance_embedding: bool,
        appearance_embedding_dim: int,
        spatial_distortion: Optional[SpatialDistortion],
        density_activation: Callable,
        linear_decoder: bool,
        linear_decoder_layers: Optional[int],
        num_images: Optional[int],
        lagrangian: bool = False,
    ) -> None:
        super().__init__()

        self.aabb = nn.Parameter(aabb, requires_grad=False)
        self.spatial_distortion = spatial_distortion
        self.grid_config = grid_config

        self.multiscale_res_multipliers: List[int] = multiscale_res or [1]
        self.concat_features = concat_features_across_scales
        self.density_activation = density_activation
        self.linear_decoder = linear_decoder
        self.lagrangian = lagrangian

        # 1. Init planes
        self.grids = nn.ModuleList()
        self.feature_dim = 0
        for res in self.multiscale_res_multipliers:
            # initialize coordinate grid
            config = self.grid_config[0].copy()
            # Resolution fix: multi-res only on spatial planes
            config["resolution"] = [
                r * res for r in config["resolution"][:3]
            ] + config["resolution"][3:]
            gp = init_grid_param(
                grid_nd=config["grid_dimensions"],
                in_dim=config["input_coordinate_dim"],
                out_dim=config["output_coordinate_dim"],
                reso=config["resolution"],
            )
            # shape[1] is out-dim - Concatenate over feature len for each scale
            if self.concat_features:
                self.feature_dim += gp[-1].shape[1]
            else:
                self.feature_dim = gp[-1].shape[1]
            self.grids.append(gp)
            self.feature_dim = self.feature_dim * 3

        log.info(f"Initialized model grids: {self.grids}")

        self.delta_grids = nn.ModuleList()
        for res in self.multiscale_res_multipliers:
            config = self.grid_config[0].copy()
            config["resolution"] = [
                r * res for r in config["resolution"][:3]
            ] + config["resolution"][3:]
            gp = init_delta_grid_param(
                grid_nd=config["grid_dimensions"],
                in_dim=config["input_coordinate_dim"],
                out_dim=config["output_coordinate_dim"],
                reso=config["resolution"],
            )
            self.delta_grids.append(gp)

        self.use_delta_features = False
        self.delta_add_back = True

        if self.lagrangian:
            self.motion_net = tcnn.NetworkWithInputEncoding(
                n_input_dims=3,
                n_output_dims=3,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 5,
                },
                encoding_config={
                    "otype": "Frequency",
                    "n_frequencies": 3,
                }
            )

        # 2. Init appearance code-related parameters
        self.use_average_appearance_embedding = True  # for test-time
        self.use_appearance_embedding = use_appearance_embedding
        self.num_images = num_images
        self.appearance_embedding = None
        if use_appearance_embedding:
            assert self.num_images is not None
            self.appearance_embedding_dim = appearance_embedding_dim
            # this will initialize as normal_(0.0, 1.0)
            self.appearance_embedding = nn.Embedding(self.num_images, self.appearance_embedding_dim)
            self.test_appearance_idx = -1
        else:
            self.appearance_embedding_dim = 0


        self.direction_encoder = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "SphericalHarmonics",
                "degree": 6,
            },
        )

        # 3. Init decoder network
        if self.linear_decoder:
            assert linear_decoder_layers is not None
            # The NN learns a basis that is used instead of spherical harmonics
            # Input is an encoded view direction, output is weights for
            # combining the color features into RGB
            # This architecture is based on instant-NGP
            self.color_basis = tcnn.Network(
                n_input_dims=3 + self.appearance_embedding_dim,#self.direction_encoder.n_output_dims,
                n_output_dims=3 * self.feature_dim,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": linear_decoder_layers,
                },
            )
            # sigma_net just does a linear transformation on the features to get density
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=1,
                network_config={
                    "otype": "CutlassMLP",
                    "activation": "None",
                    "output_activation": "None",
                    "n_neurons": 128,
                    "n_hidden_layers": 0,
                },
            )
        else:
            self.geo_feat_dim = 15
            self.sigma_net = tcnn.Network(
                n_input_dims=self.feature_dim,
                n_output_dims=self.geo_feat_dim + 1,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 1,
                },
            )
            self.in_dim_color = (
                    self.direction_encoder.n_output_dims
                    + self.geo_feat_dim
                    + self.appearance_embedding_dim
            )
            self.color_net = tcnn.Network(
                n_input_dims=self.in_dim_color,
                n_output_dims=3,
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "Sigmoid",
                    "n_neurons": 64,
                    "n_hidden_layers": 2,
                },
            )
            self.feature_norm = None

        max_resolution = 2048
        base_resolution = 16
        n_levels = 16
        per_level_scale = np.exp(
            (np.log(max_resolution) - np.log(base_resolution)) / (n_levels - 1)
        ).tolist()
        log2_hashmap_size = 19

        if False:
            self.mlp_base = tcnn.NetworkWithInputEncoding(
                n_input_dims=3,
                n_output_dims=1 + self.geo_feat_dim,
                encoding_config={
                    "otype": "HashGrid",
                    "n_levels": n_levels,
                    "n_features_per_level": 2,
                    "log2_hashmap_size": log2_hashmap_size,
                    "base_resolution": base_resolution,
                    "per_level_scale": per_level_scale,
                },
                network_config={
                    "otype": "FullyFusedMLP",
                    "activation": "ReLU",
                    "output_activation": "None",
                    "n_neurons": 64,
                    "n_hidden_layers": 3,
                },
            )

    def get_density(self, pts: torch.Tensor, timestamps: Optional[torch.Tensor] = None):
        """Computes and returns the densities."""
        if self.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.aabb)

        n_rays, n_samples = pts.shape[:2]

        if timestamps is not None:
            timestamps = timestamps[:, None].expand(-1, n_samples)[..., None]  # [n_rays, n_samples, 1]
            pts = torch.cat((pts, timestamps), dim=-1)  # [n_rays, n_samples, 4]

        #for multires_grids in self.grids:
        #    for i, grid in enumerate(multires_grids):
        #        if i in [2, 4, 5]:
        #            grid.data[:, :, 0] = torch.ones_like(grid.data[:, :, 0])

        pts = pts.reshape(-1, pts.shape[-1])

        if self.lagrangian:
            delta = self.motion_net(pts)
            pts = pts + delta

        features = interpolate_ms_features(
            pts, ms_grids=self.grids,  # noqa
            grid_dimensions=self.grid_config[0]["grid_dimensions"],
            concat_features=self.concat_features, num_levels=None)

        if self.use_delta_features:
            delta_features = interpolate_ms_features(
                pts, ms_grids=self.delta_grids,  # noqa
                grid_dimensions=self.grid_config[0]["grid_dimensions"],
                concat_features=self.concat_features, num_levels=None)
            if self.delta_add_back:
                features = delta_features + features
            else:
                features = delta_features

        if len(features) < 1:
            features = torch.zeros((0, 1)).to(features.device)
        if self.linear_decoder:
            density_before_activation = self.sigma_net(features)  # [batch, 1]
        else:
            features = self.sigma_net(features)
            #features = self.mlp_base((pts + 1.0) / 2)

            features, density_before_activation = torch.split(
                features, [self.geo_feat_dim, 1], dim=-1)

        if self.feature_norm is not None:
            features = self.feature_norm(features)

        density = self.density_activation(
            density_before_activation.to(pts)
        ).view(n_rays, n_samples, 1)

        #density = self.density_activation(
        #    self.mlp_base((pts + 1.0) / 2)
        #).view(n_rays, n_samples, 1)

        return density, features

    def forward(self,
                pts: torch.Tensor,
                directions: torch.Tensor,
                timestamps: Optional[torch.Tensor] = None):
        camera_indices = None
        if self.use_appearance_embedding:
            if timestamps is None:
                raise AttributeError("timestamps (appearance-ids) are not provided.")
            camera_indices = timestamps
            timestamps = None

        density, features = self.get_density(pts, timestamps)
        n_rays, n_samples = pts.shape[:2]

        directions = directions.view(-1, 1, 3).expand(pts.shape).reshape(-1, 3)
        if not self.linear_decoder:
            directions = get_normalized_directions(directions)
            encoded_directions = self.direction_encoder(directions)

        if self.linear_decoder:
            color_features = [features]
        else:
            color_features = [encoded_directions, features.view(-1, self.geo_feat_dim)]

        if self.use_appearance_embedding:
            if camera_indices.dtype == torch.float32:
                # Interpolate between two embeddings. Currently they are hardcoded below.
                #emb1_idx, emb2_idx = 100, 121  # trevi
                emb1_idx, emb2_idx = 11, 142  # sacre
                emb_fn = self.appearance_embedding
                emb1 = emb_fn(torch.full_like(camera_indices, emb1_idx, dtype=torch.long))
                emb1 = emb1.view(emb1.shape[0], emb1.shape[2])
                emb2 = emb_fn(torch.full_like(camera_indices, emb2_idx, dtype=torch.long))
                emb2 = emb2.view(emb2.shape[0], emb2.shape[2])
                embedded_appearance = torch.lerp(emb1, emb2, camera_indices)
            elif self.training:
                #if torch.sum(camera_indices < 0) > 0:
                #    embedded_appearance = torch.ones(
                #        (camera_indices.shape[0], self.appearance_embedding_dim), device=directions.device
                #    ) * self.appearance_embedding.weight.mean(dim=0)
                #else:
                #    embedded_appearance = self.appearance_embedding(camera_indices)
                embedded_appearance = self.appearance_embedding(camera_indices)
            else:
                if hasattr(self, "test_appearance_embedding"):
                    embedded_appearance = self.test_appearance_embedding(camera_indices)
                #elif self.use_average_appearance_embedding:
                #    embedded_appearance = torch.ones(
                #        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                #    ) * self.appearance_embedding.weight.mean(dim=0)
                else:
                    #embedded_appearance = torch.zeros(
                    #    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    #)
                    embedded_appearance = torch.ones(
                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
                    ) * self.appearance_embedding.weight[self.test_appearance_idx]

            if self.training:
                # expand embedded_appearance from n_rays, dim to n_rays*n_samples, dim
                ea_dim = embedded_appearance.shape[-1]
                embedded_appearance = embedded_appearance.view(-1, 1, ea_dim).expand(n_rays, n_samples, -1).reshape(-1, ea_dim)
            if not self.linear_decoder:
                color_features.append(embedded_appearance)

        color_features = torch.cat(color_features, dim=-1)

        if self.linear_decoder:
            if self.use_appearance_embedding:
                basis_values = self.color_basis(torch.cat([directions, embedded_appearance], dim=-1))
            else:
                basis_values = self.color_basis(directions)  # [batch, color_feature_len * 3]
            basis_values = basis_values.view(color_features.shape[0], 3, -1)  # [batch, 3, color_feature_len]
            rgb = torch.sum(color_features[:, None, :] * basis_values, dim=-1)  # [batch, 3]
            rgb = rgb.to(directions)
            rgb = torch.sigmoid(rgb).view(n_rays, n_samples, 3)
        else:
            rgb = self.color_net(color_features).to(directions).view(n_rays, n_samples, 3)

        rgb = rgb * (1 + 2 * 0.001) - 0.001

        return {"rgb": rgb, "density": density}

    def get_params(self):
        field_params = {k: v for k, v in self.grids.named_parameters(prefix="grids")}
        nn_params = [
            self.sigma_net.named_parameters(prefix="sigma_net"),
            self.direction_encoder.named_parameters(prefix="direction_encoder"),
        ]
        if self.linear_decoder:
            nn_params.append(self.color_basis.named_parameters(prefix="color_basis"))
        else:
            nn_params.append(self.color_net.named_parameters(prefix="color_net"))
        nn_params = {k: v for plist in nn_params for k, v in plist}
        other_params = {k: v for k, v in self.named_parameters() if (
            k not in nn_params.keys() and k not in field_params.keys()
        )}
        return {
            "nn": list(nn_params.values()),
            "field": list(field_params.values()),
            "other": list(other_params.values()),
        }

magnerf/models/kplane_field.py
3d-motion-magnification/3d-motion-mag
from typing import List, Sequence, Optional, Union, Dict, Tuple

import numpy as np
import torch
import torch.nn as nn

from magnerf.models.density_fields import KPlaneDensityField
from magnerf.models.kplane_field import KPlaneField
from magnerf.ops.activations import init_density_activation
from magnerf.raymarching.ray_samplers import (
    UniformLinDispPiecewiseSampler, UniformSampler,
    ProposalNetworkSampler, RayBundle, RaySamples
)
from magnerf.raymarching.spatial_distortions import SceneContraction, SpatialDistortion
from magnerf.utils.timer import CudaTimer
from magnerf.models.kplane_field import interpolate_ms_features, normalize_aabb

def exp_map_SO3xR3(tangent_vector):
    """Compute the exponential map of the direct product group `SO(3) x R^3`.

    This can be used for learning pose deltas on SE(3), and is generally faster than `exp_map_SE3`.

    Args:
        tangent_vector: Tangent vector; length-3 translations, followed by an `so(3)` tangent vector.
    Returns:
        [R|t] transformation matrices.
    """
    # code for SO3 map grabbed from pytorch3d and stripped down to bare-bones
    log_rot = tangent_vector[:, 3:]
    nrms = (log_rot * log_rot).sum(1)
    rot_angles = torch.clamp(nrms, 1e-4).sqrt()
    rot_angles_inv = 1.0 / rot_angles
    fac1 = rot_angles_inv * rot_angles.sin()
    fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
    skews = torch.zeros((log_rot.shape[0], 3, 3), dtype=log_rot.dtype, device=log_rot.device)
    skews[:, 0, 1] = -log_rot[:, 2]
    skews[:, 0, 2] = log_rot[:, 1]
    skews[:, 1, 0] = log_rot[:, 2]
    skews[:, 1, 2] = -log_rot[:, 0]
    skews[:, 2, 0] = -log_rot[:, 1]
    skews[:, 2, 1] = log_rot[:, 0]
    skews_square = torch.bmm(skews, skews)

    ret = torch.zeros(tangent_vector.shape[0], 3, 4, dtype=tangent_vector.dtype, device=tangent_vector.device)
    ret[:, :3, :3] = (
        fac1[:, None, None] * skews
        + fac2[:, None, None] * skews_square
        + torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
    )

    # Compute the translation
    ret[:, :3, 3] = tangent_vector[:, :3]
    return ret


def pose_multiply(pose_a, pose_b):
    """Multiply two pose matrices, A @ B.

    Args:
        pose_a: Left pose matrix, usually a transformation applied to the right.
        pose_b: Right pose matrix, usually a camera pose that will be transformed by pose_a.

    Returns:
        Camera pose matrix where pose_a was applied to pose_b.
    """
    R1, t1 = pose_a[..., :3, :3], pose_a[..., :3, 3:]
    R2, t2 = pose_b[..., :3, :3], pose_b[..., :3, 3:]
    R = R1.matmul(R2)
    t = t1 + R1.matmul(t2)
    return torch.cat([R, t], dim=-1)


class LowrankModel(nn.Module):
    def __init__(self,
                 grid_config: Union[str, List[Dict]],
                 # boolean flags
                 is_ndc: bool,
                 is_contracted: bool,
                 aabb: torch.Tensor,
                 # Model arguments
                 multiscale_res: Sequence[int],
                 density_activation: Optional[str] = 'trunc_exp',
                 concat_features_across_scales: bool = False,
                 linear_decoder: bool = True,
                 linear_decoder_layers: Optional[int] = 1,
                 # Spatial distortion
                 global_translation: Optional[torch.Tensor] = None,
                 global_scale: Optional[torch.Tensor] = None,
                 # proposal-sampling arguments
                 num_proposal_iterations: int = 1,
                 use_same_proposal_network: bool = False,
                 proposal_net_args_list: List[Dict] = None,
                 num_proposal_samples: Optional[Tuple[int]] = None,
                 num_samples: Optional[int] = None,
                 single_jitter: bool = False,
                 proposal_warmup: int = 5000,
                 proposal_update_every: int = 5,
                 use_proposal_weight_anneal: bool = True,
                 proposal_weights_anneal_max_num_iters: int = 1000,
                 proposal_weights_anneal_slope: float = 10.0,
                 # appearance embedding (phototourism)
                 use_appearance_embedding: bool = False,
                 appearance_embedding_dim: int = 0,
                 num_images: Optional[int] = None,
                 opaque_background: bool = False,
                 **kwargs,
                 ):
        super().__init__()
        if isinstance(grid_config, str):
            self.config: List[Dict] = eval(grid_config)
        else:
            self.config: List[Dict] = grid_config
        self.multiscale_res = multiscale_res
        self.is_ndc = is_ndc
        self.is_contracted = is_contracted
        self.concat_features_across_scales = concat_features_across_scales
        self.linear_decoder = linear_decoder
        self.linear_decoder_layers = linear_decoder_layers
        self.density_act = init_density_activation(density_activation)
        self.timer = CudaTimer(enabled=False)

        self.spatial_distortion: Optional[SpatialDistortion] = None
        if self.is_contracted:
            self.spatial_distortion = SceneContraction(
                order=float('inf'), global_scale=global_scale,
                global_translation=global_translation)

        self.field = KPlaneField(
            aabb,
            grid_config=self.config,
            concat_features_across_scales=self.concat_features_across_scales,
            multiscale_res=self.multiscale_res,
            use_appearance_embedding=use_appearance_embedding,
            appearance_embedding_dim=appearance_embedding_dim,
            spatial_distortion=self.spatial_distortion,
            density_activation=self.density_act,
            linear_decoder=self.linear_decoder,
            linear_decoder_layers=self.linear_decoder_layers,
            num_images=num_images,
            lagrangian=kwargs.get('lagrangian')
        )
        if num_images:
            self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_images, 6)))

        # Initialize proposal-sampling nets
        self.density_fns = []
        self.num_proposal_iterations = num_proposal_iterations
        self.proposal_net_args_list = proposal_net_args_list
        self.proposal_warmup = proposal_warmup
        self.proposal_update_every = proposal_update_every
        self.use_proposal_weight_anneal = use_proposal_weight_anneal
        self.proposal_weights_anneal_max_num_iters = proposal_weights_anneal_max_num_iters
        self.proposal_weights_anneal_slope = proposal_weights_anneal_slope
        self.proposal_networks = torch.nn.ModuleList()
        if use_same_proposal_network:
            assert len(self.proposal_net_args_list) == 1, "Only one proposal network is allowed."
            prop_net_args = self.proposal_net_args_list[0]
            network = KPlaneDensityField(
                aabb, spatial_distortion=self.spatial_distortion,
                density_activation=self.density_act, linear_decoder=self.linear_decoder, **prop_net_args)
            self.proposal_networks.append(network)
            self.density_fns.extend([network.get_density for _ in range(self.num_proposal_iterations)])
        else:
            for i in range(self.num_proposal_iterations):
                prop_net_args = self.proposal_net_args_list[min(i, len(self.proposal_net_args_list) - 1)]
                network = KPlaneDensityField(
                    aabb, spatial_distortion=self.spatial_distortion,
                    density_activation=self.density_act, linear_decoder=self.linear_decoder, **prop_net_args,
                )
                self.proposal_networks.append(network)
            self.density_fns.extend([network.get_density for network in self.proposal_networks])

        update_schedule = lambda step: np.clip(
            np.interp(step, [0, self.proposal_warmup], [0, self.proposal_update_every]),
            1,
            self.proposal_update_every,
        )
        if self.is_contracted or self.is_ndc:
            initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
        else:
            initial_sampler = UniformSampler(single_jitter=single_jitter)

        self.proposal_sampler = ProposalNetworkSampler(
            num_nerf_samples_per_ray=num_samples,
            num_proposal_samples_per_ray=num_proposal_samples,
            num_proposal_network_iterations=self.num_proposal_iterations,
            single_jitter=single_jitter,
            update_sched=update_schedule,
            initial_sampler=initial_sampler
        )
        if not isinstance(opaque_background, bool):
            opaque_background = eval(opaque_background)
        self.opaque_background = opaque_background

    def step_before_iter(self, step):
        if self.use_proposal_weight_anneal:
            # anneal the weights of the proposal network before doing PDF sampling
            N = self.proposal_weights_anneal_max_num_iters
            # https://arxiv.org/pdf/2111.12077.pdf eq. 18
            train_frac = np.clip(step / N, 0, 1)
            bias = lambda x, b: (b * x) / ((b - 1) * x + 1)
            anneal = bias(train_frac, self.proposal_weights_anneal_slope)
            self.proposal_sampler.set_anneal(anneal)

    def step_after_iter(self, step):
        if self.use_proposal_weight_anneal:
            self.proposal_sampler.step_cb(step)

    @staticmethod
    def render_rgb(rgb: torch.Tensor, weights: torch.Tensor, bg_color: Optional[torch.Tensor]):
        comp_rgb = torch.sum(weights * rgb, dim=-2)
        accumulated_weight = torch.sum(weights, dim=-2)
        if bg_color is None:
            pass
        else:
            comp_rgb = comp_rgb + (1.0 - accumulated_weight) * bg_color
        return comp_rgb

    @staticmethod
    def render_depth(weights: torch.Tensor, ray_samples: RaySamples, rays_d: torch.Tensor):
        steps = (ray_samples.starts + ray_samples.ends) / 2
        one_minus_transmittance = torch.sum(weights, dim=-2)
        depth = torch.sum(weights * steps, dim=-2) + one_minus_transmittance * rays_d[..., -1:]
        return depth

    @staticmethod
    def render_accumulation(weights: torch.Tensor):
        accumulation = torch.sum(weights, dim=-2)
        return accumulation

    def forward(self, rays_o, rays_d, bg_color, near_far: torch.Tensor,
                timestamps=None, c2w=None, camera_dirs=None, indices=None):
        """
        rays_o : [batch, 3]
        rays_d : [batch, 3]
        timestamps : [batch]
        near_far : [batch, 2]
        """
        # Fix shape for near-far
        nears, fars = torch.split(near_far, [1, 1], dim=-1)
        if nears.shape[0] != rays_o.shape[0]:
            ones = torch.ones_like(rays_o[..., 0:1])
            nears = ones * nears
            fars = ones * fars

        if indices is not None:
            camera_opt_to_camera = exp_map_SO3xR3(self.pose_adjustment[indices, :])
            c2w = c2w.to(camera_opt_to_camera.device)
            camera_dirs = camera_dirs.to(camera_opt_to_camera.device)
            c2w = pose_multiply(c2w, camera_opt_to_camera)
            directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1)
            rays_d = directions / torch.linalg.norm(directions, dim=-1, keepdims=True)
            rays_o = torch.broadcast_to(c2w[:, :3, -1], rays_d.shape)

        ray_bundle = RayBundle(origins=rays_o, directions=rays_d, nears=nears, fars=fars)
        # Note: proposal sampler mustn't use timestamps (=camera-IDs) with appearance-embedding,
        #       since the appearance embedding should not affect density. We still pass them in the
        #       call below, but they will not be used as long as density-field resolutions
        #       are be 3D.

        ray_samples, weights_list, ray_samples_list = self.proposal_sampler.generate_ray_samples(
            ray_bundle, timestamps=timestamps, density_fns=self.density_fns)

        field_out = self.field(ray_samples.get_positions(), ray_bundle.directions, timestamps)
        rgb, density = field_out["rgb"], field_out["density"]

        weights = ray_samples.get_weights(density, opaque_background=self.opaque_background)
        weights_list.append(weights)
        ray_samples_list.append(ray_samples)

        rgb = self.render_rgb(rgb=rgb, weights=weights, bg_color=bg_color)
        depth = self.render_depth(weights=weights, ray_samples=ray_samples, rays_d=ray_bundle.directions)
        accumulation = self.render_accumulation(weights=weights)
        outputs = {
            "rgb": rgb,
            "accumulation": accumulation,
            "depth": depth,
        }

        # These use a lot of GPU memory, so we avoid storing them for eval.
        if self.training:
            outputs["weights_list"] = weights_list
            outputs["ray_samples_list"] = ray_samples_list
        for i in range(self.num_proposal_iterations):
            outputs[f"prop_depth_{i}"] = self.render_depth(
                weights=weights_list[i], ray_samples=ray_samples_list[i], rays_d=ray_bundle.directions)
        return outputs


    def get_params(self, lr: float):
        model_params = self.field.get_params()
        pn_params = [pn.get_params() for pn in self.proposal_networks]
        field_params = model_params["field"] + [p for pnp in pn_params for p in pnp["field"]]
        nn_params = model_params["nn"] + [p for pnp in pn_params for p in pnp["nn"]]
        other_params = model_params["other"] + [p for pnp in pn_params for p in pnp["other"]]
        return [
            {"params": field_params, "lr": lr},
            {"params": nn_params, "lr": lr},
            {"params": other_params, "lr": lr},
        ]

    def partial_forward_0(self, rays_o, rays_d, bg_color, near_far: torch.Tensor, timestamps=None):
        """
        rays_o : [batch, 3]
        rays_d : [batch, 3]
        timestamps : [batch]
        near_far : [batch, 2]
        """
        # Fix shape for near-far
        nears, fars = torch.split(near_far, [1, 1], dim=-1)
        if nears.shape[0] != rays_o.shape[0]:
            ones = torch.ones_like(rays_o[..., 0:1])
            nears = ones * nears
            fars = ones * fars

        ray_bundle = RayBundle(origins=rays_o, directions=rays_d, nears=nears, fars=fars)
        # Note: proposal sampler mustn't use timestamps (=camera-IDs) with appearance-embedding,
        #       since the appearance embedding should not affect density. We still pass them in the
        #       call below, but they will not be used as long as density-field resolutions
        #       are be 3D.
        ray_samples, weights_list, ray_samples_list = self.proposal_sampler.generate_ray_samples(
            ray_bundle, timestamps=timestamps, density_fns=self.density_fns)

        return ray_samples, ray_bundle


    def partial_forward_1(self, ray_samples, ray_bundle):
        field_out = self.field(ray_samples.get_positions(), ray_bundle.directions)
        rgbs, density = field_out["rgb"], field_out["density"]
        weights = ray_samples.get_weights(density)
        
        return rgbs, weights


    def partial_forward_2(self, rgbs, weights, bg_color):
        rgb = self.render_rgb(rgb=rgbs, weights=weights, bg_color=bg_color)

        return rgb


    def partial_forward_lagrangian_1(self, pts):
        if self.field.spatial_distortion is not None:
            pts = self.spatial_distortion(pts)
            pts = pts / 2  # from [-2, 2] to [-1, 1]
        else:
            pts = normalize_aabb(pts, self.field.aabb)

        pts = pts.reshape(-1, pts.shape[-1])

        delta = self.field.motion_net(pts)

        return delta, pts

    def partial_forward_lagrangian_2(self, pts, directions, n_rays, n_samples, ray_samples, bg_color):

        features = interpolate_ms_features(
            pts, ms_grids=self.field.grids,  # noqa
            grid_dimensions=self.field.grid_config[0]["grid_dimensions"],
            concat_features=self.field.concat_features, num_levels=None)
        if len(features) < 1:
            features = torch.zeros((0, 1)).to(features.device)
        if self.field.linear_decoder:
            density_before_activation = self.sigma_net(features)  # [batch, 1]
        else:
            features = self.field.sigma_net(features)
            features, density_before_activation = torch.split(
                features, [self.field.geo_feat_dim, 1], dim=-1)

        density = self.field.density_activation(
            density_before_activation.to(pts)
        ).view(n_rays, n_samples, 1)

        directions = (directions + 1.0) / 2.0
        encoded_directions = self.field.direction_encoder(directions)

        color_features = [encoded_directions, features.view(-1, self.field.geo_feat_dim)]

        color_features = torch.cat(color_features, dim=-1)
        rgb = self.field.color_net(color_features).to(directions).view(n_rays, n_samples, 3)

        weights = ray_samples.get_weights(density)

        rgb = self.render_rgb(rgb=rgb, weights=weights, bg_color=bg_color)

        return rgb

magnerf/models/lowrank_model.py
3d-motion-magnification/3d-motion-mag
import torch
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd

__all__ = (
    "trunc_exp",
    "init_density_activation",
)

from torch.nn import functional as F


class TruncatedExponential(Function):  # pylint: disable=abstract-method
    # Implementation from torch-ngp:
    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, x):  # pylint: disable=arguments-differ
        ctx.save_for_backward(x)
        return torch.exp(x)

    @staticmethod
    @custom_bwd
    def backward(ctx, g):  # pylint: disable=arguments-differ
        x = ctx.saved_tensors[0]
        return g * torch.exp(torch.clamp(x, min=-15, max=15))


trunc_exp = TruncatedExponential.apply


def init_density_activation(activation_type: str):
    if activation_type == 'trunc_exp':
        return lambda x: trunc_exp(x - 1)
    elif activation_type == 'relu':
        return F.relu
    else:
        raise ValueError(activation_type)

magnerf/ops/activations.py
3d-motion-magnification/3d-motion-mag
""" Functions for image loading, saving, and color map changes """
#################################################################################
# Copyright (c) 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# SPDX-FileCopyrightText: Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES
# SPDX-License-Identifier: BSD-3-Clause
#################################################################################

# Visualizing and Communicating Errors in Rendered Images
# Ray Tracing Gems II, 2021,
# by Pontus Andersson, Jim Nilsson, and Tomas Akenine-Moller.
# Pointer to the chapter: https://research.nvidia.com/publication/2021-08_Visualizing-and-Communicating.

# Visualizing Errors in Rendered High Dynamic Range Images
# Eurographics 2021,
# by Pontus Andersson, Jim Nilsson, Peter Shirley, and Tomas Akenine-Moller.
# Pointer to the paper: https://research.nvidia.com/publication/2021-05_HDR-FLIP.

# FLIP: A Difference Evaluator for Alternating Images
# High Performance Graphics 2020,
# by Pontus Andersson, Jim Nilsson, Tomas Akenine-Moller,
# Magnus Oskarsson, Kalle Astrom, and Mark D. Fairchild.
# Pointer to the paper: https://research.nvidia.com/publication/2020-07_FLIP.

# Code by Pontus Andersson, Jim Nilsson, and Tomas Akenine-Moller.

import numpy as np
from PIL import Image


def HWCtoCHW(x):
    """
    Transforms an image from HxWxC layout to CxHxW

    :param x: image with HxWxC layout
    :return: image with CxHxW layout
    """
    return np.rollaxis(x, 2)


def CHWtoHWC(x):
    """
    Transforms an image from CxHxW layout to HxWxC

    :param x: image with CxHxW layout
    :return: image with HxWxC layout
    """
    return np.swapaxes(np.swapaxes(x, 0, 1), 1, 2)


def save_image(img_file, img):
    """
    Saves image as png

    :param img_file: image's filename
    :param img: float image (in the [0,1] range) to save
    """
    img_array = (np.round(np.clip(img, 0.0, 1.0) * 255.0)).astype(np.uint8)
    im = Image.fromarray(img_array)
    im.save(img_file)


def load_image_array(img_file):
    """
    Loads an image (in the [0, 255] range) and transforms it into a numpy array and into the [0, 1] range

    :param img_file: image's filename
    :return: float image (in the [0,1] range) on CxHxW layout
    """
    img = Image.open(img_file, 'r').convert('RGB')
    img = np.asarray(img).astype(np.float32)
    img = HWCtoCHW(img)
    img = img / 255.0
    return img


def read_exr(filename):
    """
    Read color data from EXR image file. Set negative values and nans to 0.

    :param filename: string describing file path
    :return: RGB image in float32 format (with HxWxC layout)
    """
    import OpenEXR as exr
    import Imath

    exrfile = exr.InputFile(filename)
    header = exrfile.header()

    dw = header['dataWindow']
    isize = (dw.max.y - dw.min.y + 1, dw.max.x - dw.min.x + 1)

    channelData = dict()

    # Convert all channels in the image to numpy arrays
    for c in header['channels']:
        C = exrfile.channel(c, Imath.PixelType(Imath.PixelType.FLOAT))
        C = np.frombuffer(C, dtype=np.float32)
        C = np.reshape(C, isize)

        channelData[c] = C

    if len(channelData) == 1:
        channelData['R'] = channelData['G'] = channelData['B'] = channelData[
            next(iter(channelData))]

    colorChannels = ['R', 'G', 'B']
    img = np.concatenate([channelData[c][..., np.newaxis] for c in colorChannels], axis=2)

    return np.maximum(np.nan_to_num(np.array(img)),
                      0.0)  # added maximum to avoid negative values in images


def index2color(index_map, color_map):
    """
    Transforms grayscale index map to colors in a color map

    :param index_map: integer matrix containing integer versions of the FLIP errors, scaled to [0,255]
    :param color_map: matrix of size 256x3 where each row contains a color
    :return: error map in color map colors
    """
    dim = index_map.shape
    index_map = index_map.flatten().astype(int)
    column_stacked_colors = color_map[index_map, :]
    heat_map = np.reshape(column_stacked_colors.transpose(), (3, dim[0], dim[1]))
    return heat_map


def get_magma_map():
    """
    Return the magma map, described here: https://bids.github.io/colormap/

    :return: matrix of size 256x3 where each row contains a color in the magma map
    """
    color_map = [[0.001462, 0.000466, 0.013866],
                 [0.002258, 0.001295, 0.018331],
                 [0.003279, 0.002305, 0.023708],
                 [0.004512, 0.003490, 0.029965],
                 [0.005950, 0.004843, 0.037130],
                 [0.007588, 0.006356, 0.044973],
                 [0.009426, 0.008022, 0.052844],
                 [0.011465, 0.009828, 0.060750],
                 [0.013708, 0.011771, 0.068667],
                 [0.016156, 0.013840, 0.076603],
                 [0.018815, 0.016026, 0.084584],
                 [0.021692, 0.018320, 0.092610],
                 [0.024792, 0.020715, 0.100676],
                 [0.028123, 0.023201, 0.108787],
                 [0.031696, 0.025765, 0.116965],
                 [0.035520, 0.028397, 0.125209],
                 [0.039608, 0.031090, 0.133515],
                 [0.043830, 0.033830, 0.141886],
                 [0.048062, 0.036607, 0.150327],
                 [0.052320, 0.039407, 0.158841],
                 [0.056615, 0.042160, 0.167446],
                 [0.060949, 0.044794, 0.176129],
                 [0.065330, 0.047318, 0.184892],
                 [0.069764, 0.049726, 0.193735],
                 [0.074257, 0.052017, 0.202660],
                 [0.078815, 0.054184, 0.211667],
                 [0.083446, 0.056225, 0.220755],
                 [0.088155, 0.058133, 0.229922],
                 [0.092949, 0.059904, 0.239164],
                 [0.097833, 0.061531, 0.248477],
                 [0.102815, 0.063010, 0.257854],
                 [0.107899, 0.064335, 0.267289],
                 [0.113094, 0.065492, 0.276784],
                 [0.118405, 0.066479, 0.286321],
                 [0.123833, 0.067295, 0.295879],
                 [0.129380, 0.067935, 0.305443],
                 [0.135053, 0.068391, 0.315000],
                 [0.140858, 0.068654, 0.324538],
                 [0.146785, 0.068738, 0.334011],
                 [0.152839, 0.068637, 0.343404],
                 [0.159018, 0.068354, 0.352688],
                 [0.165308, 0.067911, 0.361816],
                 [0.171713, 0.067305, 0.370771],
                 [0.178212, 0.066576, 0.379497],
                 [0.184801, 0.065732, 0.387973],
                 [0.191460, 0.064818, 0.396152],
                 [0.198177, 0.063862, 0.404009],
                 [0.204935, 0.062907, 0.411514],
                 [0.211718, 0.061992, 0.418647],
                 [0.218512, 0.061158, 0.425392],
                 [0.225302, 0.060445, 0.431742],
                 [0.232077, 0.059889, 0.437695],
                 [0.238826, 0.059517, 0.443256],
                 [0.245543, 0.059352, 0.448436],
                 [0.252220, 0.059415, 0.453248],
                 [0.258857, 0.059706, 0.457710],
                 [0.265447, 0.060237, 0.461840],
                 [0.271994, 0.060994, 0.465660],
                 [0.278493, 0.061978, 0.469190],
                 [0.284951, 0.063168, 0.472451],
                 [0.291366, 0.064553, 0.475462],
                 [0.297740, 0.066117, 0.478243],
                 [0.304081, 0.067835, 0.480812],
                 [0.310382, 0.069702, 0.483186],
                 [0.316654, 0.071690, 0.485380],
                 [0.322899, 0.073782, 0.487408],
                 [0.329114, 0.075972, 0.489287],
                 [0.335308, 0.078236, 0.491024],
                 [0.341482, 0.080564, 0.492631],
                 [0.347636, 0.082946, 0.494121],
                 [0.353773, 0.085373, 0.495501],
                 [0.359898, 0.087831, 0.496778],
                 [0.366012, 0.090314, 0.497960],
                 [0.372116, 0.092816, 0.499053],
                 [0.378211, 0.095332, 0.500067],
                 [0.384299, 0.097855, 0.501002],
                 [0.390384, 0.100379, 0.501864],
                 [0.396467, 0.102902, 0.502658],
                 [0.402548, 0.105420, 0.503386],
                 [0.408629, 0.107930, 0.504052],
                 [0.414709, 0.110431, 0.504662],
                 [0.420791, 0.112920, 0.505215],
                 [0.426877, 0.115395, 0.505714],
                 [0.432967, 0.117855, 0.506160],
                 [0.439062, 0.120298, 0.506555],
                 [0.445163, 0.122724, 0.506901],
                 [0.451271, 0.125132, 0.507198],
                 [0.457386, 0.127522, 0.507448],
                 [0.463508, 0.129893, 0.507652],
                 [0.469640, 0.132245, 0.507809],
                 [0.475780, 0.134577, 0.507921],
                 [0.481929, 0.136891, 0.507989],
                 [0.488088, 0.139186, 0.508011],
                 [0.494258, 0.141462, 0.507988],
                 [0.500438, 0.143719, 0.507920],
                 [0.506629, 0.145958, 0.507806],
                 [0.512831, 0.148179, 0.507648],
                 [0.519045, 0.150383, 0.507443],
                 [0.525270, 0.152569, 0.507192],
                 [0.531507, 0.154739, 0.506895],
                 [0.537755, 0.156894, 0.506551],
                 [0.544015, 0.159033, 0.506159],
                 [0.550287, 0.161158, 0.505719],
                 [0.556571, 0.163269, 0.505230],
                 [0.562866, 0.165368, 0.504692],
                 [0.569172, 0.167454, 0.504105],
                 [0.575490, 0.169530, 0.503466],
                 [0.581819, 0.171596, 0.502777],
                 [0.588158, 0.173652, 0.502035],
                 [0.594508, 0.175701, 0.501241],
                 [0.600868, 0.177743, 0.500394],
                 [0.607238, 0.179779, 0.499492],
                 [0.613617, 0.181811, 0.498536],
                 [0.620005, 0.183840, 0.497524],
                 [0.626401, 0.185867, 0.496456],
                 [0.632805, 0.187893, 0.495332],
                 [0.639216, 0.189921, 0.494150],
                 [0.645633, 0.191952, 0.492910],
                 [0.652056, 0.193986, 0.491611],
                 [0.658483, 0.196027, 0.490253],
                 [0.664915, 0.198075, 0.488836],
                 [0.671349, 0.200133, 0.487358],
                 [0.677786, 0.202203, 0.485819],
                 [0.684224, 0.204286, 0.484219],
                 [0.690661, 0.206384, 0.482558],
                 [0.697098, 0.208501, 0.480835],
                 [0.703532, 0.210638, 0.479049],
                 [0.709962, 0.212797, 0.477201],
                 [0.716387, 0.214982, 0.475290],
                 [0.722805, 0.217194, 0.473316],
                 [0.729216, 0.219437, 0.471279],
                 [0.735616, 0.221713, 0.469180],
                 [0.742004, 0.224025, 0.467018],
                 [0.748378, 0.226377, 0.464794],
                 [0.754737, 0.228772, 0.462509],
                 [0.761077, 0.231214, 0.460162],
                 [0.767398, 0.233705, 0.457755],
                 [0.773695, 0.236249, 0.455289],
                 [0.779968, 0.238851, 0.452765],
                 [0.786212, 0.241514, 0.450184],
                 [0.792427, 0.244242, 0.447543],
                 [0.798608, 0.247040, 0.444848],
                 [0.804752, 0.249911, 0.442102],
                 [0.810855, 0.252861, 0.439305],
                 [0.816914, 0.255895, 0.436461],
                 [0.822926, 0.259016, 0.433573],
                 [0.828886, 0.262229, 0.430644],
                 [0.834791, 0.265540, 0.427671],
                 [0.840636, 0.268953, 0.424666],
                 [0.846416, 0.272473, 0.421631],
                 [0.852126, 0.276106, 0.418573],
                 [0.857763, 0.279857, 0.415496],
                 [0.863320, 0.283729, 0.412403],
                 [0.868793, 0.287728, 0.409303],
                 [0.874176, 0.291859, 0.406205],
                 [0.879464, 0.296125, 0.403118],
                 [0.884651, 0.300530, 0.400047],
                 [0.889731, 0.305079, 0.397002],
                 [0.894700, 0.309773, 0.393995],
                 [0.899552, 0.314616, 0.391037],
                 [0.904281, 0.319610, 0.388137],
                 [0.908884, 0.324755, 0.385308],
                 [0.913354, 0.330052, 0.382563],
                 [0.917689, 0.335500, 0.379915],
                 [0.921884, 0.341098, 0.377376],
                 [0.925937, 0.346844, 0.374959],
                 [0.929845, 0.352734, 0.372677],
                 [0.933606, 0.358764, 0.370541],
                 [0.937221, 0.364929, 0.368567],
                 [0.940687, 0.371224, 0.366762],
                 [0.944006, 0.377643, 0.365136],
                 [0.947180, 0.384178, 0.363701],
                 [0.950210, 0.390820, 0.362468],
                 [0.953099, 0.397563, 0.361438],
                 [0.955849, 0.404400, 0.360619],
                 [0.958464, 0.411324, 0.360014],
                 [0.960949, 0.418323, 0.359630],
                 [0.963310, 0.425390, 0.359469],
                 [0.965549, 0.432519, 0.359529],
                 [0.967671, 0.439703, 0.359810],
                 [0.969680, 0.446936, 0.360311],
                 [0.971582, 0.454210, 0.361030],
                 [0.973381, 0.461520, 0.361965],
                 [0.975082, 0.468861, 0.363111],
                 [0.976690, 0.476226, 0.364466],
                 [0.978210, 0.483612, 0.366025],
                 [0.979645, 0.491014, 0.367783],
                 [0.981000, 0.498428, 0.369734],
                 [0.982279, 0.505851, 0.371874],
                 [0.983485, 0.513280, 0.374198],
                 [0.984622, 0.520713, 0.376698],
                 [0.985693, 0.528148, 0.379371],
                 [0.986700, 0.535582, 0.382210],
                 [0.987646, 0.543015, 0.385210],
                 [0.988533, 0.550446, 0.388365],
                 [0.989363, 0.557873, 0.391671],
                 [0.990138, 0.565296, 0.395122],
                 [0.990871, 0.572706, 0.398714],
                 [0.991558, 0.580107, 0.402441],
                 [0.992196, 0.587502, 0.406299],
                 [0.992785, 0.594891, 0.410283],
                 [0.993326, 0.602275, 0.414390],
                 [0.993834, 0.609644, 0.418613],
                 [0.994309, 0.616999, 0.422950],
                 [0.994738, 0.624350, 0.427397],
                 [0.995122, 0.631696, 0.431951],
                 [0.995480, 0.639027, 0.436607],
                 [0.995810, 0.646344, 0.441361],
                 [0.996096, 0.653659, 0.446213],
                 [0.996341, 0.660969, 0.451160],
                 [0.996580, 0.668256, 0.456192],
                 [0.996775, 0.675541, 0.461314],
                 [0.996925, 0.682828, 0.466526],
                 [0.997077, 0.690088, 0.471811],
                 [0.997186, 0.697349, 0.477182],
                 [0.997254, 0.704611, 0.482635],
                 [0.997325, 0.711848, 0.488154],
                 [0.997351, 0.719089, 0.493755],
                 [0.997351, 0.726324, 0.499428],
                 [0.997341, 0.733545, 0.505167],
                 [0.997285, 0.740772, 0.510983],
                 [0.997228, 0.747981, 0.516859],
                 [0.997138, 0.755190, 0.522806],
                 [0.997019, 0.762398, 0.528821],
                 [0.996898, 0.769591, 0.534892],
                 [0.996727, 0.776795, 0.541039],
                 [0.996571, 0.783977, 0.547233],
                 [0.996369, 0.791167, 0.553499],
                 [0.996162, 0.798348, 0.559820],
                 [0.995932, 0.805527, 0.566202],
                 [0.995680, 0.812706, 0.572645],
                 [0.995424, 0.819875, 0.579140],
                 [0.995131, 0.827052, 0.585701],
                 [0.994851, 0.834213, 0.592307],
                 [0.994524, 0.841387, 0.598983],
                 [0.994222, 0.848540, 0.605696],
                 [0.993866, 0.855711, 0.612482],
                 [0.993545, 0.862859, 0.619299],
                 [0.993170, 0.870024, 0.626189],
                 [0.992831, 0.877168, 0.633109],
                 [0.992440, 0.884330, 0.640099],
                 [0.992089, 0.891470, 0.647116],
                 [0.991688, 0.898627, 0.654202],
                 [0.991332, 0.905763, 0.661309],
                 [0.990930, 0.912915, 0.668481],
                 [0.990570, 0.920049, 0.675675],
                 [0.990175, 0.927196, 0.682926],
                 [0.989815, 0.934329, 0.690198],
                 [0.989434, 0.941470, 0.697519],
                 [0.989077, 0.948604, 0.704863],
                 [0.988717, 0.955742, 0.712242],
                 [0.988367, 0.962878, 0.719649],
                 [0.988033, 0.970012, 0.727077],
                 [0.987691, 0.977154, 0.734536],
                 [0.987387, 0.984288, 0.742002],
                 [0.987053, 0.991438, 0.749504]]
    return np.asarray(color_map)


def get_viridis_map():
    """
    Return the viridis map, described here: https://bids.github.io/colormap/

    :return: matrix of size 256x3 where each row contains a color in the virdis map
    """
    color_map = [[0.267004, 0.004874, 0.329415],
                 [0.268510, 0.009605, 0.335427],
                 [0.269944, 0.014625, 0.341379],
                 [0.271305, 0.019942, 0.347269],
                 [0.272594, 0.025563, 0.353093],
                 [0.273809, 0.031497, 0.358853],
                 [0.274952, 0.037752, 0.364543],
                 [0.276022, 0.044167, 0.370164],
                 [0.277018, 0.050344, 0.375715],
                 [0.277941, 0.056324, 0.381191],
                 [0.278791, 0.062145, 0.386592],
                 [0.279566, 0.067836, 0.391917],
                 [0.280267, 0.073417, 0.397163],
                 [0.280894, 0.078907, 0.402329],
                 [0.281446, 0.084320, 0.407414],
                 [0.281924, 0.089666, 0.412415],
                 [0.282327, 0.094955, 0.417331],
                 [0.282656, 0.100196, 0.422160],
                 [0.282910, 0.105393, 0.426902],
                 [0.283091, 0.110553, 0.431554],
                 [0.283197, 0.115680, 0.436115],
                 [0.283229, 0.120777, 0.440584],
                 [0.283187, 0.125848, 0.444960],
                 [0.283072, 0.130895, 0.449241],
                 [0.282884, 0.135920, 0.453427],
                 [0.282623, 0.140926, 0.457517],
                 [0.282290, 0.145912, 0.461510],
                 [0.281887, 0.150881, 0.465405],
                 [0.281412, 0.155834, 0.469201],
                 [0.280868, 0.160771, 0.472899],
                 [0.280255, 0.165693, 0.476498],
                 [0.279574, 0.170599, 0.479997],
                 [0.278826, 0.175490, 0.483397],
                 [0.278012, 0.180367, 0.486697],
                 [0.277134, 0.185228, 0.489898],
                 [0.276194, 0.190074, 0.493001],
                 [0.275191, 0.194905, 0.496005],
                 [0.274128, 0.199721, 0.498911],
                 [0.273006, 0.204520, 0.501721],
                 [0.271828, 0.209303, 0.504434],
                 [0.270595, 0.214069, 0.507052],
                 [0.269308, 0.218818, 0.509577],
                 [0.267968, 0.223549, 0.512008],
                 [0.266580, 0.228262, 0.514349],
                 [0.265145, 0.232956, 0.516599],
                 [0.263663, 0.237631, 0.518762],
                 [0.262138, 0.242286, 0.520837],
                 [0.260571, 0.246922, 0.522828],
                 [0.258965, 0.251537, 0.524736],
                 [0.257322, 0.256130, 0.526563],
                 [0.255645, 0.260703, 0.528312],
                 [0.253935, 0.265254, 0.529983],
                 [0.252194, 0.269783, 0.531579],
                 [0.250425, 0.274290, 0.533103],
                 [0.248629, 0.278775, 0.534556],
                 [0.246811, 0.283237, 0.535941],
                 [0.244972, 0.287675, 0.537260],
                 [0.243113, 0.292092, 0.538516],
                 [0.241237, 0.296485, 0.539709],
                 [0.239346, 0.300855, 0.540844],
                 [0.237441, 0.305202, 0.541921],
                 [0.235526, 0.309527, 0.542944],
                 [0.233603, 0.313828, 0.543914],
                 [0.231674, 0.318106, 0.544834],
                 [0.229739, 0.322361, 0.545706],
                 [0.227802, 0.326594, 0.546532],
                 [0.225863, 0.330805, 0.547314],
                 [0.223925, 0.334994, 0.548053],
                 [0.221989, 0.339161, 0.548752],
                 [0.220057, 0.343307, 0.549413],
                 [0.218130, 0.347432, 0.550038],
                 [0.216210, 0.351535, 0.550627],
                 [0.214298, 0.355619, 0.551184],
                 [0.212395, 0.359683, 0.551710],
                 [0.210503, 0.363727, 0.552206],
                 [0.208623, 0.367752, 0.552675],
                 [0.206756, 0.371758, 0.553117],
                 [0.204903, 0.375746, 0.553533],
                 [0.203063, 0.379716, 0.553925],
                 [0.201239, 0.383670, 0.554294],
                 [0.199430, 0.387607, 0.554642],
                 [0.197636, 0.391528, 0.554969],
                 [0.195860, 0.395433, 0.555276],
                 [0.194100, 0.399323, 0.555565],
                 [0.192357, 0.403199, 0.555836],
                 [0.190631, 0.407061, 0.556089],
                 [0.188923, 0.410910, 0.556326],
                 [0.187231, 0.414746, 0.556547],
                 [0.185556, 0.418570, 0.556753],
                 [0.183898, 0.422383, 0.556944],
                 [0.182256, 0.426184, 0.557120],
                 [0.180629, 0.429975, 0.557282],
                 [0.179019, 0.433756, 0.557430],
                 [0.177423, 0.437527, 0.557565],
                 [0.175841, 0.441290, 0.557685],
                 [0.174274, 0.445044, 0.557792],
                 [0.172719, 0.448791, 0.557885],
                 [0.171176, 0.452530, 0.557965],
                 [0.169646, 0.456262, 0.558030],
                 [0.168126, 0.459988, 0.558082],
                 [0.166617, 0.463708, 0.558119],
                 [0.165117, 0.467423, 0.558141],
                 [0.163625, 0.471133, 0.558148],
                 [0.162142, 0.474838, 0.558140],
                 [0.160665, 0.478540, 0.558115],
                 [0.159194, 0.482237, 0.558073],
                 [0.157729, 0.485932, 0.558013],
                 [0.156270, 0.489624, 0.557936],
                 [0.154815, 0.493313, 0.557840],
                 [0.153364, 0.497000, 0.557724],
                 [0.151918, 0.500685, 0.557587],
                 [0.150476, 0.504369, 0.557430],
                 [0.149039, 0.508051, 0.557250],
                 [0.147607, 0.511733, 0.557049],
                 [0.146180, 0.515413, 0.556823],
                 [0.144759, 0.519093, 0.556572],
                 [0.143343, 0.522773, 0.556295],
                 [0.141935, 0.526453, 0.555991],
                 [0.140536, 0.530132, 0.555659],
                 [0.139147, 0.533812, 0.555298],
                 [0.137770, 0.537492, 0.554906],
                 [0.136408, 0.541173, 0.554483],
                 [0.135066, 0.544853, 0.554029],
                 [0.133743, 0.548535, 0.553541],
                 [0.132444, 0.552216, 0.553018],
                 [0.131172, 0.555899, 0.552459],
                 [0.129933, 0.559582, 0.551864],
                 [0.128729, 0.563265, 0.551229],
                 [0.127568, 0.566949, 0.550556],
                 [0.126453, 0.570633, 0.549841],
                 [0.125394, 0.574318, 0.549086],
                 [0.124395, 0.578002, 0.548287],
                 [0.123463, 0.581687, 0.547445],
                 [0.122606, 0.585371, 0.546557],
                 [0.121831, 0.589055, 0.545623],
                 [0.121148, 0.592739, 0.544641],
                 [0.120565, 0.596422, 0.543611],
                 [0.120092, 0.600104, 0.542530],
                 [0.119738, 0.603785, 0.541400],
                 [0.119512, 0.607464, 0.540218],
                 [0.119423, 0.611141, 0.538982],
                 [0.119483, 0.614817, 0.537692],
                 [0.119699, 0.618490, 0.536347],
                 [0.120081, 0.622161, 0.534946],
                 [0.120638, 0.625828, 0.533488],
                 [0.121380, 0.629492, 0.531973],
                 [0.122312, 0.633153, 0.530398],
                 [0.123444, 0.636809, 0.528763],
                 [0.124780, 0.640461, 0.527068],
                 [0.126326, 0.644107, 0.525311],
                 [0.128087, 0.647749, 0.523491],
                 [0.130067, 0.651384, 0.521608],
                 [0.132268, 0.655014, 0.519661],
                 [0.134692, 0.658636, 0.517649],
                 [0.137339, 0.662252, 0.515571],
                 [0.140210, 0.665859, 0.513427],
                 [0.143303, 0.669459, 0.511215],
                 [0.146616, 0.673050, 0.508936],
                 [0.150148, 0.676631, 0.506589],
                 [0.153894, 0.680203, 0.504172],
                 [0.157851, 0.683765, 0.501686],
                 [0.162016, 0.687316, 0.499129],
                 [0.166383, 0.690856, 0.496502],
                 [0.170948, 0.694384, 0.493803],
                 [0.175707, 0.697900, 0.491033],
                 [0.180653, 0.701402, 0.488189],
                 [0.185783, 0.704891, 0.485273],
                 [0.191090, 0.708366, 0.482284],
                 [0.196571, 0.711827, 0.479221],
                 [0.202219, 0.715272, 0.476084],
                 [0.208030, 0.718701, 0.472873],
                 [0.214000, 0.722114, 0.469588],
                 [0.220124, 0.725509, 0.466226],
                 [0.226397, 0.728888, 0.462789],
                 [0.232815, 0.732247, 0.459277],
                 [0.239374, 0.735588, 0.455688],
                 [0.246070, 0.738910, 0.452024],
                 [0.252899, 0.742211, 0.448284],
                 [0.259857, 0.745492, 0.444467],
                 [0.266941, 0.748751, 0.440573],
                 [0.274149, 0.751988, 0.436601],
                 [0.281477, 0.755203, 0.432552],
                 [0.288921, 0.758394, 0.428426],
                 [0.296479, 0.761561, 0.424223],
                 [0.304148, 0.764704, 0.419943],
                 [0.311925, 0.767822, 0.415586],
                 [0.319809, 0.770914, 0.411152],
                 [0.327796, 0.773980, 0.406640],
                 [0.335885, 0.777018, 0.402049],
                 [0.344074, 0.780029, 0.397381],
                 [0.352360, 0.783011, 0.392636],
                 [0.360741, 0.785964, 0.387814],
                 [0.369214, 0.788888, 0.382914],
                 [0.377779, 0.791781, 0.377939],
                 [0.386433, 0.794644, 0.372886],
                 [0.395174, 0.797475, 0.367757],
                 [0.404001, 0.800275, 0.362552],
                 [0.412913, 0.803041, 0.357269],
                 [0.421908, 0.805774, 0.351910],
                 [0.430983, 0.808473, 0.346476],
                 [0.440137, 0.811138, 0.340967],
                 [0.449368, 0.813768, 0.335384],
                 [0.458674, 0.816363, 0.329727],
                 [0.468053, 0.818921, 0.323998],
                 [0.477504, 0.821444, 0.318195],
                 [0.487026, 0.823929, 0.312321],
                 [0.496615, 0.826376, 0.306377],
                 [0.506271, 0.828786, 0.300362],
                 [0.515992, 0.831158, 0.294279],
                 [0.525776, 0.833491, 0.288127],
                 [0.535621, 0.835785, 0.281908],
                 [0.545524, 0.838039, 0.275626],
                 [0.555484, 0.840254, 0.269281],
                 [0.565498, 0.842430, 0.262877],
                 [0.575563, 0.844566, 0.256415],
                 [0.585678, 0.846661, 0.249897],
                 [0.595839, 0.848717, 0.243329],
                 [0.606045, 0.850733, 0.236712],
                 [0.616293, 0.852709, 0.230052],
                 [0.626579, 0.854645, 0.223353],
                 [0.636902, 0.856542, 0.216620],
                 [0.647257, 0.858400, 0.209861],
                 [0.657642, 0.860219, 0.203082],
                 [0.668054, 0.861999, 0.196293],
                 [0.678489, 0.863742, 0.189503],
                 [0.688944, 0.865448, 0.182725],
                 [0.699415, 0.867117, 0.175971],
                 [0.709898, 0.868751, 0.169257],
                 [0.720391, 0.870350, 0.162603],
                 [0.730889, 0.871916, 0.156029],
                 [0.741388, 0.873449, 0.149561],
                 [0.751884, 0.874951, 0.143228],
                 [0.762373, 0.876424, 0.137064],
                 [0.772852, 0.877868, 0.131109],
                 [0.783315, 0.879285, 0.125405],
                 [0.793760, 0.880678, 0.120005],
                 [0.804182, 0.882046, 0.114965],
                 [0.814576, 0.883393, 0.110347],
                 [0.824940, 0.884720, 0.106217],
                 [0.835270, 0.886029, 0.102646],
                 [0.845561, 0.887322, 0.099702],
                 [0.855810, 0.888601, 0.097452],
                 [0.866013, 0.889868, 0.095953],
                 [0.876168, 0.891125, 0.095250],
                 [0.886271, 0.892374, 0.095374],
                 [0.896320, 0.893616, 0.096335],
                 [0.906311, 0.894855, 0.098125],
                 [0.916242, 0.896091, 0.100717],
                 [0.926106, 0.897330, 0.104071],
                 [0.935904, 0.898570, 0.108131],
                 [0.945636, 0.899815, 0.112838],
                 [0.955300, 0.901065, 0.118128],
                 [0.964894, 0.902323, 0.123941],
                 [0.974417, 0.903590, 0.130215],
                 [0.983868, 0.904867, 0.136897],
                 [0.993248, 0.906157, 0.143936]]
    return np.asarray(color_map)

magnerf/ops/flip/data.py
3d-motion-magnification/3d-motion-mag
import glob
import os
from typing import List, Optional

import cv2
import torch
import torchvision.transforms.functional as TF
from PIL import Image
import logging as log
import numpy as np
import imageio.v3 as iio
import subprocess

def write_png(path, data):
    """Writes an PNG image to some path.

    Args:
        path (str): Path to save the PNG.
        data (np.array): HWC image.

    Returns:
        (void): Writes to path.
    """
    Image.fromarray(data).save(path)


def read_png(file_name: str, resize_h: Optional[int] = None, resize_w: Optional[int] = None) -> torch.Tensor:
    """Reads a PNG image from path, potentially resizing it.
    """
    img = Image.open(file_name).convert('RGB')  # PIL outputs BGR by default
    if resize_h is not None and resize_w is not None:
        img.resize((resize_w, resize_h), Image.LANCZOS)
    img = TF.to_tensor(img)  # TF converts to C, H, W
    img = img.permute(1, 2, 0).contiguous()  # H, W, C
    return img


def glob_imgs(path, exts=None):
    """Utility to find images in some path.

    Args:
        path (str): Path to search images in.
        exts (list of str): List of extensions to try.

    Returns:
        (list of str): List of paths that were found.
    """
    if exts is None:
        exts = ['*.png', '*.PNG', '*.jpg', '*.jpeg', '*.JPG', '*.JPEG']
    imgs = []
    for ext in exts:
        imgs.extend(glob.glob(os.path.join(path, ext)))
    return imgs


def _write_video_to_file(file_name, frames: List[np.ndarray]):
    log.info(f"Saving video ({len(frames)} frames) to {file_name}")
    # Photo tourism image sizes differ
    sizes = np.array([frame.shape[:2] for frame in frames])
    same_size_frames = np.unique(sizes, axis=0).shape[0] == 1
    if same_size_frames:
        height, width = frames[0].shape[:2]
        video = cv2.VideoWriter(
            file_name, cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
        for img in frames:
            if isinstance(img, torch.Tensor):
                img = img.numpy()
            video.write(img[:, :, ::-1])  # opencv uses BGR instead of RGB
        cv2.destroyAllWindows()
        video.release()
    else:
        height = sizes[:, 0].max()
        width = sizes[:, 1].max()
        video = cv2.VideoWriter(
            file_name, cv2.VideoWriter_fourcc(*'mp4v'), 5, (width, height))
        for img in frames:
            image = np.zeros((height, width, 3), dtype=np.uint8)
            h, w = img.shape[:2]
            if isinstance(img, torch.Tensor):
                img = img.numpy()
            image[(height-h)//2:(height-h)//2+h, (width-w)//2:(width-w)//2+w, :] = img
            video.write(image[:, :, ::-1])  # opencv uses BGR instead of RGB
        cv2.destroyAllWindows()
        video.release()

def write_video_to_file(file_name, frames: List[np.ndarray]):
    log.info(f"Saving video ({len(frames)} frames) to {file_name}")
    # Photo tourism image sizes differ
    sizes = np.array([frame.shape[:2] for frame in frames])
    same_size_frames = np.unique(sizes, axis=0).shape[0] == 1

    folder = '/'.join(file_name.split('/')[:-1])
    frame_dir = f'{folder}/frames'
    os.makedirs(frame_dir, exist_ok=True)

    if same_size_frames:
        for i, img in enumerate(frames):
            iio.imwrite(
                f"{frame_dir}/{i:03d}.png",
                img,
            )
    else:
        height = sizes[:, 0].max()
        width = sizes[:, 1].max()
        for i, img in enumerate(frames):
            image = np.zeros((height, width, 3), dtype=np.uint8)
            h, w = img.shape[:2]
            if isinstance(img, torch.Tensor):
                img = img.numpy()
            image[(height-h)//2:(height-h)//2+h, (width-w)//2:(width-w)//2+w, :] = img
            iio.imwrite(
                f"{frame_dir}/{i:03d}.png",
                img,
            )

    subprocess.call(['ffmpeg',
        '-framerate', '30', 
        '-i', f'{frame_dir}/%03d.png',
        '-vcodec', 'libx264',
        '-y',
        '-crf', '8',
        '-vf', 'pad=ceil(iw/2)*2:ceil(ih/2)*2',
        '-pix_fmt', 'yuv420p',
        f'{file_name}'
    ])


def read_mp4(file_name: str) -> List[torch.Tensor]:
    all_frames = iio.imread(
        file_name, plugin='pyav', format='rgb24', constant_framerate=True, thread_count=2
    )
    out_frames = [torch.from_numpy(f) for f in all_frames]
    return out_frames

magnerf/ops/image/io.py
3d-motion-magnification/3d-motion-mag
import math
import tempfile
import subprocess
import os
import re
from typing import List

import numpy as np
import skimage.metrics
import torch
from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure
import lpips

from .io import write_video_to_file, write_png


ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0)

""" A module for image based metrics """


def psnr(rgb, gts):
    """Calculate the PSNR metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.Tensor): Image tensor of shape [H,W3]

    Returns:
        (float): The PSNR score
    """
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    assert (rgb.shape[-1] == 3)
    assert (gts.shape[-1] == 3)

    mse = torch.mean((rgb[..., :3] - gts[..., :3]) ** 2).item()
    return 10 * math.log10(1.0 / mse)


def ssim(rgb, gts):
    """
    Modified from https://github.com/google/mipnerf/blob/16e73dfdb52044dcceb47cda5243a686391a6e0f/internal/math.py#L58
    """
    filter_size = 11
    filter_sigma = 1.5
    k1 = 0.01
    k2 = 0.03
    max_val = 1.0
    rgb = rgb.cpu().numpy()
    gts = gts.cpu().numpy()
    assert len(rgb.shape) == 3
    assert rgb.shape[-1] == 3
    assert rgb.shape == gts.shape
    import scipy.signal

    # Construct a 1D Gaussian blur filter.
    hw = filter_size // 2
    shift = (2 * hw - filter_size + 1) / 2
    f_i = ((np.arange(filter_size) - hw + shift) / filter_sigma)**2
    filt = np.exp(-0.5 * f_i)
    filt /= np.sum(filt)

    # Blur in x and y (faster than the 2D convolution).
    def convolve2d(z, f):
        return scipy.signal.convolve2d(z, f, mode='valid')

    filt_fn = lambda z: np.stack([
        convolve2d(convolve2d(z[..., i], filt[:, None]), filt[None, :])
        for i in range(z.shape[-1])], -1)
    mu0 = filt_fn(rgb)
    mu1 = filt_fn(gts)
    mu00 = mu0 * mu0
    mu11 = mu1 * mu1
    mu01 = mu0 * mu1
    sigma00 = filt_fn(rgb**2) - mu00
    sigma11 = filt_fn(gts**2) - mu11
    sigma01 = filt_fn(rgb * gts) - mu01

    # Clip the variances and covariances to valid values.
    # Variance must be non-negative:
    sigma00 = np.maximum(0., sigma00)
    sigma11 = np.maximum(0., sigma11)
    sigma01 = np.sign(sigma01) * np.minimum(
        np.sqrt(sigma00 * sigma11), np.abs(sigma01))
    c1 = (k1 * max_val)**2
    c2 = (k2 * max_val)**2
    numer = (2 * mu01 + c1) * (2 * sigma01 + c2)
    denom = (mu00 + mu11 + c1) * (sigma00 + sigma11 + c2)
    ssim_map = numer / denom
    return np.mean(ssim_map)


def ssim_old(rgb, gts):
    """Calculate the SSIM metric.

    Assumes the RGB image is in [0,1]

    Args:
        rgb (torch.Tensor): Image tensor of shape [H,W,3]
        gts (torch.Tensor): Image tensor of shape [H,W,3]

    Returns:
        (float): The SSIM score
    """
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    return skimage.metrics.structural_similarity(
        rgb[..., :3].cpu().numpy(),
        gts[..., :3].cpu().numpy(),
        channel_axis=2,
        gaussian_weights=True,
        sigma=1.5,
        use_sample_covariance=False)


def msssim(rgb, gts):
    assert (rgb.max() <= 1.05 and rgb.min() >= -0.05)
    assert (gts.max() <= 1.05 and gts.min() >= -0.05)
    return ms_ssim(torch.permute(rgb[None, ...], (0, 3, 1, 2)),
                   torch.permute(gts[None, ...], (0, 3, 1, 2))).item()


__LPIPS__ = {}


def init_lpips(net_name, device):
    return lpips.LPIPS(net=net_name, version='0.1').eval().to(device)


def rgb_lpips(rgb, gts, net_name='alex', device='cpu'):
    if net_name not in __LPIPS__:
        __LPIPS__[net_name] = init_lpips(net_name, device)
    gts = gts.permute([2, 0, 1]).contiguous().to(device)
    rgb = rgb.permute([2, 0, 1]).contiguous().to(device)
    return __LPIPS__[net_name](gts, rgb, normalize=True).item()


def jod(pred_frames: List[np.ndarray], gt_frames: List[np.ndarray]) -> float:
    with tempfile.TemporaryDirectory() as tmpdir:
        file_pred = os.path.join(tmpdir, "pred.mp4")
        write_video_to_file(file_pred, pred_frames)
        file_gt = os.path.join(tmpdir, "gt.mp4")
        write_video_to_file(file_gt, gt_frames)
        result = subprocess.check_output([
            'fvvdp', '--test', file_pred, '--ref', file_gt, '--gpu', '0',
            '--display', 'standard_fhd'])
        result = float(result.decode().strip().split('=')[1])
    return result


def flip(pred_frames: List[np.ndarray], gt_frames: List[np.ndarray], interval: int = 10) -> float:
    def extract_from_result(text: str, prompt: str):
        m = re.search(prompt, text)
        return float(m.group(1))

    all_results = []
    with tempfile.TemporaryDirectory() as tmpdir:
        pred_fname = os.path.join(tmpdir, "pred.png")
        gt_fname = os.path.join(tmpdir, "gt.png")
        for i in range(len(pred_frames)):
            if (i % interval) != 0:
                continue
            write_png(pred_fname, pred_frames[i])
            write_png(gt_fname, gt_frames[i])
            result = subprocess.check_output(
                ['python', 'plenoxels/ops/flip/flip.py', '--reference', gt_fname, '--test', pred_fname]
            ).decode()
            all_results.append(extract_from_result(result, r'Mean: (\d+\.\d+)'))
    return sum(all_results) / len(all_results)

magnerf/ops/image/metrics.py
3d-motion-magnification/3d-motion-mag
import torch
from torch.nn import functional as F

from datetime import datetime

def grid_sample_wrapper(grid: torch.Tensor, coords: torch.Tensor, align_corners: bool = True) -> torch.Tensor:
    grid_dim = coords.shape[-1]

    if grid.dim() == grid_dim + 1:
        # no batch dimension present, need to add it
        grid = grid.unsqueeze(0)
    if coords.dim() == 2:
        coords = coords.unsqueeze(0)

    if grid_dim == 2 or grid_dim == 3:
        grid_sampler = F.grid_sample
    else:
        raise NotImplementedError(f"Grid-sample was called with {grid_dim}D data but is only "
                                  f"implemented for 2 and 3D data.")

    coords = coords.view([coords.shape[0]] + [1] * (grid_dim - 1) + list(coords.shape[1:]))
    
    now = datetime.now()

    if False and int(now.second) % 30 == 0:
        print(f'x: {coords[..., 0].min()}, {coords[..., 0].max()}')
        print(f'y: {coords[..., 1].min()}, {coords[..., 1].max()}')

    B, feature_dim = grid.shape[:2]
    n = coords.shape[-2]
    interp = grid_sampler(
        grid,  # [B, feature_dim, reso, ...]
        coords,  # [B, 1, ..., n, grid_dim]
        align_corners=align_corners,
        mode='bilinear', padding_mode='border')
    interp = interp.view(B, feature_dim, n).transpose(-1, -2)  # [B, n, feature_dim]
    interp = interp.squeeze()  # [B?, n, feature_dim?]
    return interp

magnerf/ops/interpolation.py
3d-motion-magnification/3d-motion-mag
import torch


def ray_samples_to_sdist(ray_samples):
    """Convert ray samples to s space"""
    starts = ray_samples.spacing_starts
    ends = ray_samples.spacing_ends
    sdist = torch.cat([starts[..., 0], ends[..., -1:, 0]], dim=-1)  # (num_rays, num_samples + 1)
    return sdist


def lossfun_distortion(t, w):
    """
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L142
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L266
    """
    ut = (t[..., 1:] + t[..., :-1]) / 2
    dut = torch.abs(ut[..., :, None] - ut[..., None, :])
    loss_inter = torch.sum(w * torch.sum(w[..., None, :] * dut, dim=-1), dim=-1)

    loss_intra = torch.sum(w**2 * (t[..., 1:] - t[..., :-1]), dim=-1) / 3

    return loss_inter + loss_intra


def distortion_loss(weights_list, ray_samples_list):
    """From mipnerf360"""
    # takes samples and weights as proposed by the 'nerf' network
    c = ray_samples_to_sdist(ray_samples_list[-1])
    w = weights_list[-1][..., 0]
    loss = torch.mean(lossfun_distortion(c, w))
    return loss

magnerf/ops/losses/distortion_loss.py
3d-motion-magnification/3d-motion-mag
"""
Implements the histogram loss from mipnerf360
"""
import torch


EPS = 1.0e-7


def outer(
    t0_starts: torch.Tensor,
    t0_ends: torch.Tensor,
    t1_starts: torch.Tensor,
    t1_ends: torch.Tensor,
    y1: torch.Tensor,
) -> torch.Tensor:
    """Faster version of
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L64
    Args:
        t0_starts: start of the interval edges
        t0_ends: end of the interval edges
        t1_starts: start of the interval edges
        t1_ends: end of the interval edges
        y1: weights
    """
    cy1 = torch.cat([torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1)

    idx_lo = torch.searchsorted(t1_starts.contiguous(), t0_starts.contiguous(), right=True) - 1
    idx_lo = torch.clamp(idx_lo, min=0, max=y1.shape[-1] - 1)
    idx_hi = torch.searchsorted(t1_ends.contiguous(), t0_ends.contiguous(), right=True)
    idx_hi = torch.clamp(idx_hi, min=0, max=y1.shape[-1] - 1)
    cy1_lo = torch.take_along_dim(cy1[..., :-1], idx_lo, dim=-1)
    cy1_hi = torch.take_along_dim(cy1[..., 1:], idx_hi, dim=-1)
    y0_outer = cy1_hi - cy1_lo

    return y0_outer


def lossfun_outer(
    t: torch.Tensor,  # [..., "num_samples+1"],
    w: torch.Tensor,  # [..., "num_samples"],
    t_env: torch.Tensor,  # [..., "num_samples+1"],
    w_env: torch.Tensor,  # [..., "num_samples"],
):
    """
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L136
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L80
    Args:
        t: interval edges
        w: weights
        t_env: interval edges of the upper bound enveloping histogram (from proposal model)
        w_env: weights that should upper bound the inner (t,w) histogram (from proposal model)
    """
    w_outer = outer(t[..., :-1], t[..., 1:], t_env[..., :-1], t_env[..., 1:], w_env)
    return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS)


def ray_samples_to_sdist(ray_samples):
    """Convert ray samples to s space"""
    starts = ray_samples.spacing_starts
    ends = ray_samples.spacing_ends
    sdist = torch.cat([starts[..., 0], ends[..., -1:, 0]], dim=-1)  # (num_rays, num_samples + 1)
    return sdist


def interlevel_loss(weights_list, ray_samples_list):
    """Calculates the proposal loss in the MipNeRF-360 paper.
    https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/model.py#L515
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/train_utils.py#L133
    """
    c = ray_samples_to_sdist(ray_samples_list[-1]).detach()
    w = weights_list[-1][..., 0].detach()
    loss_interlevel = 0.0
    for ray_samples, weights in zip(ray_samples_list[:-1], weights_list[:-1]):
        sdist = ray_samples_to_sdist(ray_samples)
        cp = sdist  # (num_rays, num_samples + 1)
        wp = weights[..., 0]  # (num_rays, num_samples)
        loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp))
    return loss_interlevel

magnerf/ops/losses/histogram_loss.py
3d-motion-magnification/3d-motion-mag
import math

import torch
import torch.optim


def get_cosine_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    eta_min: float = 0.0,
    num_cycles: float = 0.999,
    last_epoch: int = -1,
):
    """
    https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/optimization.py#L129
    """
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(eta_min, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)


def get_log_linear_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    eta_min: float = 2e-5,
    eta_max: float = 1e-2,
    num_cycles: float = 0.999,
    last_epoch: int = -1,
):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))  # in [0,1]
        return math.exp(progress * math.log(eta_min) + (1 - progress) * math.log(eta_max)) / eta_max

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)


def get_step_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    milestones,
    gamma: float,
    num_warmup_steps: int,
    last_epoch: int = -1,
):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        out = 1.0
        for m in milestones:
            if current_step < m:
                break
            out *= gamma
        return out
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch)

magnerf/ops/lr_scheduling.py
3d-motion-magnification/3d-motion-mag
"""The ray samplers are almost completely copied from NeRF-studio

https://github.com/nerfstudio-project/nerfstudio/blob/628e4fe1a638e7fb3b7ad33d4d91a4b1d63a9b68/nerfstudio/model_components/ray_samplers.py

Copyright 2022 The Nerfstudio Team. All rights reserved.

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

    http://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.
"""

from abc import abstractmethod
from dataclasses import dataclass
from typing import Optional, Callable, Tuple, List

import torch
from torch import nn


@dataclass
class RaySamples:
    """xyz coordinate for ray origin."""
    origins: torch.Tensor  # [bs:..., 3]
    """Direction of ray."""
    directions: torch.Tensor  # [bs:..., 3]
    """Where the frustum starts along a ray."""
    starts: torch.Tensor  # [bs:..., 1]
    """Where the frustum ends along a ray."""
    ends: torch.Tensor  # [bs:..., 1]
    """"width" of each sample."""
    deltas: Optional[torch.Tensor] = None  # [bs, ...?, 1]
    """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling."""
    spacing_starts: Optional[torch.Tensor] = None  # [bs, ...?, num_samples, 1]
    """End of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling."""
    spacing_ends: Optional[torch.Tensor] = None  # [bs, ...?, num_samples, 1]
    """Function to convert bins to euclidean distance."""
    spacing_to_euclidean_fn: Optional[Callable] = None

    def get_positions(self) -> torch.Tensor:
        """Calulates "center" position of frustum. Not weighted by mass.
        Returns:
            xyz positions (..., 3).
        """
        return self.origins + self.directions * (self.starts + self.ends) / 2  # world space

    def get_weights2(self, densities: torch.Tensor) -> torch.Tensor:
        densities = densities.squeeze(2)
        deltas = self.deltas.squeeze(2)
        delta_mask = deltas > 0
        deltas = deltas[delta_mask]

        delta_density = torch.zeros_like(densities)
        delta_density[delta_mask] = deltas * densities[delta_mask]
        alphas = 1 - torch.exp(-delta_density)

        transmittance = torch.cat(
            (
                torch.ones(alphas.shape[0], 1, device=alphas.device),
                torch.cumprod(1.0 - alphas, dim=-1)
            ), dim=-1
        )
        weights = alphas * transmittance[:, :-1]
        return weights[..., None]

    def get_weights(self, densities: torch.Tensor, opaque_background=False) -> torch.Tensor:
        """Return weights based on predicted densities
        Args:
            densities: Predicted densities for samples along ray (..., num_samples, 1)
        Returns:
            Weights for each sample  (..., num_samples, 1)
        """
        delta_mask = self.deltas > 0
        deltas = self.deltas[delta_mask]

        delta_density = torch.zeros_like(densities)
        delta_density[delta_mask] = deltas * densities[delta_mask]

        if opaque_background:
            # Equivalent to making the final t-interval infinitely wide.
            delta_density = torch.cat(
                [
                    delta_density[..., :-1, :],
                    torch.ones_like(delta_density[..., -1:, :]) * 1e18
                ],dim=-2)

        alphas = 1 - torch.exp(-delta_density)

        transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)
        transmittance = torch.cat(
            [torch.zeros((*transmittance.shape[:1], 1, 1), device=densities.device), transmittance], dim=-2
        )
        transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
        weights = alphas * transmittance  # [..., "num_samples"]
        return weights


@dataclass
class RayBundle:
    """A bundle of ray parameters."""

    """Ray origins (XYZ)"""
    origins: torch.Tensor  # [..., 3]
    """Unit ray direction vector"""
    directions: torch.Tensor  # [..., 3]
    """Distance along ray to start sampling"""
    nears: Optional[torch.Tensor] = None  # [..., 1]
    """Rays Distance along ray to stop sampling"""
    fars: Optional[torch.Tensor] = None  # [..., 1]

    def __len__(self):
        num_rays = torch.numel(self.origins) // self.origins.shape[-1]
        return num_rays

    def get_ray_samples(
        self,
        bin_starts: torch.Tensor,
        bin_ends: torch.Tensor,
        spacing_starts: Optional[torch.Tensor] = None,
        spacing_ends: Optional[torch.Tensor] = None,
        spacing_to_euclidean_fn: Optional[Callable] = None,
    ) -> RaySamples:
        """Produces samples for each ray by projection points along the ray direction. Currently samples uniformly.
        Args:
            bin_starts: Distance from origin to start of bin.
                TensorType["bs":..., "num_samples", 1]
            bin_ends: Distance from origin to end of bin.
        Returns:
            Samples projected along ray.
        """
        deltas = bin_ends - bin_starts
        return RaySamples(
            origins=self.origins[..., None, :],  # [..., 1, 3]
            directions=self.directions[..., None, :],  # [..., 1, 3]
            starts=bin_starts,  # [..., num_samples, 1]  world
            ends=bin_ends,  # [..., num_samples, 1]      world
            deltas=deltas,  # [..., num_samples, 1]  world coo
            spacing_starts=spacing_starts,  # [..., num_samples, 1]
            spacing_ends=spacing_ends,  # [..., num_samples, 1]
            spacing_to_euclidean_fn=spacing_to_euclidean_fn,
        )


class Sampler(nn.Module):
    """Generate Samples
    Args:
        num_samples: number of samples to take
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.num_samples = num_samples

    @abstractmethod
    def generate_ray_samples(self) -> RaySamples:
        """Generate Ray Samples"""

    def forward(self, *args, **kwargs) -> RaySamples:
        """Generate ray samples"""
        return self.generate_ray_samples(*args, **kwargs)


class SpacedSampler(Sampler):
    """Sample points according to a function.
    Args:
        num_samples: Number of samples per ray
        spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
        spacing_fn_inv: The inverse of spacing_fn.
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        spacing_fn: Callable,
        spacing_fn_inv: Callable,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(num_samples=num_samples)
        self.train_stratified = train_stratified
        self.single_jitter = single_jitter
        self.spacing_fn = spacing_fn
        self.spacing_fn_inv = spacing_fn_inv

    # noinspection PyMethodOverriding
    def generate_ray_samples(
        self,
        ray_bundle: RayBundle,
        num_samples: Optional[int] = None,
    ) -> RaySamples:
        """Generates position samples accoring to spacing function.
        Args:
            ray_bundle: Ray-origins, directions, etc.
            num_samples: Number of samples per ray
        Returns:
            Positions and deltas for samples along a ray
        """
        num_samples = num_samples or self.num_samples
        assert num_samples is not None
        num_rays = ray_bundle.origins.shape[0]

        bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...]  # [1, num_samples+1]

        # TODO More complicated than it needs to be.
        if self.train_stratified and self.training:
            if self.single_jitter:
                t_rand = torch.rand((num_rays, 1), dtype=bins.dtype, device=bins.device)
            else:
                t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)
            bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0
            bin_upper = torch.cat([bin_centers, bins[..., -1:]], -1)
            bin_lower = torch.cat([bins[..., :1], bin_centers], -1)
            bins = bin_lower + (bin_upper - bin_lower) * t_rand
        else:
            bins = bins.repeat(num_rays, 1)

        # s_near, s_far in [0, 1]
        s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars))
        spacing_to_euclidean_fn = lambda x: self.spacing_fn_inv(x * s_far + (1 - x) * s_near)
        # euclidean = world
        euclidean_bins = spacing_to_euclidean_fn(bins)  # [num_rays, num_samples+1]

        return ray_bundle.get_ray_samples(
            bin_starts=euclidean_bins[..., :-1, None],  # world [near, far]
            bin_ends=euclidean_bins[..., 1:, None],     # world [near, far]
            spacing_starts=bins[..., :-1, None],        # [0, 1]
            spacing_ends=bins[..., 1:, None],           # [0, 1]
            spacing_to_euclidean_fn=spacing_to_euclidean_fn,
        )


class UniformSampler(SpacedSampler):
    """Sample uniformly along a ray
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: x,
            spacing_fn_inv=lambda x: x,
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class LinearDisparitySampler(SpacedSampler):
    """Sample linearly in disparity along a ray
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: 1 / x,
            spacing_fn_inv=lambda x: 1 / x,
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class UniformLinDispPiecewiseSampler(SpacedSampler):
    """Piecewise sampler along a ray that allocates the first half of the samples uniformly and the second half
    using linearly in disparity spacing.
    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified=True,
        single_jitter=False,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: torch.where(x < 1, x / 2, 1 - 1 / (2 * x)),
            spacing_fn_inv=lambda x: torch.where(x < 0.5, 2 * x, 1 / (2 - 2 * x)),
            train_stratified=train_stratified,
            single_jitter=single_jitter,
        )


class PDFSampler(Sampler):
    """Sample based on probability distribution
    Args:
        num_samples: Number of samples per ray
        train_stratified: Randomize location within each bin during training.
        single_jitter: Use a same random jitter for all samples along a ray. Defaults to False
        include_original: Add original samples to ray.
        histogram_padding: Amount to weights prior to computing PDF.
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        train_stratified: bool = True,
        single_jitter: bool = False,
        include_original: bool = True,
        histogram_padding: float = 0.01,
    ) -> None:
        super().__init__(num_samples=num_samples)
        self.train_stratified = train_stratified
        self.include_original = include_original
        self.histogram_padding = histogram_padding
        self.single_jitter = single_jitter

    # noinspection PyMethodOverriding
    def generate_ray_samples(
        self,
        ray_bundle: RayBundle,
        ray_samples: Optional[RaySamples] = None,
        weights: Optional[torch.Tensor] = None,
        num_samples: Optional[int] = None,
        eps: float = 1e-5,
    ) -> RaySamples:
        """Generates position samples given a distribution.
        Args:
            ray_bundle: Ray-origins, directions, etc.
            ray_samples: Existing ray samples
            weights: Weights for each bin  [..., "num_samples", 1]
            num_samples: Number of samples per ray
            eps: Small value to prevent numerical issues.
        Returns:
            Positions and deltas for samples along a ray
        """
        if ray_samples is None or ray_bundle is None:
            raise ValueError("ray_samples must be provided")

        num_samples = num_samples or self.num_samples
        assert num_samples is not None
        num_bins = num_samples + 1

        weights = weights[..., 0] + self.histogram_padding

        # Add small offset to rays with zero weight to prevent NaNs
        weights_sum = torch.sum(weights, dim=-1, keepdim=True)
        padding = torch.relu(eps - weights_sum)
        weights = weights + padding / weights.shape[-1]
        weights_sum += padding

        pdf = weights / weights_sum
        cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1))
        cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)

        if self.train_stratified and self.training:
            # Stratified samples between 0 and 1
            u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
            u = u.expand((*cdf.shape[:-1], num_bins))
            if self.single_jitter:
                rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins
            else:
                rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins
            u = u + rand
        else:
            # Uniform samples between 0 and 1
            u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
            u = u + 1.0 / (2 * num_bins)
            u = u.expand(size=(*cdf.shape[:-1], num_bins))
        u = u.contiguous()

        assert (
            ray_samples.spacing_starts is not None and ray_samples.spacing_ends is not None
        ), "ray_sample spacing_starts and spacing_ends must be provided"
        assert ray_samples.spacing_to_euclidean_fn is not None, "ray_samples.spacing_to_euclidean_fn must be provided"
        existing_bins = torch.cat(
            [
                ray_samples.spacing_starts[..., 0],
                ray_samples.spacing_ends[..., -1:, 0],
            ],
            dim=-1,
        )  # [0, 1]

        inds = torch.searchsorted(cdf, u, right=True)
        below = torch.clamp(inds - 1, 0, existing_bins.shape[-1] - 1)
        above = torch.clamp(inds, 0, existing_bins.shape[-1] - 1)
        cdf_g0 = torch.gather(cdf, -1, below)
        bins_g0 = torch.gather(existing_bins, -1, below)
        cdf_g1 = torch.gather(cdf, -1, above)
        bins_g1 = torch.gather(existing_bins, -1, above)

        t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1)
        bins = bins_g0 + t * (bins_g1 - bins_g0)

        if self.include_original:
            bins, _ = torch.sort(torch.cat([existing_bins, bins], -1), -1)

        # Stop gradients
        bins = bins.detach()

        euclidean_bins = ray_samples.spacing_to_euclidean_fn(bins)

        return ray_bundle.get_ray_samples(
            bin_starts=euclidean_bins[..., :-1, None],
            bin_ends=euclidean_bins[..., 1:, None],
            spacing_starts=bins[..., :-1, None],
            spacing_ends=bins[..., 1:, None],
            spacing_to_euclidean_fn=ray_samples.spacing_to_euclidean_fn,
        )


class ProposalNetworkSampler(Sampler):
    """Sampler that uses a proposal network to generate samples."""
    def __init__(
        self,
        num_proposal_samples_per_ray: Tuple[int] = (64,),
        num_nerf_samples_per_ray: int = 32,
        num_proposal_network_iterations: int = 2,
        single_jitter: bool = False,
        update_sched: Callable = lambda x: 1,
        initial_sampler: Optional[Sampler] = None,
    ) -> None:
        super().__init__()
        self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
        self.num_nerf_samples_per_ray = num_nerf_samples_per_ray
        self.num_proposal_network_iterations = num_proposal_network_iterations
        self.update_sched = update_sched
        if self.num_proposal_network_iterations < 1:
            raise ValueError("num_proposal_network_iterations must be >= 1")

        # samplers
        if initial_sampler is None:
            initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
        self.initial_sampler = initial_sampler
        self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)

        self._anneal = 1.0
        self._steps_since_update = 0
        self._step = 0

    def set_anneal(self, anneal: float) -> None:
        """Set the anneal value for the proposal network."""
        self._anneal = anneal

    def step_cb(self, step):
        """Callback to register a training step has passed. This is used to keep track of the sampling schedule"""
        self._step = step
        self._steps_since_update += 1

    def generate_ray_samples(
        self,
        ray_bundle: Optional[RayBundle] = None,
        timestamps: Optional[float] = None,
        density_fns: Optional[List[Callable]] = None,
    ) -> Tuple[RaySamples, List, List]:
        assert ray_bundle is not None
        assert density_fns is not None
        assert len(density_fns) == self.num_proposal_network_iterations

        weights_list = []
        ray_samples_list = []

        n = self.num_proposal_network_iterations
        weights = None
        ray_samples = None
        updated = self._steps_since_update > self.update_sched(self._step) or self._step < 10
        for i_level in range(n + 1):
            is_prop = i_level < n
            num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray
            if i_level == 0:
                # Uniform sampling because we need to start with some samples
                ray_samples = self.initial_sampler(ray_bundle, num_samples=num_samples)
            else:
                # PDF sampling based on the last samples and their weights
                # Perform annealing to the weights. This will be a no-op if self._anneal is 1.0.
                assert weights is not None
                annealed_weights = torch.pow(weights, self._anneal)
                ray_samples = self.pdf_sampler(ray_bundle, ray_samples, annealed_weights, num_samples=num_samples)
            if is_prop:
                if updated:
                    # always update on the first step or the inf check in grad scaling crashes
                    density = density_fns[i_level](ray_samples.get_positions(), timestamps)  # world space
                else:
                    with torch.no_grad():
                        density = density_fns[i_level](ray_samples.get_positions(), timestamps)
                if type(density) is tuple:
                    density = density[0]
                weights = ray_samples.get_weights(density)
                weights_list.append(weights)  # (num_rays, num_samples)
                ray_samples_list.append(ray_samples)
        if updated:
            self._steps_since_update = 0

        assert ray_samples is not None
        return ray_samples, weights_list, ray_samples_list

    def __str__(self):
        return (f"ProposalNetworkSampler("
                f"num_proposal_samples_per_ray={self.num_proposal_samples_per_ray}, "
                f"num_nerf_samples_per_ray={self.num_nerf_samples_per_ray}, "
                f"num_proposal_network_iterations={self.num_proposal_network_iterations})")

magnerf/raymarching/ray_samplers.py
3d-motion-magnification/3d-motion-mag
from typing import Optional, Union

import torch
import torch.nn as nn


class SpatialDistortion(nn.Module):
    """Apply spatial distortions"""

    def forward(
        self, positions: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            positions: Sample to distort (shape: batch-size, ..., 3)
        Returns:
            distorted sample - same shape
        """

import numpy as np
import struct
def write_pointcloud(filename,xyz_points,rgb_points=None):
    assert xyz_points.shape[1] == 3,'Input XYZ points should be Nx3 float array'
    if rgb_points is None:
        rgb_points = np.ones(xyz_points.shape).astype(np.uint8)*255
    assert xyz_points.shape == rgb_points.shape,'Input RGB colors should be Nx3 float array and have same size as input XYZ points'

    # Write header of .ply file
    fid = open(filename,'wb')
    fid.write(bytes('ply\n', 'utf-8'))
    fid.write(bytes('format binary_little_endian 1.0\n', 'utf-8'))
    fid.write(bytes('element vertex %d\n'%xyz_points.shape[0], 'utf-8'))
    fid.write(bytes('property float x\n', 'utf-8'))
    fid.write(bytes('property float y\n', 'utf-8'))
    fid.write(bytes('property float z\n', 'utf-8'))
    fid.write(bytes('property uchar red\n', 'utf-8'))
    fid.write(bytes('property uchar green\n', 'utf-8'))
    fid.write(bytes('property uchar blue\n', 'utf-8'))
    fid.write(bytes('end_header\n', 'utf-8'))

    # Write 3D points to .ply file
    for i in range(xyz_points.shape[0]):
        fid.write(bytearray(struct.pack("fffccc",xyz_points[i,0],xyz_points[i,1],xyz_points[i,2],
                                        rgb_points[i,0].tobytes(),rgb_points[i,1].tobytes(),
                                        rgb_points[i,2].tobytes())))
    fid.close()


class SceneContraction(SpatialDistortion):
    """Contract unbounded space using the contraction was proposed in MipNeRF-360.
        We use the following contraction equation:
        .. math::
            f(x) = \\begin{cases}
                x & ||x|| \\leq 1 \\\\
                (2 - \\frac{1}{||x||})(\\frac{x}{||x||}) & ||x|| > 1
            \\end{cases}
        If the order is not specified, we use the Frobenius norm, this will contract the space to a sphere of
        radius 1. If the order is L_inf (order=float("inf")), we will contract the space to a cube of side length 2.
        If using voxel based encodings such as the Hash encoder, we recommend using the L_inf norm.
        Args:
            order: Order of the norm. Default to the Frobenius norm. Must be set to None for Gaussians.
    """

    def __init__(self,
                 order: Optional[Union[float, int]] = None,
                 global_translation: Optional[torch.Tensor] = None,
                 global_scale: Optional[torch.Tensor] = None,
                 ) -> None:
        super().__init__()
        self.order = order

        if global_translation is None:
            global_translation = torch.tensor([0.0, 0.0, 0.0])
        self.global_translation = nn.Parameter(global_translation, requires_grad=False)
        if global_scale is None:
            global_scale = torch.tensor([1.0, 1.0, 1.0])

        #self.global_scale = nn.Parameter(global_scale, requires_grad=False)

        #aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0]) * 4.0
        #aabb = global_scale
        self.aabb = nn.Parameter(torch.tensor(global_scale), requires_grad=False)
        

    def forward(self, positions):
        # Apply global scale and translation
        #positions = (
        #    positions * self.global_scale[None, None, :]
        #    + self.global_translation[None, None, :]
        #)

        aabb_min, aabb_max = torch.split(self.aabb, 3, dim=-1)
        positions = (positions - aabb_min) / (aabb_max - aabb_min)
        positions = positions * 2 - 1  # aabb is at [-1, 1]

        mag = torch.linalg.norm(positions, ord=self.order, dim=-1)
        mask = mag >= 1
        x_new = positions.clone()
        #x_new[mask] = (2 - (1 / mag[mask][..., None])) * (positions[mask] / mag[mask][..., None])
        x_new[mask] = (1.5 - (1 / mag[mask][..., None])) * (positions[mask] / mag[mask][..., None])

        x_new = (x_new / 1.5) * 2

        return x_new

    def __str__(self):
        return (f"SceneContraction(global_translation={self.global_translation}, "
                f"global_scale={self.global_scale})")

magnerf/raymarching/spatial_distortions.py
3d-motion-magnification/3d-motion-mag
import abc
import random
import logging as log
import math
import os
from copy import copy
from typing import Iterable, Optional, Union, Dict, Tuple, Sequence, MutableMapping

import numpy as np
import torch
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter

from magnerf.utils.timer import CudaTimer
from magnerf.utils.ema import EMA
from magnerf.models.lowrank_model import LowrankModel
from magnerf.utils.my_tqdm import tqdm
from magnerf.ops.image import metrics
from magnerf.ops.image.io import write_png
from magnerf.runners.regularization import Regularizer
from magnerf.ops.lr_scheduling import (
    get_cosine_schedule_with_warmup, get_step_schedule_with_warmup
)


class BaseTrainer(abc.ABC):
    def __init__(self,
                 train_data_loader: Iterable,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 device: Union[str, torch.device],
                 **kwargs):
        self.train_data_loader = train_data_loader
        self.num_steps = num_steps
        self.train_fp16 = train_fp16
        self.save_every = save_every
        self.valid_every = int(valid_every)
        self.save_outputs = save_outputs
        self.device = device
        self.eval_batch_size = kwargs.get('eval_batch_size', 8192)
        self.extra_args = kwargs
        self.timer = CudaTimer(enabled=False)

        self.log_dir = os.path.join(logdir, expname)
        os.makedirs(self.log_dir, exist_ok=True)
        #self.writer = SummaryWriter(log_dir=self.log_dir)

        self.global_step: Optional[int] = None
        self.loss_info: Optional[Dict[str, EMA]] = None

        self.model = self.init_model(**self.extra_args)
        self.optimizer = self.init_optim(**self.extra_args)
        self.scheduler = self.init_lr_scheduler(**self.extra_args)
        self.criterion = torch.nn.MSELoss(reduction='mean')
        self.regularizers = self.init_regularizers(**self.extra_args)
        self.gscaler = torch.cuda.amp.GradScaler(enabled=self.train_fp16)

        self.model.to(self.device)

    @abc.abstractmethod
    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        self.model.eval()
        return None  # noqa

    def train_step(self, data, **kwargs) -> bool:
        self.model.train()
        data = self._move_data_to_device(data)
        if "timestamps" not in data:
            data["timestamps"] = None
        self.timer.check("move-to-device")

        with torch.cuda.amp.autocast(enabled=self.train_fp16):
            fwd_out = self.model(
                data['rays_o'], data['rays_d'], bg_color=data['bg_color'],
                near_far=data['near_fars'], timestamps=data['timestamps'],
                c2w=data['c2w'] if 'c2w' in data else None,
                camera_dirs=data['camera_dirs'] if 'camera_dirs' in data else None,
                indices=data['indices'] if 'indices' in data else None,)

            self.timer.check("model-forward")
            # Reconstruction loss
            recon_loss = self.criterion(fwd_out['rgb'], data['imgs'])
            # Regularization
            loss = recon_loss

            if "bg_mask" in data:
                bg_acc = fwd_out['accumulation'] * data['bg_mask']
                bg_loss = torch.nn.functional.l1_loss(bg_acc, torch.zeros_like(bg_acc))
                loss = loss + 1e-4 * bg_loss

            for r in self.regularizers:
                reg_loss = r.regularize(self.model, model_out=fwd_out)
                loss = loss + reg_loss
            self.timer.check("regularizaion-forward")
        # Update weights
        self.optimizer.zero_grad(set_to_none=True)
        self.gscaler.scale(loss).backward()
        self.timer.check("backward")
        self.gscaler.step(self.optimizer)
        scale = self.gscaler.get_scale()
        self.gscaler.update()
        self.timer.check("scaler-step")

        # Report on losses
        if self.global_step % self.calc_metrics_every == 0:
            with torch.no_grad():
                recon_loss_val = recon_loss.item()
                self.loss_info[f"mse"].update(recon_loss_val)
                self.loss_info[f"psnr"].update(-10 * math.log10(recon_loss_val))
                self.loss_info[f"lr"].update(self.lr)
                for r in self.regularizers:
                    r.report(self.loss_info)

        return scale <= self.gscaler.get_scale()

    def post_step(self, progress_bar):
        self.model.step_after_iter(self.global_step)
        if self.global_step % self.calc_metrics_every == 0:
            progress_bar.set_postfix_str(
                losses_to_postfix(self.loss_info, lr=self.lr), refresh=False)
            for loss_name, loss_val in self.loss_info.items():
                #self.writer.add_scalar(f"train/loss/{loss_name}", loss_val.value, self.global_step)
                if self.timer.enabled:
                    tsum = 0.
                    tstr = "Timings: "
                    for tname, tval in self.timer.timings.items():
                        tstr += f"{tname}={tval:.1f}ms  "
                        tsum += tval
                    tstr += f"tot={tsum:.1f}ms"
                    log.info(tstr)
        progress_bar.update(1)
        if self.valid_every > -1 and self.global_step % self.valid_every == 0:
            print()
            self.validate()
        if self.save_every > -1 and self.global_step % self.save_every == 0:
            print()
            self.save_model()

    def pre_epoch(self):
        self.loss_info = self.init_epoch_info()

    def train(self):
        """Override this if some very specific training procedure is needed."""
        if self.global_step is None:
            self.global_step = 0
        log.info(f"Starting training from step {self.global_step + 1}")
        pb = tqdm(initial=self.global_step, total=self.num_steps)
        try:
            self.pre_epoch()
            batch_iter = iter(self.train_data_loader)
            while self.global_step < self.num_steps:
                self.timer.reset()
                self.model.step_before_iter(self.global_step)
                self.global_step += 1
                self.timer.check("step-before-iter")
                try:
                    data = next(batch_iter)
                    self.timer.check("dloader-next")
                except StopIteration:
                    self.pre_epoch()
                    batch_iter = iter(self.train_data_loader)
                    data = next(batch_iter)
                    log.info("Reset data-iterator")

                try:
                    step_successful = self.train_step(data)
                except StopIteration:
                    self.pre_epoch()
                    batch_iter = iter(self.train_data_loader)
                    log.info("Reset data-iterator")
                    step_successful = True

                if step_successful and self.scheduler is not None:
                    self.scheduler.step()
                for r in self.regularizers:
                    r.step(self.global_step)
                self.post_step(progress_bar=pb)
                self.timer.check("after-step")
        finally:
            pb.close()
            #self.writer.close()

    def _move_data_to_device(self, data):
        data["rays_o"] = data["rays_o"].to(self.device)
        data["rays_d"] = data["rays_d"].to(self.device)
        data["imgs"] = data["imgs"].to(self.device)
        data["near_fars"] = data["near_fars"].to(self.device)
        if "timestamps" in data:
            data["timestamps"] = data["timestamps"].to(self.device)
        if "bg_mask" in data:
            data["bg_mask"] = data["bg_mask"].to(self.device)
        bg_color = data["bg_color"]
        if isinstance(bg_color, torch.Tensor):
            bg_color = bg_color.to(self.device)
        data["bg_color"] = bg_color
        return data

    def _normalize_err(self, preds: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
        err = torch.abs(preds - gt)
        err = err.mean(-1, keepdim=True)  # mean over channels
        # normalize between 0, 1 where 1 corresponds to the 90th percentile
        # err = err.clamp_max(torch.quantile(err, 0.9))
        err = self._normalize_01(err)
        return err.repeat(1, 1, 3)

    @staticmethod
    def _normalize_01(t: torch.Tensor) -> torch.Tensor:
        return (t - t.min()) / t.max()

    def _normalize_depth(self, depth: torch.Tensor, img_h: int, img_w: int) -> torch.Tensor:
        return (
            self._normalize_01(depth)
        ).cpu().reshape(img_h, img_w)[..., None]

    def calc_metrics(self, preds: torch.Tensor, gt: torch.Tensor):
        if gt.shape[-1] == 4:
            gt = gt[..., :3] * gt[..., 3:] + (1.0 - gt[..., 3:])

        err = (gt - preds) ** 2
        return {
            "mse": torch.mean(err),
            "psnr": metrics.psnr(preds, gt),
            "ssim": metrics.ssim(preds, gt),
            "ms-ssim": metrics.msssim(preds, gt),
            #"alex_lpips": metrics.rgb_lpips(preds, gt, net_name='alex', device=err.device),
            #"vgg_lpips": metrics.rgb_lpips(preds, gt, net_name='vgg', device=err.device)
        }

    def evaluate_metrics(self,
                         gt: Optional[torch.Tensor],
                         preds: MutableMapping[str, torch.Tensor],
                         dset,
                         img_idx: int,
                         name: Optional[str] = None,
                         save_outputs: bool = True) -> Tuple[dict, np.ndarray, Optional[np.ndarray]]:
        if isinstance(dset.img_h, int):
            img_h, img_w = dset.img_h, dset.img_w
        else:
            img_h, img_w = dset.img_h[img_idx], dset.img_w[img_idx]
        preds_rgb = (
            preds["rgb"]
            .reshape(img_h, img_w, 3)
            .cpu()
            .clamp(0, 1)
        )
        if not torch.isfinite(preds_rgb).all():
            log.warning(f"Predictions have {torch.isnan(preds_rgb).sum()} NaNs, "
                        f"{torch.isinf(preds_rgb).sum()} infs.")
            preds_rgb = torch.nan_to_num(preds_rgb, nan=0.0)
        out_img = preds_rgb
        summary = dict()

        out_depth = None
        if "depth" in preds:
            out_depth = preds["depth"].cpu().reshape(img_h, img_w)[..., None]
        preds.pop("depth")

        # This is used for proposal-depth keys
        for k in preds.keys():
            if "depth" in k:
                prop_depth = preds[k].cpu().reshape(img_h, img_w)[..., None]
                out_depth = torch.cat((out_depth, prop_depth), dim=1) if out_depth is not None else prop_depth

        if gt is not None:
            gt = gt.reshape(img_h, img_w, -1).cpu()
            if gt.shape[-1] == 4:
                gt = gt[..., :3] * gt[..., 3:] + (1.0 - gt[..., 3:])
            summary.update(self.calc_metrics(preds_rgb, gt))
            out_img = torch.cat((out_img, gt), dim=1)
            out_img = torch.cat((out_img, self._normalize_err(preds_rgb, gt)), dim=1)

        out_img_np: np.ndarray = (out_img * 255.0).byte().numpy()
        out_depth_np: Optional[np.ndarray] = None
        if out_depth is not None:
            out_depth = self._normalize_01(out_depth)
            out_depth_np = (out_depth * 255.0).repeat(1, 1, 3).byte().numpy()

        if save_outputs:
            out_name = f"step{self.global_step}-{img_idx}"
            if name is not None and name != "":
                out_name += "-" + name
            write_png(os.path.join(self.log_dir, out_name + ".png"), out_img_np)
            if out_depth is not None:
                depth_name = out_name + "-depth"
                write_png(os.path.join(self.log_dir, depth_name + ".png"), out_depth_np)

        return summary, out_img_np, out_depth_np

    @abc.abstractmethod
    def validate(self):
        pass

    def report_test_metrics(self, scene_metrics: Dict[str, Sequence[float]], extra_name: Optional[str]):
        log_text = f"step {self.global_step}/{self.num_steps}"
        if extra_name is not None:
            log_text += f" | {extra_name}"
        scene_metrics_agg: Dict[str, float] = {}
        for k in scene_metrics:
            ak = f"{k}_{extra_name}"
            scene_metrics_agg[ak] = np.mean(np.asarray(scene_metrics[k])).item()
            log_text += f" | {k}: {scene_metrics_agg[ak]:.4f}"
            #self.writer.add_scalar(f"test/{ak}", scene_metrics_agg[ak], self.global_step)

        log.info(log_text)
        return scene_metrics_agg

    def get_save_dict(self):
        return {
            "model": self.model.state_dict(),
            #"optimizer": self.optimizer.state_dict(),
            #"lr_scheduler": self.scheduler.state_dict() if self.scheduler is not None else None,
            "global_step": self.global_step
        }

    def save_model(self):
        model_fname = os.path.join(self.log_dir, f'model.pth')
        log.info(f'Saving model checkpoint to: {model_fname}')
        torch.save(self.get_save_dict(), model_fname)

    def load_model(self, checkpoint_data, training_needed: bool = True):
        self.model.load_state_dict(checkpoint_data["model"], strict=False)
        log.info("=> Loaded model state from checkpoint")

        #if training_needed:
        #    self.optimizer.load_state_dict(checkpoint_data["optimizer"])
        #    log.info("=> Loaded optimizer state from checkpoint")

        #if training_needed and self.scheduler is not None:
        #    self.scheduler.load_state_dict(checkpoint_data['lr_scheduler'])
        #    log.info("=> Loaded scheduler state from checkpoint")

        self.global_step = checkpoint_data["global_step"]
        log.info(f"=> Loaded step {self.global_step} from checkpoints")

    @abc.abstractmethod
    def init_epoch_info(self) -> Dict[str, EMA]:
        pass

    # noinspection PyUnresolvedReferences,PyProtectedMember
    def init_lr_scheduler(self, **kwargs) -> Optional[torch.optim.lr_scheduler._LRScheduler]:
        eta_min = 0
        lr_sched = None
        max_steps = self.num_steps
        scheduler_type = kwargs['scheduler_type']
        log.info(f"Initializing LR Scheduler of type {scheduler_type} with "
                 f"{max_steps} maximum steps.")
        if scheduler_type == "cosine":
            lr_sched = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer,
                T_max=max_steps,
                eta_min=eta_min)
        elif scheduler_type == "warmup_cosine":
            lr_sched = get_cosine_schedule_with_warmup(
                self.optimizer, num_warmup_steps=512, num_training_steps=max_steps)
        elif scheduler_type == "step":
            lr_sched = torch.optim.lr_scheduler.MultiStepLR(
                self.optimizer,
                milestones=[
                    max_steps // 2,
                    max_steps * 3 // 4,
                    max_steps * 5 // 6,
                    max_steps * 9 // 10,
                ],
                gamma=0.33)
        elif scheduler_type == "warmup_step":
            lr_sched = get_step_schedule_with_warmup(
                self.optimizer, milestones=[
                    max_steps // 2,
                    max_steps * 3 // 4,
                    max_steps * 5 // 6,
                    max_steps * 9 // 10,
                ],
                gamma=0.33,
                num_warmup_steps=512)
        return lr_sched

    def init_optim(self, **kwargs) -> torch.optim.Optimizer:
        optim_type = kwargs['optim_type']
        if optim_type == 'adam':
            optim = torch.optim.Adam(params=self.model.get_params(kwargs['lr']), eps=1e-15)
        else:
            raise NotImplementedError()
        return optim

    @abc.abstractmethod
    def init_model(self, **kwargs) -> torch.nn.Module:
        pass

    def get_regularizers(self, **kwargs) -> Sequence[Regularizer]:
        return ()

    def init_regularizers(self, **kwargs):
        # Keep only the regularizers with a positive weight
        regularizers = [r for r in self.get_regularizers(**kwargs) if r.weight > 0]
        return regularizers

    @property
    def lr(self):
        return self.optimizer.param_groups[0]['lr']

    @property
    def calc_metrics_every(self):
        return 1


def losses_to_postfix(loss_dict: Dict[str, EMA], lr: Optional[float]) -> str:
    pfix = [f"{lname}={lval}" for lname, lval in loss_dict.items()]
    if lr is not None:
        pfix.append(f"lr={lr:.2e}")
    return "  ".join(pfix)


def init_dloader_random(_):
    seed = torch.initial_seed() % 2**32  # worker-specific seed initialized by pytorch
    np.random.seed(seed)
    random.seed(seed)


def initialize_model(
        runner: 'StaticTrainer',
        **kwargs) -> LowrankModel:
    """Initialize a `LowrankModel` according to the **kwargs parameters.

    Args:
        runner: The runner object which will hold the model.
                Needed here to fetch dataset parameters.
        **kwargs: Extra parameters to pass to the model

    Returns:
        Initialized LowrankModel.
    """
    extra_args = copy(kwargs)
    extra_args.pop('global_scale', None)
    extra_args.pop('global_translation', None)

    dset = runner.test_dataset
    try:
        global_translation = dset.global_translation
    except AttributeError:
        global_translation = None
    try:
        global_scale = dset.global_scale
    except AttributeError:
        global_scale = None

    num_images = None
    if runner.train_dataset is not None:
        try:
            num_images = runner.train_dataset.num_images
        except AttributeError:
            num_images = None
    else:
        try:
            num_images = runner.test_dataset.num_images
        except AttributeError:
            num_images = None
    if hasattr(dset, "use_app_emb"):
        use_app_emb = getattr(dset, "use_app_emb")
    else:
        use_app_emb = False
    if use_app_emb:
        print("Using appearance embedding")

    model = LowrankModel(
        grid_config=extra_args.pop("grid_config"),
        aabb=dset.scene_bbox,
        is_ndc=dset.is_ndc,
        is_contracted=dset.is_contracted,
        global_scale=global_scale,
        global_translation=global_translation,
        use_appearance_embedding=use_app_emb,
        num_images=num_images,
        **extra_args)
    log.info(f"Initialized {model.__class__} model with "
             f"{sum(np.prod(p.shape) for p in model.parameters()):,} parameters, "
             f"using ndc {model.is_ndc} and contraction {model.is_contracted}. "
             f"Linear decoder: {model.linear_decoder}.")
    return model

magnerf/runners/base_trainer.py
3d-motion-magnification/3d-motion-mag
import abc
import os
from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.optim.lr_scheduler
from torch import nn

from magnerf.models.lowrank_model import LowrankModel
from magnerf.ops.losses.histogram_loss import interlevel_loss
from magnerf.raymarching.ray_samplers import RaySamples


def compute_plane_tv(t):
    batch_size, c, h, w = t.shape
    count_h = batch_size * c * (h - 1) * w
    count_w = batch_size * c * h * (w - 1)
    h_tv = torch.square(t[..., 1:, :] - t[..., :h-1, :]).sum()
    w_tv = torch.square(t[..., :, 1:] - t[..., :, :w-1]).sum()
    return 2 * (h_tv / count_h + w_tv / count_w)  # This is summing over batch and c instead of avg


def compute_plane_smoothness(t):
    batch_size, c, h, w = t.shape
    # Convolve with a second derivative filter, in the time dimension which is dimension 2
    first_difference = t[..., 1:, :] - t[..., :h-1, :]  # [batch, c, h-1, w]
    second_difference = first_difference[..., 1:, :] - first_difference[..., :h-2, :]  # [batch, c, h-2, w]
    # Take the L2 norm of the result
    return torch.square(second_difference).mean()


class Regularizer():
    def __init__(self, reg_type, initialization):
        self.reg_type = reg_type
        self.initialization = initialization
        self.weight = float(self.initialization)
        self.last_reg = None

    def step(self, global_step):
        pass

    def report(self, d):
        if self.last_reg is not None:
            d[self.reg_type].update(self.last_reg.item())

    def regularize(self, *args, **kwargs) -> torch.Tensor:
        out = self._regularize(*args, **kwargs) * self.weight
        self.last_reg = out.detach()
        return out

    @abc.abstractmethod
    def _regularize(self, *args, **kwargs) -> torch.Tensor:
        raise NotImplementedError()

    def __str__(self):
        return f"Regularizer({self.reg_type}, weight={self.weight})"


class PlaneTV(Regularizer):
    def __init__(self, initial_value, what: str = 'field'):
        if what not in {'field', 'motion_field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field", "motion_field", or "proposal_network" '
                             f'but {what} was passed.')
        name = f'planeTV-{what[:2]}'
        super().__init__(name, initial_value)
        self.what = what

    def step(self, global_step):
        pass

    def _regularize(self, model: LowrankModel, **kwargs):
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        elif self.what == 'motion_field':
            multi_res_grids = model.field.motion_grids
        else:
            raise NotImplementedError(self.what)
        total = 0
        # Note: input to compute_plane_tv should be of shape [batch_size, c, h, w]
        for grids in multi_res_grids:
            if len(grids) == 3:
                spatial_grids = [0, 1, 2]
            else:
                spatial_grids = [0, 1, 3]  # These are the spatial grids; the others are spatiotemporal
            for grid_id in spatial_grids:
                _grid = grids[grid_id]
                if len(_grid.shape) > 4:
                    for t in range(_grid.shape[-1]):
                        total += compute_plane_tv(_grid[..., t])
                else:
                    total += compute_plane_tv(_grid)
            for _grid in grids:
                # grid: [1, c, h, w]
                if len(_grid.shape) > 4:
                    for t in range(_grid.shape[-1]):
                        total += compute_plane_tv(_grid[..., t])
                else:
                    total += compute_plane_tv(_grid)
        return total


class TimeSmoothness(Regularizer):
    def __init__(self, initial_value, what: str = 'field'):
        if what not in {'field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field" or "proposal_network" '
                             f'but {what} was passed.')
        name = f'time-smooth-{what[:2]}'
        super().__init__(name, initial_value)
        self.what = what

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        else:
            raise NotImplementedError(self.what)
        total = 0
        # model.grids is 6 x [1, rank * F_dim, reso, reso]
        for grids in multi_res_grids:
            if len(grids) == 3:
                #time_grids = []
                time_grids = [0, 1, 2]
            else:
                time_grids = [2, 4, 5]
            for grid_id in time_grids:
                _grid = grids[grid_id]
                if len(_grid.shape) > 4:
                    gs = _grid.shape
                    B, C, T = gs[0], gs[1], gs[-1]
                    total += compute_plane_smoothness(_grid.view(B, C, -1, T).permute(0, 1, 3, 2))
                else:
                    total += compute_plane_smoothness(_grid)
        return torch.as_tensor(total)


class HistogramLoss(Regularizer):
    def __init__(self, initial_value):
        super().__init__('histogram-loss', initial_value)

        self.visualize = False
        self.count = 0

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        if self.visualize:
            if self.count % 500 == 0:
                prop_idx = 0
                fine_idx = 1
                # proposal info
                weights_proposal = model_out["weights_list"][prop_idx].detach().cpu().numpy()
                spacing_starts_proposal = model_out["ray_samples_list"][prop_idx].spacing_starts
                spacing_ends_proposal = model_out["ray_samples_list"][prop_idx].spacing_ends
                sdist_proposal = torch.cat([
                    spacing_starts_proposal[..., 0],
                    spacing_ends_proposal[..., -1:, 0]
                ], dim=-1).detach().cpu().numpy()

                # fine info
                weights_fine = model_out["weights_list"][fine_idx].detach().cpu().numpy()
                spacing_starts_fine = model_out["ray_samples_list"][fine_idx].spacing_starts
                spacing_ends_fine = model_out["ray_samples_list"][fine_idx].spacing_ends
                sdist_fine = torch.cat([
                    spacing_starts_fine[..., 0],
                    spacing_ends_fine[..., -1:, 0]
                ], dim=-1).detach().cpu().numpy()

                for i in range(10):  # plot 10 rays
                    fix, ax1 = plt.subplots()

                    delta = np.diff(sdist_proposal[i], axis=-1)
                    ax1.bar(sdist_proposal[i, :-1], weights_proposal[i].squeeze() / delta, width=delta, align="edge", label='proposal', alpha=0.7, color="b")
                    ax1.legend()
                    ax2 = ax1.twinx()

                    delta = np.diff(sdist_fine[i], axis=-1)
                    ax2.bar(sdist_fine[i, :-1], weights_fine[i].squeeze() / delta, width=delta, align="edge", label='fine', alpha=0.3, color='r')
                    ax2.legend()
                    os.makedirs(f'histogram_loss/{self.count}', exist_ok=True)
                    plt.savefig(f'./histogram_loss/{self.count}/batch_{i}.png')
                    plt.close()
                    plt.cla()
                    plt.clf()
            self.count += 1
        return interlevel_loss(model_out['weights_list'], model_out['ray_samples_list'])


class L1ProposalNetwork(Regularizer):
    def __init__(self, initial_value):
        super().__init__('l1-proposal-network', initial_value)

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        grids = [p.grids for p in model.proposal_networks]
        total = 0.0
        for pn_grids in grids:
            for grid in pn_grids:
                total += torch.abs(grid).mean()
        return torch.as_tensor(total)


class L1MotionField(Regularizer):
    def __init__(self, initial_value):
        super().__init__('l1-motion-field', initial_value)

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        grids = model.field.motion_grids
        total = 0.0
        for pn_grids in grids:
            for grid in pn_grids:
                total += torch.abs(grid).mean()
        return torch.as_tensor(total)


class DepthTV(Regularizer):
    def __init__(self, initial_value):
        super().__init__('tv-depth', initial_value)

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        depth = model_out['depth']
        tv = compute_plane_tv(
            depth.reshape(64, 64)[None, None, :, :]
        )
        return tv


class L1TimePlanes(Regularizer):
    def __init__(self, initial_value, what='field'):
        if what not in {'field', 'proposal_network'}:
            raise ValueError(f'what must be one of "field" or "proposal_network" '
                             f'but {what} was passed.')
        super().__init__(f'l1-time-{what[:2]}', initial_value)
        self.what = what

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        # model.grids is 6 x [1, rank * F_dim, reso, reso]
        multi_res_grids: Sequence[nn.ParameterList]
        if self.what == 'field':
            multi_res_grids = model.field.grids
        elif self.what == 'proposal_network':
            multi_res_grids = [p.grids for p in model.proposal_networks]
        else:
            raise NotImplementedError(self.what)

        total = 0.0
        for grids in multi_res_grids:
            if len(grids) == 3:
                continue
            else:
                # These are the spatiotemporal grids
                spatiotemporal_grids = [2, 4, 5]
            for grid_id in spatiotemporal_grids:
                total += torch.abs(1 - grids[grid_id]).mean()
        return torch.as_tensor(total)


class DistortionLoss(Regularizer):
    def __init__(self, initial_value):
        super().__init__('distortion-loss', initial_value)

    def _regularize(self, model: LowrankModel, model_out, **kwargs) -> torch.Tensor:
        w = model_out['weights_list'][-1]
        rs: RaySamples = model_out['ray_samples_list'][-1]
        m = (rs.starts + rs.ends) / 2
        interval = rs.deltas

        loss_uni = (1/3) * (interval * w.pow(2)).sum(dim=-1).mean()
        wm = (w * m)
        w_cumsum = w.cumsum(dim=-1)
        wm_cumsum = wm.cumsum(dim=-1)
        loss_bi_0 = wm[..., 1:] * w_cumsum[..., :-1]
        loss_bi_1 = w[..., 1:] * wm_cumsum[..., :-1]
        loss_bi = 2 * (loss_bi_0 - loss_bi_1).sum(dim=-1).mean()
        return loss_bi + loss_uni

class L1GridChange(Regularizer):
    def __init__(self, initial_value):
        super().__init__('l1-dgrid', initial_value)

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        grids = model.field.delta_grids
        total = 0.0
        for pn_grids in grids:
            for grid in pn_grids:
                total += torch.abs(grid).mean()
        return torch.as_tensor(total)

class TVGridChange(Regularizer):
    def __init__(self, initial_value):
        super().__init__('tv-dgrid', initial_value)

    def _regularize(self, model: LowrankModel, **kwargs) -> torch.Tensor:
        grids = model.field.delta_grids
        total = 0.0
        for pn_grids in grids:
            for grid in pn_grids:
                total += compute_plane_tv(grid)

        return torch.as_tensor(total)

magnerf/runners/regularization.py
3d-motion-magnification/3d-motion-mag
import math
import os
from collections import defaultdict
from typing import Dict, MutableMapping, Union, Any

import pandas as pd
import torch
import torch.utils.data

from magnerf.datasets import SyntheticNerfDataset, HandheldDataset, TripodDataset
from magnerf.models.lowrank_model import LowrankModel
from magnerf.utils.ema import EMA
from magnerf.utils.my_tqdm import tqdm
from magnerf.utils.parse_args import parse_optint
from .base_trainer import BaseTrainer, init_dloader_random, initialize_model
from .regularization import (
    PlaneTV, HistogramLoss, L1ProposalNetwork, DepthTV, DistortionLoss, L1GridChange, TVGridChange
)


class StaticTrainer(BaseTrainer):
    def __init__(self,
                 tr_loader: torch.utils.data.DataLoader,
                 ts_dset: torch.utils.data.TensorDataset,
                 tr_dset: torch.utils.data.TensorDataset,
                 num_steps: int,
                 logdir: str,
                 expname: str,
                 train_fp16: bool,
                 save_every: int,
                 valid_every: int,
                 save_outputs: bool,
                 device: Union[str, torch.device],
                 **kwargs
                 ):
        self.test_dataset = ts_dset
        self.train_dataset = tr_dset
        self.is_ndc = self.test_dataset.is_ndc
        self.is_contracted = self.test_dataset.is_contracted

        super().__init__(
            train_data_loader=tr_loader,
            num_steps=num_steps,
            logdir=logdir,
            expname=expname,
            train_fp16=train_fp16,
            save_every=save_every,
            valid_every=int(valid_every),
            save_outputs=save_outputs,
            device=device,
            **kwargs
        )

    def eval_step(self, data, **kwargs) -> MutableMapping[str, torch.Tensor]:
        """
        Note that here `data` contains a whole image. we need to split it up before tracing
        for memory constraints.
        """
        super().eval_step(data, **kwargs)
        batch_size = self.eval_batch_size

        if "timestamps" not in data:
            data["timestamps"] = None

        channels = {"rgb", "depth", "proposal_depth"}
        with torch.cuda.amp.autocast(enabled=self.train_fp16), torch.no_grad():
            rays_o = data["rays_o"]
            rays_d = data["rays_d"]
            # near_far and bg_color are constant over mini-batches
            near_far = data["near_fars"].to(self.device)
            bg_color = data["bg_color"]
            if isinstance(bg_color, torch.Tensor):
                bg_color = bg_color.to(self.device)
            timestamps = data['timestamps']
            if isinstance(timestamps, torch.Tensor):
                timestamps = timestamps.to(self.device)

            preds = defaultdict(list)
            for b in range(math.ceil(rays_o.shape[0] / batch_size)):
                rays_o_b = rays_o[b * batch_size: (b + 1) * batch_size].to(self.device)
                rays_d_b = rays_d[b * batch_size: (b + 1) * batch_size].to(self.device)
                outputs = self.model(rays_o_b, rays_d_b, near_far=near_far,
                                     bg_color=bg_color, timestamps=timestamps)
                for k, v in outputs.items():
                    if k in channels or "depth" in k:
                        preds[k].append(v.cpu())
        return {k: torch.cat(v, 0) for k, v in preds.items()}

    def train_step(self, data: Dict[str, Union[int, torch.Tensor]], **kwargs):
        return super().train_step(data, **kwargs)

    def post_step(self, progress_bar):
        super().post_step(progress_bar)

    def pre_epoch(self):
        super().pre_epoch()
        self.train_dataset.reset_iter()

    @torch.no_grad()
    def validate(self):
        dataset = self.test_dataset
        per_scene_metrics = defaultdict(list)
        pb = tqdm(total=len(dataset), desc=f"Test scene {dataset.name}")
        for img_idx, data in enumerate(dataset):
            ts_render = self.eval_step(data)
            out_metrics, _, _ = self.evaluate_metrics(
                data["imgs"], ts_render, dset=dataset, img_idx=img_idx,
                name=None, save_outputs=self.save_outputs)
            for k, v in out_metrics.items():
                per_scene_metrics[k].append(v)
            pb.set_postfix_str(f"PSNR={out_metrics['psnr']:.2f}", refresh=False)
            pb.update(1)
        pb.close()
        val_metrics = [
            self.report_test_metrics(per_scene_metrics, extra_name="")
        ]
        df = pd.DataFrame.from_records(val_metrics)
        df.to_csv(os.path.join(self.log_dir, f"test_metrics_step{self.global_step}.csv"))

    def get_save_dict(self):
        base_save_dict = super().get_save_dict()
        return base_save_dict

    def load_model(self, checkpoint_data, training_needed: bool = True):
        super().load_model(checkpoint_data, training_needed)

    def init_epoch_info(self):
        ema_weight = 0.9  # higher places higher weight to new observations
        loss_info = defaultdict(lambda: EMA(ema_weight))
        return loss_info

    def init_model(self, **kwargs) -> LowrankModel:
        return initialize_model(self, **kwargs)

    def get_regularizers(self, **kwargs):
        regularizers = [
            PlaneTV(kwargs.get('plane_tv_weight', 0.0), what='field'),
            PlaneTV(kwargs.get('plane_tv_weight_proposal_net', 0.0), what='proposal_network'),
            HistogramLoss(kwargs.get('histogram_loss_weight', 0.0)),
            L1ProposalNetwork(kwargs.get('l1_proposal_net_weight', 0.0)),
            DepthTV(kwargs.get('depth_tv_weight', 0.0)),
            DistortionLoss(kwargs.get('distortion_loss_weight', 0.0)),
            L1GridChange(kwargs.get('gridchange_loss_weight', 0.0)),
            TVGridChange(kwargs.get('tv_gridchange_loss_weight', 0.0)),
        ]

        return regularizers

    @property
    def calc_metrics_every(self):
        return 5


def decide_dset_type(dd) -> str:
    if ("chair" in dd or "drums" in dd or "ficus" in dd or "hotdog" in dd
            or "lego" in dd or "materials" in dd or "mic" in dd
            or "ship" in dd):
        return "synthetic"
    elif ("fern" in dd or "flower" in dd or "fortress" in dd
          or "horns" in dd or "leaves" in dd or "orchids" in dd
          or "room" in dd or "trex" in dd):
        return "llff"
    else:
        raise RuntimeError(f"data_dir {dd} not recognized.")


def init_tr_data(data_downsample: float, data_dir: str, **kwargs):
    batch_size = int(kwargs['batch_size'])
    dset_type = kwargs["dset_type"] if "dset_type" in kwargs else "synthetic"

    if dset_type == "synthetic":
        max_tr_frames = parse_optint(kwargs.get('max_tr_frames'))
        radius = kwargs["radius"] if "radius" in kwargs else 1.3
        near_far = kwargs["near_far"] if "near_far" in kwargs else None
        dset = SyntheticNerfDataset(
            data_dir, split='train', downsample=data_downsample,
            max_frames=max_tr_frames, batch_size=batch_size, radius=radius, near_far=near_far)
    elif dset_type == "handheld":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = HandheldDataset(
            data_dir, split='train', downsample=int(data_downsample), hold_every=hold_every,
            batch_size=batch_size,
            contraction=kwargs.get('contract', 'True'),
            train_cutoff=int(kwargs.get('train_cutoff', -1)),
            dynamic_start=int(kwargs.get('dynamic_start', 0)),
            aabb=kwargs.get('bbox'),
            use_app=kwargs.get('use_app_embedding', False),
            near_plane=float(kwargs.get('near_plane', 0.2)),
            far_plane=float(kwargs.get('far_plane', 30.0)),
            bg_color_aug=kwargs.get('bg_color_aug', 'black'),
        )
    elif dset_type == "tripod":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = TripodDataset(
            data_dir, split='train', downsample=int(data_downsample), hold_every=hold_every,
            batch_size=batch_size,
            contraction=kwargs.get('contract', False),
            frame_offset=int(kwargs.get('frame_offset', 2)),
            aabb=kwargs.get('bbox'),
            use_app=kwargs.get('use_app_embedding', False),
            near_plane=float(kwargs.get('near_plane', 0.2)),
            far_plane=float(kwargs.get('far_plane', 30.0)),
            bg_color_aug=kwargs.get('bg_color_aug', 'black'),
        )
    else:
        raise ValueError(f"Dataset type {dset_type} invalid.")
    dset.reset_iter()

    tr_loader = torch.utils.data.DataLoader(
        dset, num_workers=4, prefetch_factor=4, pin_memory=True,
        batch_size=None, worker_init_fn=init_dloader_random)

    return {
        "tr_dset": dset,
        "tr_loader": tr_loader,
    }


def init_ts_data(data_downsample: float, data_dir: str, split: str, **kwargs):

    #dset_type = decide_dset_type(data_dir)
    dset_type = kwargs["dset_type"] if "dset_type" in kwargs else "synthetic"

    if dset_type == "synthetic":
        max_ts_frames = parse_optint(kwargs.get('max_ts_frames'))
        dset = SyntheticNerfDataset(
            data_dir, split=split, downsample=1, max_frames=max_ts_frames)
    elif dset_type == "handheld":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = HandheldDataset(
            data_dir, split=split, downsample=data_downsample, hold_every=hold_every,
            contraction=kwargs.get('contract', 'True'),
            train_cutoff=int(kwargs.get('train_cutoff', -1)),
            dynamic_start=int(kwargs.get('dynamic_start', 0)),
            aabb=kwargs.get('bbox'),
            start_interp_index=int(kwargs.get('start_interp_index', 5)),
            use_app=kwargs.get('use_app_embedding', 'False'),
            near_plane_offset=float(kwargs.get('near_plane_offset', 0.0)),
            near_plane=float(kwargs.get('near_plane', 0.2)),
            far_plane=float(kwargs.get('far_plane', 30.0)),
            bg_color_aug=kwargs.get('bg_color_aug', 'black'),
        )
    elif dset_type == "tripod":
        hold_every = parse_optint(kwargs.get('hold_every'))
        dset = TripodDataset(
            data_dir, split=split, downsample=data_downsample, hold_every=hold_every,
            contraction=kwargs.get('contract', False),
            frame_offset=int(kwargs.get('frame_offset', 2)),
            aabb=kwargs.get('bbox'),
            start_interp_index=int(kwargs.get('start_interp_index', 5)),
            use_app=kwargs.get('use_app_embedding', 'False'),
            near_plane_offset=float(kwargs.get('near_plane_offset', 0.0)),
            near_plane=float(kwargs.get('near_plane', 0.2)),
            far_plane=float(kwargs.get('far_plane', 30.0)),
            bg_color_aug=kwargs.get('bg_color_aug', 'black'),
        )
    else:
        raise ValueError(f"Dataset type {dset_type} invalid.")
    return {"ts_dset": dset}


def load_data(data_downsample, data_dir, validate_only, render_only, **kwargs):
    od: Dict[str, Any] = {}
    if not validate_only:
        od.update(init_tr_data(data_downsample, data_dir, **kwargs))
    else:
        od.update(tr_loader=None, tr_dset=None)
    test_split = 'render' if render_only else 'test'
    od.update(init_ts_data(data_downsample, data_dir, split=test_split, **kwargs))
    return od

magnerf/runners/static_trainer.py
3d-motion-magnification/3d-motion-mag
import argparse
import importlib.util
import logging
import os
import pprint
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from typing import List, Dict, Any
import tempfile

import numpy as np
import imageio


def get_freer_gpu():
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_fname = os.path.join(tmpdir, "tmp")
        os.system(f'nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >"{tmp_fname}"')
        if os.path.isfile(tmp_fname):
            memory_available = [int(x.split()[2]) for x in open(tmp_fname, 'r').readlines()]
            if len(memory_available) > 0:
                return np.argmax(memory_available)
    return None  # The grep doesn't work with all GPUs. If it fails we ignore it.

gpu = get_freer_gpu()
if gpu is not None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu)
    print(f"CUDA_VISIBLE_DEVICES set to {gpu}")
else:
    print(f"Did not set GPU.")

import torch
import torch.utils.data
from magnerf.runners import static_trainer
from magnerf.utils.create_rendering import render_to_path
from magnerf.utils.parse_args import parse_optfloat


def setup_logging(log_level=logging.INFO):
    handlers = [logging.StreamHandler(sys.stdout)]
    logging.basicConfig(level=log_level,
                        format='%(asctime)s|%(levelname)8s| %(message)s',
                        handlers=handlers,
                        force=True)


def load_data(model_type: str, data_downsample, data_dirs, validate_only: bool, render_only: bool, **kwargs):
    data_downsample = parse_optfloat(data_downsample, default_val=1.0)

    return static_trainer.load_data(
        data_downsample, data_dirs, validate_only=validate_only,
        render_only=render_only, **kwargs)


def init_trainer(model_type: str, **kwargs):
    from magnerf.runners import static_trainer
    return static_trainer.StaticTrainer(**kwargs)


def save_config(config):
    log_dir = os.path.join(config['logdir'], config['expname'])
    os.makedirs(log_dir, exist_ok=True)

    with open(os.path.join(log_dir, 'config.py'), 'wt') as out:
        out.write('config = ' + pprint.pformat(config))

    with open(os.path.join(log_dir, 'config.csv'), 'w') as f:
        for key in config.keys():
            f.write("%s\t%s\n" % (key, config[key]))


def vis_grid_features(trainer, base_logdir):
    new_field_feature_imgs = {}
    for n, p in trainer.model.field.named_parameters():
        if 'grids' in n and 'delta' not in n:
            new_field_feature_imgs[n] = p.data

    planes = new_field_feature_imgs['grids.0.0'].squeeze(0).cpu()
    out = planes[:3].permute(1, 2, 0)
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    mask = np.zeros_like(out)
    h, w = mask.shape[0], mask.shape[1]
    mask[(h//4):-(h//4), (w//4):-(w//4)] = 1.0
    imageio.imwrite(f'{base_logdir}/xy.png', out)
    #imageio.imwrite(f'{base_logdir}/xy_m.png', out * mask)

    planes = new_field_feature_imgs['grids.0.1'].squeeze(0).cpu()
    out = planes[:3].permute(1, 2, 0)
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    imageio.imwrite(f'{base_logdir}/xz.png', out)
    #imageio.imwrite(f'{base_logdir}/xz_m.png', out * mask)

    planes = new_field_feature_imgs['grids.0.2'].squeeze(0).cpu()
    out = planes[:3].permute(1, 2, 0)
    out = (out - out.min()) / (out.max() - out.min())
    out = np.uint8(out * 255)
    imageio.imwrite(f'{base_logdir}/yz.png', out)
    #imageio.imwrite(f'{base_logdir}/yz_m.png', out * mask)


def main():
    setup_logging()

    p = argparse.ArgumentParser(description="")

    p.add_argument('--render-only', action='store_true')
    p.add_argument('--validate-only', action='store_true')
    p.add_argument('--spacetime-only', action='store_true')
    p.add_argument('--config-path', type=str, required=True)
    p.add_argument('--log-dir', type=str, default=None)
    p.add_argument('--timestep', type=str, default="000")
    p.add_argument('--seed', type=int, default=0)
    p.add_argument('--c2f', action='store_true')
    p.add_argument('override', nargs=argparse.REMAINDER)

    args = p.parse_args()

    # Set random seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Import config
    spec = importlib.util.spec_from_file_location(os.path.basename(args.config_path), args.config_path)
    cfg = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(cfg)
    config: Dict[str, Any] = cfg.config
    # Process overrides from argparse into config
    # overrides can be passed from the command line as key=value pairs. E.g.
    # python plenoxels/main.py --config-path plenoxels/config/cfg.py max_ts_frames=200
    # note that all values are strings, so code should assume incorrect data-types for anything
    # that's derived from config - and should not a string.
    overrides: List[str] = args.override
    overrides_dict = {ovr.split("=")[0]: ovr.split("=")[1] for ovr in overrides}

    base_logdir = config["logdir"]
    expname = config["expname"]
    if "runname" in config:
        runname = config["runname"]
    else:
        runname = '.'
    data_dir = config["data_dirs"]

    if "expname" in overrides_dict:
        expname = overrides_dict["expname"]
    if "runname" in overrides_dict:
        runname = overrides_dict["runname"]
    if "data_dirs" in overrides_dict:
        data_dir = overrides_dict["data_dirs"]

    overrides_dict["data_dirs"] = f"{data_dir}/{expname}/{args.timestep}"

    expname = f"{expname}/{runname}"
    base_logdir = f"{base_logdir}/{expname}/output"
    overrides_dict["expname"] = f"{expname}/output/{args.timestep}"
    config.update(overrides_dict)

    model_type = "static"
    validate_only = args.validate_only
    render_only = args.render_only
    spacetime_only = args.spacetime_only
    if validate_only and render_only:
        raise ValueError("render_only and validate_only are mutually exclusive.")
    if render_only and spacetime_only:
        raise ValueError("render_only and spacetime_only are mutually exclusive.")
    if validate_only and spacetime_only:
        raise ValueError("validate_only and spacetime_only are mutually exclusive.")

    if 'bbox' in config and isinstance(config["bbox"], str):
        config.update({'bbox': list(map(float, config["bbox"].split(',')))})
    if isinstance(config["num_samples"], str):
        config.update({'num_samples': int(config["num_samples"])})

    pprint.pprint(config)
    if not validate_only or render_only:
        save_config(config)

    data = load_data(model_type, validate_only=validate_only, render_only=render_only or spacetime_only, **config)

    config.update(data)
    config.update({'lr': float(config["lr"])})
    trainer = init_trainer(model_type, **config)

    finetune_iters = int(overrides_dict.get('finetune_iters', 10001))
    finetune_lr = config.get('finetune_lr', None)
    print(finetune_lr)
    if args.timestep != '000' and not render_only:
        config.update(
            {'num_steps': finetune_iters, 
             'save_every': finetune_iters-1,
             'valid_every': finetune_iters-1,
             }
        )
        if finetune_lr is not None:
            config.update({'lr': finetune_lr})

        trainer = init_trainer(model_type, **config)
        checkpoint_path = os.path.join(base_logdir, "000", "model.pth")

        ckpt = torch.load(checkpoint_path)
        trainer.model.load_state_dict(ckpt['model'])

        for n, p in trainer.model.proposal_networks.named_parameters():
            p.requires_grad = False
        for n, p in trainer.model.field.named_parameters():
            if not 'delta_grids' in n:
                p.requires_grad = False
        for n, p in trainer.model.named_parameters():
            print(f'{n} | {p.shape} | Trainable: {p.requires_grad}')

        trainer.model.field.use_delta_features = True

    if validate_only:
        checkpoint_path = os.path.join(base_logdir, f"{(int(args.timestep)):03d}", "model.pth")
        ckpt = torch.load(checkpoint_path)
        trainer.model.load_state_dict(ckpt['model'])
        trainer.validate()

    elif render_only:
        checkpoint_path = os.path.join(base_logdir, f"{(int(args.timestep)):03d}", "model.pth")
        ckpt = torch.load(checkpoint_path)
        trainer.model.load_state_dict(ckpt['model'])

        render_to_path(trainer, extra_name="")
        vis_grid_features(trainer, base_logdir)

    else:
        if args.timestep != '000' or not args.c2f:
            trainer.train()
        else:
            factors = [1, 1, 1, 1, 1]
            steps = [100, 500, 1000, 2000, 20000]
            for f, s in zip(factors, steps):
                print(f'Training at factor {f} for {s} steps')
                config.update({'data_downsample': int(f)})
                data = load_data(model_type, validate_only=validate_only, render_only=render_only or spacetime_only, **config)
                for k,v in data.items():
                    setattr(trainer, k, v)
                trainer.num_steps = s
                trainer.train()
                vis_grid_features(trainer, base_logdir)
            trainer.save_model()

if __name__ == "__main__":
    main()

magnerf/train.py
3d-motion-magnification/3d-motion-mag
"""Entry point for simple renderings, given a trainer and some poses."""
import os
import logging as log
from typing import Union

import torch

from magnerf.models.lowrank_model import LowrankModel
from magnerf.utils.my_tqdm import tqdm
from magnerf.ops.image.io import write_video_to_file
from magnerf.runners.static_trainer import StaticTrainer

import matplotlib.pyplot as plt


def bandpass_filter(data_3d, mask):
    d_data_3d = data_3d - data_3d[..., 0:1]
    data_f = torch.fft.rfft(d_data_3d)
    data_f[~mask[:len(data_f)]] = 0
    filtered = torch.fft.irfft(data_f)

    return filtered

def create_bandpass_mask(lowcut, highcut, n, fs):
    freq = torch.arange(n) / n
    freq = freq * fs

    freq = freq[:(n//2 + 1)]
    mask = (freq > lowcut) * (freq <= highcut)
    mask = mask.unsqueeze(0).unsqueeze(0)

    return mask


@torch.no_grad()
def render_to_path(trainer: StaticTrainer, extra_name: str = "") -> None:
    """Render all poses in the `test_dataset`, saving them to file
    Args:
        trainer: The trainer object which is used for rendering
        extra_name: String to append to the saved file-name
    """
    dataset = trainer.test_dataset

    pb = tqdm(total=len(dataset.poses), desc=f"Rendering scene")
    frames = []
    for img_idx, data in enumerate(dataset):
        ts_render = trainer.eval_step(data)

        if isinstance(dataset.img_h, int):
            img_h, img_w = dataset.img_h, dataset.img_w
        else:
            img_h, img_w = dataset.img_h[img_idx], dataset.img_w[img_idx]
        preds_rgb = (
            ts_render["rgb"]
            .reshape(img_h, img_w, 3)
            .cpu()
            .clamp(0, 1)
            .mul(255.0)
            .byte()
            .numpy()
        )
        frames.append(preds_rgb)
        pb.update(1)
    pb.close()

    out_fname = os.path.join(trainer.log_dir, f"rendering_path_{extra_name}.mp4")
    write_video_to_file(out_fname, frames)
    log.info(f"Saved rendering path with {len(frames)} frames to {out_fname}")


def normalize_for_disp(img):
    img = img - torch.min(img)
    img = img / torch.max(img)
    return img


@torch.no_grad()
def decompose_space_time(trainer: StaticTrainer, extra_name: str = "") -> None:
    """Render space-time decomposition videos for poses in the `test_dataset`.

    The space-only part of the decomposition is obtained by setting the time-planes to 1.
    The time-only part is obtained by simple subtraction of the space-only part from the full
    rendering.

    Args:
        trainer: The trainer object which is used for rendering
        extra_name: String to append to the saved file-name
    """
    chosen_cam_idx = 60
    model: LowrankModel = trainer.model
    dataset = trainer.test_dataset

    # Store original parameters from main field and proposal-network field
    parameters = []
    for multires_grids in model.field.grids:
        parameters.append([grid.data for grid in multires_grids])
    pn_parameters = []
    for pn in model.proposal_networks:
        pn_parameters.append([grid_plane.data for grid_plane in pn.grids])

    for i in range(len(parameters)):
        for plane_idx in [2, 4, 5]:
            x = parameters[i][plane_idx].squeeze()
            x0 = x[:, 0:1, :]
            mask = create_bandpass_mask(2.9, 3.1, n = x.shape[1], fs = 30)
            delta = x - x0
            filtered_delta = bandpass_filter(delta.permute(0, 2, 1), mask.repeat(x.shape[0], x.shape[2], 1))
            filtered_delta = filtered_delta.permute(0, 2, 1)
            plt.imsave(f'{i}_{plane_idx}.png', x[0].cpu().numpy())
            new_x = x0 + 10 * filtered_delta
            parameters[i][plane_idx] = new_x.unsqueeze(0)

    camdata = None
    camdata_list = []
    for img_idx, data in enumerate(dataset):
        if img_idx == chosen_cam_idx:
            camdata = data
        camdata_list.append(data)
    if camdata is None:
        raise ValueError(f"Cam idx {chosen_cam_idx} invalid.")

    num_frames = 30#img_idx + 1
    frames = []
    full_frames, spatial_frames, diff_frames = [], [], []
    for img_idx in tqdm(range(num_frames), desc="Rendering scene with separate space and time components"):
        #camdata = camdata_list[img_idx]

        # Linearly interpolated timestamp, normalized between -1, 1
        camdata["timestamps"] = torch.Tensor([img_idx / num_frames]) * 2 - 1

        if isinstance(dataset.img_h, int):
            img_h, img_w = dataset.img_h, dataset.img_w
        else:
            img_h, img_w = dataset.img_h[img_idx], dataset.img_w[img_idx]

        # Full model: turn on time-planes
        for i in range(len(model.field.grids)):
            for plane_idx in [2, 4, 5]:
                model.field.grids[i][plane_idx].data = parameters[i][plane_idx]
        for i in range(len(model.proposal_networks)):
            for plane_idx in [2, 4, 5]:
                model.proposal_networks[i].grids[plane_idx].data = pn_parameters[i][plane_idx]
        preds = trainer.eval_step(camdata)
        full_out = preds["rgb"].reshape(img_h, img_w, 3).cpu()

        # Space-only model: turn off time-planes
        for i in range(len(model.field.grids)):
            for plane_idx in [2, 4, 5]:  # time-grids off
                model.field.grids[i][plane_idx].data = torch.ones_like(parameters[i][plane_idx])
        for i in range(len(model.proposal_networks)):
            for plane_idx in [2, 4, 5]:
                model.proposal_networks[i].grids[plane_idx].data = torch.ones_like(pn_parameters[i][plane_idx])
        preds = trainer.eval_step(camdata)
        spatial_out = preds["rgb"].reshape(img_h, img_w, 3).cpu()

        # Temporal model: full - space
        temporal_out = full_out - spatial_out
        full_frames.append(full_out)
        spatial_frames.append(spatial_out)
        diff_frames.append(temporal_out)

    full_out = torch.stack(full_frames)
    spatial_out = torch.stack(spatial_frames)
    diff_frames = torch.stack(diff_frames)
    temporal_out = normalize_for_disp(diff_frames)

    frames = torch.cat([full_out, spatial_out, temporal_out], dim=2).clamp(0, 1).mul(255.0).byte().numpy()

    out_fname = os.path.join(trainer.log_dir, f"spacetime_{extra_name}.mp4")
    write_video_to_file(out_fname, frames)
    log.info(f"Saved rendering path with {len(frames)} frames to {out_fname}")

magnerf/utils/create_rendering.py
3d-motion-magnification/3d-motion-mag

class EMA():
    def __init__(self, weighting=0.9):
        self.weighting = weighting
        self.val = None

    def update(self, val):
        if self.val is None:
            self.val = val
        else:
            self.val = self.weighting * val + (1 - self.weighting) * self.val

    @property
    def value(self):
        return self.val

    def __str__(self):
        return f"{self.val:.2e}"

magnerf/utils/ema.py
3d-motion-magnification/3d-motion-mag
"""Evaluate several metrics for a pretrained model. Handles video and static."""
import re
import glob
import os
from collections import defaultdict

import numpy as np
import torch

from magnerf.ops.image import metrics
from magnerf.ops.image.io import read_mp4, read_png


def eval_static_metrics(static_dir):
    all_file_names = glob.glob(os.path.join(static_dir, r"*.png"))
    # Collect all GT+Pred files per step
    files_per_step = defaultdict(list)
    for f in all_file_names:
        if "depth" in f:
            continue
        match = re.match(r".*step([0-9]+)-([0-9])+\.png", f)
        if match is None:
            continue
        step = int(match.group(1))
        files_per_step[step].append(f)
    steps = list(files_per_step.keys())
    max_step = max(steps)
    print(f"Evaluating static metrics for {static_dir} at step "
          f"{max_step} with {len(files_per_step[max_step])} files.")
    frames = [read_png(f) for f in files_per_step[max_step]]
    h1 = frames[0].shape[0] // 3
    pred_frames = [f[:h1] for f in frames]
    gt_frames = [f[h1:2*h1] for f in frames]

    psnrs, ssims, msssims, lpipss = [], [], [], []
    for pred, gt in zip(pred_frames, gt_frames):
        psnrs.append(metrics.psnr(pred, gt))
        ssims.append(metrics.ssim(pred, gt))
        msssims.append(metrics.msssim(pred, gt))
        lpipss.append(metrics.rgb_lpips(pred, gt, net_name="alex"))
    psnr = np.mean(psnrs)
    ssim = np.mean(ssims)
    msssim = np.mean(msssims)
    lpips = np.mean(lpipss)

    print()
    print(f"Images at {static_dir} - step {max_step}. Metrics:")
    print(f"PSNR = {psnr}")
    print(f"SSIM = {ssim}")
    print(f"MS-SSIM = {msssim}")
    print(f"Alex-LPIPS= {lpips}")
    print()


def eval_video_metrics(video_path):
    frames = read_mp4(video_path)
    h1 = frames[0].shape[0] // 3
    pred_frames = [f[:h1].numpy() for f in frames]
    gt_frames = [f[h1:2*h1].numpy() for f in frames]

    psnrs, ssims, msssims, lpipss = [], [], [], []
    for pred, gt in zip(pred_frames, gt_frames):
        pred = torch.from_numpy(pred).float().div(255)
        gt = torch.from_numpy(gt).float().div(255)
        psnrs.append(metrics.psnr(pred, gt))
        ssims.append(metrics.ssim(pred, gt))
        msssims.append(metrics.msssim(pred, gt))
        lpipss.append(metrics.rgb_lpips(pred, gt, net_name="alex"))
    psnr = np.mean(psnrs)
    ssim = np.mean(ssims)
    msssim = np.mean(msssims)
    lpips = np.mean(lpipss)

    flip = metrics.flip(pred_frames=pred_frames, gt_frames=gt_frames, interval=10)
    # jod = metrics.jod(pred_frames=pred_frames, gt_frames=gt_frames)

    print()
    print(f"Video at {video_path} metrics:")
    print(f"PSNR = {psnr}")
    print(f"SSIM = {ssim}")
    # print(f"MS-SSIM = {msssim}")
    # print(f"Alex-LPIPS = {lpips}")
    # print(f"FLIP = {flip}")
    # print(f"JOD = {jod}")
    print()
    print()


dnerf_scenes = ['hellwarrior', 'mutant', 'hook', 'bouncingballs', 'lego', 'trex', 'standup', 'jumpingjacks']
types = ['linear', 'mlp']

if __name__ == "__main__":
    for modeltype in types:
        for scene in dnerf_scenes:
            eval_video_metrics(f"logs/dnerf_{modeltype}_refactor1/{scene}_concat32_lr0.01_time0.1_tv0.0001_proptime0.001_proptv0.0001_distort0/step30000.mp4")

magnerf/utils/eval_model_metrics.py
3d-motion-magnification/3d-motion-mag
import os
from tqdm.auto import tqdm as original_tqdm
import functools
import logging as log


tqdm_disabled = os.environ.get('DISABLE_TQDM', False)
if tqdm_disabled:
    log.info("Disabling tqdm")
tqdm = original_tqdm
tqdm.__init__ = functools.partialmethod(
    tqdm.__init__, mininterval=1, disable=tqdm_disabled)

magnerf/utils/my_tqdm.py
3d-motion-magnification/3d-motion-mag
from typing import Optional


def parse_optfloat(val, default_val=None) -> Optional[float]:
    if val == "None" or val is None:
        return default_val
    return float(val)


def parse_optint(val, default_val=None) -> Optional[int]:
    if val == "None" or val is None:
        return default_val
    return int(val)

magnerf/utils/parse_args.py
3d-motion-magnification/3d-motion-mag
"""Accurate timer for CUDA code"""
import torch.cuda


class CudaTimer:
    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self.start, self.end = None, None
        self.timings = {}
        if self.enabled:
            self.start = torch.cuda.Event(enable_timing=True)
            self.end = torch.cuda.Event(enable_timing=True)
            self.prev_time_gpu = self.start.record()

    def reset(self):
        if self.enabled:
            self.start = torch.cuda.Event(enable_timing=True)
            self.end = torch.cuda.Event(enable_timing=True)
            self.prev_time_gpu = self.start.record()

    def check(self, name=None):
        if self.enabled:
            self.end.record()
            torch.cuda.synchronize()
            gpu_time = self.start.elapsed_time(self.end)
            self.timings[name] = gpu_time

            self.prev_time_gpu = self.start.record()

magnerf/utils/timer.py
3d-motion-magnification/3d-motion-mag