SIGGRAPH Asia 2023: Code for "Im4D: High-Fidelity and Real-Time Novel View Synthesis for Dynamic Scenes"
from .yacs import CfgNode as CN
import argparse
import os
from os.path import join
import numpy as np
from . import yacs
import yaml
import sys

import logging
import colorlog
formatter = colorlog.ColoredFormatter(
    "%(log_color)s%(asctime)s - %(levelname)s - %(message)s",
    log_colors={
        'DEBUG': 'yellow',   
        'INFO': 'white',    
        'WARNING': 'red',
        'ERROR': 'red,bg_white',     
        'CRITICAL': 'red,bg_white', 
    }
)
logger = logging.getLogger(__name__)


def define_basic_cfg():
    cfg = CN()
    cfg.machine = 'default'
    cfg.train_mode = sys.argv[0] == 'train_net.py'
    cfg.debug = False
    cfg.workspace = os.environ['workspace']
    logger.debug('workspace: ' + cfg.workspace)
    
    cfg.save_result = False
    cfg.clear_result = False
    cfg.save_tag = 'default'
    cfg.write_video = False
    cfg.depth_method = 'expected'
    # module
    cfg.train_dataset_module = 'lib.datasets.dtu.neus'
    cfg.test_dataset_module = 'lib.datasets.dtu.neus'
    cfg.network_module = 'lib.neworks.neus.neus'
    cfg.loss_module = 'lib.train.losses.neus'
    cfg.evaluator_module = 'lib.evaluators.neus'
    
    # experiment name
    cfg.exp_name = 'gitbranch_hello'
    cfg.exp_name_tag = ''
    cfg.grid_tag = 'default'
    cfg.pretrain = ''
    
    # training network
    cfg.distributed = False
    cfg.train_fp16 = False
    cfg.eval_fp16 = False
    
    # task
    cfg.task = 'hello'
    cfg.resume = True
    cfg.sigma_thresh = 5.
    cfg.fast_render = False
    
    # epoch
    cfg.ep_iter = -1
    cfg.save_ep = 1
    cfg.save_latest_ep = 1000
    cfg.eval_ep = 1
    cfg.log_interval = 1
    
    cfg.task_arg = CN()
    cfg.vr_weight_thresh = 0.0
    # -----------------------------------------------------------------------------
    # train
    # -----------------------------------------------------------------------------
    cfg.train = CN()
    cfg.train.epoch = 10000
    cfg.train.num_workers = 8
    cfg.train.collator = 'default'
    cfg.train.batch_sampler = 'default'
    cfg.train.sampler_meta = CN({})
    cfg.train.shuffle = True
    cfg.train.eps = 1e-8
    
    # use adam as default
    cfg.train.optim = 'adam'
    cfg.train.lr = 5e-4
    cfg.train.weight_decay = 0.
    cfg.train.scheduler = CN({'type': 'multi_step', 'milestones': [
                             80, 120, 200, 240], 'gamma': 0.5})
    cfg.train.batch_size = 4
    
    # test
    cfg.test = CN()
    cfg.test.batch_size = 1
    cfg.test.collator = 'default'
    cfg.test.epoch = -1
    cfg.test.num_workers = 0
    cfg.test.batch_sampler = 'default'
    cfg.test.sampler_meta = CN({})
    
    # trained model
    cfg.trained_model_dir = os.path.join(os.environ['workspace'], 'trained_model')
    cfg.clean_tag = 'debug'
    
    # recorder
    cfg.record_dir = os.path.join(os.environ['workspace'], 'record')
    
    # result
    cfg.result_dir = os.path.join(os.environ['workspace'], 'result')
    
    # evaluation
    cfg.skip_eval = False
    cfg.fix_random = False
    return cfg
    

def parse_cfg(cfg, args):
    if len(cfg.exp_name_tag) != 0: cfg.exp_name += ('_' + cfg.exp_name_tag)
    cfg.exp_name = cfg.exp_name.replace('FILENAME', os.path.basename(args.cfg_file).split('.')[0])
    cfg.exp_name = cfg.exp_name.replace('GITBRANCH', os.popen('git describe --all').readline().strip()[6:])
    cfg.exp_name = cfg.exp_name.replace('GITCOMMIT', os.popen('git describe --tags --always').readline().strip())
    cfg.exp_name = cfg.exp_name.replace('TODAY', os.popen('date +%Y-%m-%d').readline().strip())
    
    cfg.trained_model_dir = os.path.join(cfg.trained_model_dir, cfg.scene, cfg.task, cfg.exp_name)
    cfg.record_dir = os.path.join(cfg.record_dir, cfg.scene, cfg.task, cfg.exp_name)
    cfg.grid_dir = os.path.join(cfg.result_dir, cfg.scene, cfg.task, 'grid', cfg.grid_tag)
    cfg.result_dir = os.path.join(cfg.result_dir, cfg.scene, cfg.task, cfg.exp_name, cfg.save_tag)
    cfg.local_rank = int(os.environ['LOCAL_RANK']) if 'LOCAL_RANK' in os.environ else 0
    modules = [key for key in cfg if '_module' in key]
    for module in modules:
        logger.info(module + ' ' + cfg[module])
        cfg[module.replace('_module', '_path')] = cfg[module].replace('.', '/') + '.py'
        
def load_yaml_file(file_path):
    with open(file_path, 'r') as file:
        return yaml.safe_load(file)
    
def merge_cfg(cfg_file, cfg, loaded_cfg_files):
    with open(cfg_file, 'r') as f:
        current_cfg = yacs.load_cfg(f)
    if 'parent_cfg' in current_cfg.keys():
        cfg = merge_cfg(current_cfg.parent_cfg, cfg, loaded_cfg_files)
    if 'configs' in current_cfg.keys():
        for cfg_file_ in current_cfg.configs:
            cfg = merge_cfg(cfg_file_, cfg, loaded_cfg_files)
    cfg.merge_from_other_cfg(current_cfg)
    loaded_cfg_files.append(cfg_file)
    return cfg

def list_to_str(inputs, split='\n    '):
    ret = '['
    for item in inputs: ret += item + split
    return ret + ']'

def make_cfg(args):
    if args.cfg_file[:8] == 'configs/':
        cfg = define_basic_cfg()
        cfg.log_level = os.environ['LOG_LEVEL'] if ('LOG_LEVEL' in os.environ and os.environ['LOG_LEVEL'] in ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']) else args.log_level
        loaded_cfg_files = []
        cfg_ = merge_cfg(args.cfg_file, cfg, loaded_cfg_files)
        if args.configs is not None:
            for config in args.configs: merge_cfg(config, cfg_, loaded_cfg_files)
        logger.debug('Loaded config files (in order):')
        logger.debug(list_to_str(loaded_cfg_files))
        try: index = args.opts.index('other_opts'); cfg_.merge_from_list(args.opts[:index])
        except: cfg_.merge_from_list(args.opts)
        parse_cfg(cfg_, args)
        os.system('mkdir -p {}'.format(cfg_.result_dir))
    else:
        cfg = yaml.safe_load(open(args.cfg_file))
    if cfg.local_rank == 0:
        file_handler = logging.FileHandler(join(cfg_.result_dir, 'app.log'), mode='a' if cfg_.resume else 'w')
        file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        logger.addHandler(file_handler)
        file_handler.setFormatter(file_formatter)
    
    # LOGGING
    logger.info('=' * 30 + ' CONFIG ' + '=' * 30)
    logger.info('EXP NAME: ' + cfg.exp_name)
    if cfg.train_mode:
        logger.info('Training mode')
        logger.info('Training with float16: ' + str(cfg_.train_fp16))
        logger.info('Number of images: ' + str(cfg_.train_dataset.imgs_per_batch))
        logger.info('Number of pixels: ' + str(cfg_.num_pixels))
    else:
        logger.info('Testing mode')
        logger.info('Evaluation with float16: ' + str(cfg_.eval_fp16))
    logger.info('=' * 30 + ' CONFIG END ' + '=' * 30)
    return cfg_

def get_log_level(level_str):
    level_str = level_str.upper()
    level_map = {
        'DEBUG': logging.DEBUG,
        'INFO': logging.INFO,
        'WARNING': logging.WARNING,
        'ERROR': logging.ERROR,
        'CRITICAL': logging.CRITICAL
    }
    return level_map.get(level_str, logging.INFO)

parser = argparse.ArgumentParser()
parser.add_argument("--cfg_file", default="configs/default.yaml", type=str)
parser.add_argument("--type", type=str, default="")
parser.add_argument("--configs", type=str, action='append')
parser.add_argument('--log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the log level; It will firstly read LOG_LEVEL from environment variables', default='DEBUG')
parser.add_argument("opts", default=None, nargs=argparse.REMAINDER)
args = parser.parse_args()
logger.setLevel(get_log_level(os.environ['LOG_LEVEL'] if ('LOG_LEVEL' in os.environ and os.environ['LOG_LEVEL'] in ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']) else args.log_level))
console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
cfg = make_cfg(args)
lib/config/config.py
zju3dv/im4d
# Copyright (c) 2018-present, Facebook, Inc.
#
# 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.
##############################################################################

"""YACS -- Yet Another Configuration System is designed to be a simple
configuration management system for academic and industrial research
projects.

See README.md for usage and examples.
"""

import copy
import io
import logging
import os
from ast import literal_eval

import yaml


# Flag for py2 and py3 compatibility to use when separate code paths are necessary
# When _PY2 is False, we assume Python 3 is in use
_PY2 = False

# Filename extensions for loading configs from files
_YAML_EXTS = {"", ".yaml", ".yml"}
_PY_EXTS = {".py"}

# py2 and py3 compatibility for checking file object type
# We simply use this to infer py2 vs py3
try:
    _FILE_TYPES = (file, io.IOBase)
    _PY2 = True
except NameError:
    _FILE_TYPES = (io.IOBase,)

# CfgNodes can only contain a limited set of valid types
_VALID_TYPES = {tuple, list, str, int, float, bool}
# py2 allow for str and unicode
if _PY2:
    _VALID_TYPES = _VALID_TYPES.union({unicode})  # noqa: F821

# Utilities for importing modules from file paths
if _PY2:
    # imp is available in both py2 and py3 for now, but is deprecated in py3
    import imp
else:
    import importlib.util

logger = logging.getLogger(__name__)


class CfgNode(dict):
    """
    CfgNode represents an internal node in the configuration tree. It's a simple
    dict-like container that allows for attribute-based access to keys.
    """

    IMMUTABLE = "__immutable__"
    DEPRECATED_KEYS = "__deprecated_keys__"
    RENAMED_KEYS = "__renamed_keys__"

    def __init__(self, init_dict=None, key_list=None):
        # Recursively convert nested dictionaries in init_dict into CfgNodes
        init_dict = {} if init_dict is None else init_dict
        key_list = [] if key_list is None else key_list
        for k, v in init_dict.items():
            if type(v) is dict:
                # Convert dict to CfgNode
                init_dict[k] = CfgNode(v, key_list=key_list + [k])
            else:
                # Check for valid leaf type or nested CfgNode
                _assert_with_logging(
                    _valid_type(v, allow_cfg_node=True),
                    "Key {} with value {} is not a valid type; valid types: {}".format(
                        ".".join(key_list + [k]), type(v), _VALID_TYPES
                    ),
                )
        super(CfgNode, self).__init__(init_dict)
        # Manage if the CfgNode is frozen or not
        self.__dict__[CfgNode.IMMUTABLE] = False
        # Deprecated options
        # If an option is removed from the code and you don't want to break existing
        # yaml configs, you can add the full config key as a string to the set below.
        self.__dict__[CfgNode.DEPRECATED_KEYS] = set()
        # Renamed options
        # If you rename a config option, record the mapping from the old name to the new
        # name in the dictionary below. Optionally, if the type also changed, you can
        # make the value a tuple that specifies first the renamed key and then
        # instructions for how to edit the config file.
        self.__dict__[CfgNode.RENAMED_KEYS] = {
            # 'EXAMPLE.OLD.KEY': 'EXAMPLE.NEW.KEY',  # Dummy example to follow
            # 'EXAMPLE.OLD.KEY': (                   # A more complex example to follow
            #     'EXAMPLE.NEW.KEY',
            #     "Also convert to a tuple, e.g., 'foo' -> ('foo',) or "
            #     + "'foo:bar' -> ('foo', 'bar')"
            # ),
        }

    def __getattr__(self, name):
        if name in self:
            return self[name]
        else:
            raise AttributeError(name)

    def __setattr__(self, name, value):
        if self.is_frozen():
            raise AttributeError(
                "Attempted to set {} to {}, but CfgNode is immutable".format(
                    name, value
                )
            )

        _assert_with_logging(
            name not in self.__dict__,
            "Invalid attempt to modify internal CfgNode state: {}".format(name),
        )
        _assert_with_logging(
            _valid_type(value, allow_cfg_node=True),
            "Invalid type {} for key {}; valid types = {}".format(
                type(value), name, _VALID_TYPES
            ),
        )

        self[name] = value

    def __str__(self):
        def _indent(s_, num_spaces):
            s = s_.split("\n")
            if len(s) == 1:
                return s_
            first = s.pop(0)
            s = [(num_spaces * " ") + line for line in s]
            s = "\n".join(s)
            s = first + "\n" + s
            return s

        r = ""
        s = []
        for k, v in sorted(self.items()):
            seperator = "\n" if isinstance(v, CfgNode) else " "
            attr_str = "{}:{}{}".format(str(k), seperator, str(v))
            attr_str = _indent(attr_str, 2)
            s.append(attr_str)
        r += "\n".join(s)
        return r

    def __repr__(self):
        return "{}({})".format(self.__class__.__name__, super(CfgNode, self).__repr__())

    def dump(self):
        """Dump to a string."""
        self_as_dict = _to_dict(self)
        return yaml.safe_dump(self_as_dict)

    def merge_from_file(self, cfg_filename):
        """Load a yaml config file and merge it this CfgNode."""
        with open(cfg_filename, "r") as f:
            cfg = load_cfg(f)
        self.merge_from_other_cfg(cfg)

    def merge_from_other_cfg(self, cfg_other):
        """Merge `cfg_other` into this CfgNode."""
        _merge_a_into_b(cfg_other, self, self, [])

    def merge_from_list(self, cfg_list):
        """Merge config (keys, values) in a list (e.g., from command line) into
        this CfgNode. For example, `cfg_list = ['FOO.BAR', 0.5]`.
        """
        _assert_with_logging(
            len(cfg_list) % 2 == 0,
            "Override list has odd length: {}; it must be a list of pairs".format(
                cfg_list
            ),
        )
        root = self
        for full_key, v in zip(cfg_list[0::2], cfg_list[1::2]):
            if root.key_is_deprecated(full_key):
                continue
            if root.key_is_renamed(full_key):
                root.raise_key_rename_error(full_key)
            key_list = full_key.split(".")
            d = self
            for subkey in key_list[:-1]:
                _assert_with_logging(
                    subkey in d, "Non-existent key: {}".format(full_key)
                )
                d = d[subkey]
            subkey = key_list[-1]
            _assert_with_logging(subkey in d, "Non-existent key: {}".format(full_key))
            value = _decode_cfg_value(v)
            value = _check_and_coerce_cfg_value_type(value, d[subkey], subkey, full_key)
            d[subkey] = value

    def freeze(self):
        """Make this CfgNode and all of its children immutable."""
        self._immutable(True)

    def defrost(self):
        """Make this CfgNode and all of its children mutable."""
        self._immutable(False)

    def is_frozen(self):
        """Return mutability."""
        return self.__dict__[CfgNode.IMMUTABLE]

    def _immutable(self, is_immutable):
        """Set immutability to is_immutable and recursively apply the setting
        to all nested CfgNodes.
        """
        self.__dict__[CfgNode.IMMUTABLE] = is_immutable
        # Recursively set immutable state
        for v in self.__dict__.values():
            if isinstance(v, CfgNode):
                v._immutable(is_immutable)
        for v in self.values():
            if isinstance(v, CfgNode):
                v._immutable(is_immutable)

    def clone(self):
        """Recursively copy this CfgNode."""
        return copy.deepcopy(self)

    def register_deprecated_key(self, key):
        """Register key (e.g. `FOO.BAR`) a deprecated option. When merging deprecated
        keys a warning is generated and the key is ignored.
        """
        _assert_with_logging(
            key not in self.__dict__[CfgNode.DEPRECATED_KEYS],
            "key {} is already registered as a deprecated key".format(key),
        )
        self.__dict__[CfgNode.DEPRECATED_KEYS].add(key)

    def register_renamed_key(self, old_name, new_name, message=None):
        """Register a key as having been renamed from `old_name` to `new_name`.
        When merging a renamed key, an exception is thrown alerting to user to
        the fact that the key has been renamed.
        """
        _assert_with_logging(
            old_name not in self.__dict__[CfgNode.RENAMED_KEYS],
            "key {} is already registered as a renamed cfg key".format(old_name),
        )
        value = new_name
        if message:
            value = (new_name, message)
        self.__dict__[CfgNode.RENAMED_KEYS][old_name] = value

    def key_is_deprecated(self, full_key):
        """Test if a key is deprecated."""
        if full_key in self.__dict__[CfgNode.DEPRECATED_KEYS]:
            logger.warning("Deprecated config key (ignoring): {}".format(full_key))
            return True
        return False

    def key_is_renamed(self, full_key):
        """Test if a key is renamed."""
        return full_key in self.__dict__[CfgNode.RENAMED_KEYS]

    def raise_key_rename_error(self, full_key):
        new_key = self.__dict__[CfgNode.RENAMED_KEYS][full_key]
        if isinstance(new_key, tuple):
            msg = " Note: " + new_key[1]
            new_key = new_key[0]
        else:
            msg = ""
        raise KeyError(
            "Key {} was renamed to {}; please update your config.{}".format(
                full_key, new_key, msg
            )
        )


def load_cfg(cfg_file_obj_or_str):
    """Load a cfg. Supports loading from:
        - A file object backed by a YAML file
        - A file object backed by a Python source file that exports an attribute
          "cfg" that is either a dict or a CfgNode
        - A string that can be parsed as valid YAML
    """
    _assert_with_logging(
        isinstance(cfg_file_obj_or_str, _FILE_TYPES + (str,)),
        "Expected first argument to be of type {} or {}, but it was {}".format(
            _FILE_TYPES, str, type(cfg_file_obj_or_str)
        ),
    )
    if isinstance(cfg_file_obj_or_str, str):
        return _load_cfg_from_yaml_str(cfg_file_obj_or_str)
    elif isinstance(cfg_file_obj_or_str, _FILE_TYPES):
        return _load_cfg_from_file(cfg_file_obj_or_str)
    else:
        raise NotImplementedError("Impossible to reach here (unless there's a bug)")


def _load_cfg_from_file(file_obj):
    """Load a config from a YAML file or a Python source file."""
    _, file_extension = os.path.splitext(file_obj.name)
    if file_extension in _YAML_EXTS:
        return _load_cfg_from_yaml_str(file_obj.read())
    elif file_extension in _PY_EXTS:
        return _load_cfg_py_source(file_obj.name)
    else:
        raise Exception(
            "Attempt to load from an unsupported file type {}; "
            "only {} are supported".format(file_obj, _YAML_EXTS.union(_PY_EXTS))
        )


def _load_cfg_from_yaml_str(str_obj):
    """Load a config from a YAML string encoding."""
    cfg_as_dict = yaml.safe_load(str_obj)
    return CfgNode(cfg_as_dict)


def _load_cfg_py_source(filename):
    """Load a config from a Python source file."""
    module = _load_module_from_file("yacs.config.override", filename)
    _assert_with_logging(
        hasattr(module, "cfg"),
        "Python module from file {} must have 'cfg' attr".format(filename),
    )
    VALID_ATTR_TYPES = {dict, CfgNode}
    _assert_with_logging(
        type(module.cfg) in VALID_ATTR_TYPES,
        "Imported module 'cfg' attr must be in {} but is {} instead".format(
            VALID_ATTR_TYPES, type(module.cfg)
        ),
    )
    if type(module.cfg) is dict:
        return CfgNode(module.cfg)
    else:
        return module.cfg


def _to_dict(cfg_node):
    """Recursively convert all CfgNode objects to dict objects."""

    def convert_to_dict(cfg_node, key_list):
        if not isinstance(cfg_node, CfgNode):
            _assert_with_logging(
                _valid_type(cfg_node),
                "Key {} with value {} is not a valid type; valid types: {}".format(
                    ".".join(key_list), type(cfg_node), _VALID_TYPES
                ),
            )
            return cfg_node
        else:
            cfg_dict = dict(cfg_node)
            for k, v in cfg_dict.items():
                cfg_dict[k] = convert_to_dict(v, key_list + [k])
            return cfg_dict

    return convert_to_dict(cfg_node, [])


def _valid_type(value, allow_cfg_node=False):
    return (type(value) in _VALID_TYPES) or (allow_cfg_node and type(value) == CfgNode)


def _merge_a_into_b(a, b, root, key_list):
    """Merge config dictionary a into config dictionary b, clobbering the
    options in b whenever they are also specified in a.
    """
    _assert_with_logging(
        isinstance(a, CfgNode),
        "`a` (cur type {}) must be an instance of {}".format(type(a), CfgNode),
    )
    _assert_with_logging(
        isinstance(b, CfgNode),
        "`b` (cur type {}) must be an instance of {}".format(type(b), CfgNode),
    )

    for k, v_ in a.items():
        full_key = ".".join(key_list + [k])
        # a must specify keys that are in b
        if k not in b:
            if root.key_is_deprecated(full_key):
                continue
            elif root.key_is_renamed(full_key):
                root.raise_key_rename_error(full_key)
            else:
                v = copy.deepcopy(v_)
                v = _decode_cfg_value(v)
                b.update({k: v})
        else:
            v = copy.deepcopy(v_)
            v = _decode_cfg_value(v)
            v = _check_and_coerce_cfg_value_type(v, b[k], k, full_key)

        # Recursively merge dicts
        if isinstance(v, CfgNode):
            try:
                _merge_a_into_b(v, b[k], root, key_list + [k])
            except BaseException:
                raise
        else:
            b[k] = v


def _decode_cfg_value(v):
    """Decodes a raw config value (e.g., from a yaml config files or command
    line argument) into a Python object.
    """
    # Configs parsed from raw yaml will contain dictionary keys that need to be
    # converted to CfgNode objects
    if isinstance(v, dict):
        return CfgNode(v)
    # All remaining processing is only applied to strings
    if not isinstance(v, str):
        return v
    # Try to interpret `v` as a:
    #   string, number, tuple, list, dict, boolean, or None
    try:
        v = literal_eval(v)
    # The following two excepts allow v to pass through when it represents a
    # string.
    #
    # Longer explanation:
    # The type of v is always a string (before calling literal_eval), but
    # sometimes it *represents* a string and other times a data structure, like
    # a list. In the case that v represents a string, what we got back from the
    # yaml parser is 'foo' *without quotes* (so, not '"foo"'). literal_eval is
    # ok with '"foo"', but will raise a ValueError if given 'foo'. In other
    # cases, like paths (v = 'foo/bar' and not v = '"foo/bar"'), literal_eval
    # will raise a SyntaxError.
    except ValueError:
        pass
    except SyntaxError:
        pass
    return v


def _check_and_coerce_cfg_value_type(replacement, original, key, full_key):
    """Checks that `replacement`, which is intended to replace `original` is of
    the right type. The type is correct if it matches exactly or is one of a few
    cases in which the type can be easily coerced.
    """
    original_type = type(original)
    replacement_type = type(replacement)

    # The types must match (with some exceptions)
    if replacement_type == original_type:
        return replacement

    # Cast replacement from from_type to to_type if the replacement and original
    # types match from_type and to_type
    def conditional_cast(from_type, to_type):
        if replacement_type == from_type and original_type == to_type:
            return True, to_type(replacement)
        else:
            return False, None

    # Conditionally casts
    # list <-> tuple
    casts = [(tuple, list), (list, tuple)]
    # For py2: allow converting from str (bytes) to a unicode string
    try:
        casts.append((str, unicode))  # noqa: F821
    except Exception:
        pass

    for (from_type, to_type) in casts:
        converted, converted_value = conditional_cast(from_type, to_type)
        if converted:
            return converted_value

    raise ValueError(
        "Type mismatch ({} vs. {}) with values ({} vs. {}) for config "
        "key: {}".format(
            original_type, replacement_type, original, replacement, full_key
        )
    )


def _assert_with_logging(cond, msg):
    if not cond:
        logger.debug(msg)
    assert cond, msg


def _load_module_from_file(name, filename):
    if _PY2:
        module = imp.load_source(name, filename)
    else:
        spec = importlib.util.spec_from_file_location(name, filename)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
    return module

lib/config/yacs.py
zju3dv/im4d
from torch.utils.data.dataloader import default_collate
import torch
import numpy as np
from lib.config import cfg


def none_collate(batch):
    return batch[0]

_collators = {
    'none': none_collate,
}

def make_collator(cfg, is_train):
    collator = cfg.train.collator if is_train else cfg.test.collator
    if collator in _collators:
        return _collators[collator]
    else:
        return default_collate

lib/datasets/collate_batch.py
zju3dv/im4d
from . import samplers
import torch
import torch.utils.data
import imp
import os
from .collate_batch import make_collator
import numpy as np
import time
from lib.config.config import cfg
from torch.utils.data import DataLoader, ConcatDataset


# torch.multiprocessing.set_sharing_strategy('file_system')

def _dataset_factory(is_train, is_val):
    if is_val:
        module = cfg.val_dataset_module
        path = cfg.val_dataset_path
    elif is_train:
        module = cfg.train_dataset_module
        path = cfg.train_dataset_path
    else:
        module = cfg.test_dataset_module
        path = cfg.test_dataset_path
    dataset = imp.load_source(module, path).Dataset
    return dataset


def make_dataset(cfg, is_train=True):
    if is_train:
        args = cfg.train_dataset
        module = cfg.train_dataset_module
        path = cfg.train_dataset_path
    else:
        args = cfg.test_dataset
        module = cfg.test_dataset_module
        path = cfg.test_dataset_path
    dataset = imp.load_source(module, path).Dataset
    dataset = dataset(**args)
    return dataset


def make_data_sampler(dataset, shuffle, is_distributed):
    if is_distributed:
        return samplers.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        sampler = torch.utils.data.sampler.RandomSampler(dataset)
    else:
        sampler = torch.utils.data.sampler.SequentialSampler(dataset)
    return sampler


def make_batch_data_sampler(cfg, sampler, batch_size, drop_last, max_iter,
                            is_train):
    if is_train:
        batch_sampler = cfg.train.batch_sampler
        sampler_meta = cfg.train.sampler_meta
    else:
        batch_sampler = cfg.test.batch_sampler
        sampler_meta = cfg.test.sampler_meta

    if batch_sampler == 'default':
        batch_sampler = torch.utils.data.sampler.BatchSampler(
            sampler, batch_size, drop_last)
    elif batch_sampler == 'image_size':
        batch_sampler = samplers.ImageSizeBatchSampler(sampler, batch_size,
                                                       drop_last, sampler_meta)

    if max_iter != -1:
        batch_sampler = samplers.IterationBasedBatchSampler(
            batch_sampler, max_iter)
    return batch_sampler


def worker_init_fn(worker_id):
    np.random.seed(worker_id + (int(round(time.time() * 1000) % (2**16))))


def make_data_loader(cfg, is_train=True, is_distributed=False, max_iter=-1):
    if is_train:
        batch_size = cfg.train.batch_size
        # shuffle = True
        shuffle = cfg.train.shuffle
        drop_last = False
    else:
        batch_size = cfg.test.batch_size
        shuffle = True if is_distributed else False
        drop_last = False

    dataset = make_dataset(cfg, is_train)
    sampler = make_data_sampler(dataset, shuffle, is_distributed)
    batch_sampler = make_batch_data_sampler(cfg, sampler, batch_size,
                                            drop_last, max_iter, is_train)
    num_workers = cfg.train.num_workers if is_train else cfg.test.num_workers
    collator = make_collator(cfg, is_train)
    data_loader = DataLoader(dataset,
                            batch_sampler=batch_sampler,
                            num_workers=num_workers,
                            collate_fn=collator,
                            pin_memory=True,
                            worker_init_fn=worker_init_fn)

    return data_loader

lib/datasets/make_dataset.py
zju3dv/im4d
from torch.utils.data.sampler import Sampler
from torch.utils.data.sampler import BatchSampler
import numpy as np
import torch
import math
import torch.distributed as dist
from lib.config import cfg


class ImageSizeBatchSampler(Sampler):
    def __init__(self, sampler, batch_size, drop_last, sampler_meta):
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.strategy = sampler_meta.strategy
        self.hmin, self.wmin = sampler_meta.min_hw
        self.hmax, self.wmax = sampler_meta.max_hw
        self.divisor = 32
        if cfg.fix_random:
            np.random.seed(0)

    def generate_height_width(self):
        if self.strategy == 'origin':
            return -1, -1
        h = np.random.randint(self.hmin, self.hmax + 1)
        w = np.random.randint(self.wmin, self.wmax + 1)
        h = (h | (self.divisor - 1)) + 1
        w = (w | (self.divisor - 1)) + 1
        return h, w

    def __iter__(self):
        batch = []
        h, w = self.generate_height_width()
        for idx in self.sampler:
            batch.append((idx, h, w))
            if len(batch) == self.batch_size:
                h, w = self.generate_height_width()
                yield batch
                batch = []
        if len(batch) > 0 and not self.drop_last:
            yield batch

    def __len__(self):
        if self.drop_last:
            return len(self.sampler) // self.batch_size
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size


class IterationBasedBatchSampler(BatchSampler):
    """
    Wraps a BatchSampler, resampling from it until
    a specified number of iterations have been sampled
    """

    def __init__(self, batch_sampler, num_iterations, start_iter=0):
        self.batch_sampler = batch_sampler
        self.sampler = self.batch_sampler.sampler
        self.num_iterations = num_iterations
        self.start_iter = start_iter

    def __iter__(self):
        iteration = self.start_iter
        while iteration <= self.num_iterations:
            for batch in self.batch_sampler:
                iteration += 1
                if iteration > self.num_iterations:
                    break
                yield batch

    def __len__(self):
        return self.num_iterations


class DistributedSampler(Sampler):
    """Sampler that restricts data loading to a subset of the dataset.
    It is especially useful in conjunction with
    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
    process can pass a DistributedSampler instance as a DataLoader sampler,
    and load a subset of the original dataset that is exclusive to it.
    .. note::
        Dataset is assumed to be of constant size.
    Arguments:
        dataset: Dataset used for sampling.
        num_replicas (optional): Number of processes participating in
            distributed training.
        rank (optional): Rank of the current process within num_replicas.
    """

    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle

    def __iter__(self):
        if self.shuffle:
            # deterministically shuffle based on epoch
            g = torch.Generator()
            g.manual_seed(self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        else:
            indices = torch.arange(len(self.dataset)).tolist()

        # add extra samples to make it evenly divisible
        indices += indices[: (self.total_size - len(indices))]
        assert len(indices) == self.total_size

        # subsample
        offset = self.num_samples * self.rank
        indices = indices[offset:offset+self.num_samples]
        assert len(indices) == self.num_samples

        return iter(indices)

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):
        self.epoch = epoch

lib/datasets/samplers.py
zju3dv/im4d
import numpy as np
import os
from lib.utils.base_utils import get_sample
from lib.utils.data_utils import load_image_bytes, load_image_from_bytes, load_resize_undist_im_bytes, load_img, load_resize_undist_calib_im_bytes, get_ratios, transform_points_opengl_to_opencv, get_near_far
from lib.utils.parallel_utils import parallel_execution
from lib.config import cfg, logger
import imageio
import cv2
from tqdm import tqdm
from lib.config import cfg
from lib.utils import data_utils
from lib.utils import base_utils
from lib.utils.easyvv import easy_utils
import trimesh
from lib.utils.base_utils import get_bound_corners
from os.path import join
from termcolor import colored
from torch.utils.data.dataloader import default_collate
from lib.utils.data_utils import parse_cameras_genebody, save_img
from lib.utils.pixel_sampling.pixel_sampler import PixelSampler
from lib.utils import rend_utils
from lib.config.yacs import CfgNode as CN

timer = base_utils.perf_timer(sync_cuda=False)

class Dataset:
    def __init__(self, **kwargs):
        super(Dataset, self).__init__()
        # Prepare camera intrinsic and extrinsic parameters
        # Prepare the metadata for rendering
        # Prepare bounding box or point cloud for near-far determination
        # Precompute some things in advance
        self.cfg = CN(kwargs)
        self.prepare_metainfo()
        self.prepare_camera() # prepare camera intrinsics and extrinsics
        self.prepare_render_meta() # prepare render meta for the rendered images
        if self.preload_data: self.preload_data_func()
        self.load_bboxes() # load bboxes
        self.load_pointcloud()
        self.precompute() # precompute some things in advance, e.g., near_far etc.

    def prepare_metainfo(self):
        self.data_root = os.path.join(cfg.workspace, self.cfg.data_root, cfg.scene)
        self.split = self.cfg.split
        self.start_frame = self.cfg.frame_sample[0]
        self.frame_len = self.cfg.frame_sample[1] - self.cfg.frame_sample[0]
        self.crop_h_w = self.cfg.get('crop_h_w', [-1, -1])
        self.imgs_per_batch = self.cfg.get('imgs_per_batch', 1)
        self.shift_pixel = self.cfg.get('shift_pixel', False)
        self.preload_data = self.cfg.get('preload_data', False)
        self.have_msk = self.cfg.get('msk_dir', 'none') != 'none'
        self.ignore_dist_k3 = self.cfg.get('ignore_dist_k3', False)
        self.img_dir = self.cfg.get('img_dir', 'images')
        self.msk_dir = self.cfg.get('msk_dir', 'masks')
        self.img_frame_format = self.cfg.get('img_frame_format', '{:06d}.jpg')
        self.msk_frame_format = self.cfg.get('msk_frame_format', '{:06d}.png')
        self.scene_type = self.cfg.get('scene_type', 'bbox_scene')
        self.near_far = np.array(self.cfg.get('near_far', [0.1, 100.]))
        self.bbox_type = self.cfg.get('bbox_type', 'NONE')
        self.render_path = cfg.get('render_path', False)
        self.path_type = self.cfg.get('path_type', 'NERFSTUDIO')
        self.path_name = self.cfg.get('path_name', '2023-07-13_220007.json')

    def prepare_camera(self):
        intri_path = join(self.data_root, self.cfg.get('intri_file', 'intri.yml'))
        extri_path = join(self.data_root, self.cfg.get('extri_file', 'extri.yml'))
        logger.debug('Loading camera from: {}, {}'.format(intri_path, extri_path))
        cams = easy_utils.read_camera(intri_path, extri_path, read_rvec=self.cfg.get('read_rvec', True))
        self.basenames = sorted([k for k in cams])
        cam_len = len(self.basenames)
        self.ixts = np.array([cams[cam_id]['K'] for cam_id in self.basenames]).reshape(cam_len, 3, 3).astype(np.float32)
        self.orig_ixts = self.ixts.copy()
        exts = np.array([cams[cam_id]['RT'] for cam_id in self.basenames]).reshape(cam_len, 3, 4).astype(np.float32)
        exts_ones = np.zeros_like(exts[:, :1, :])
        exts_ones[..., 3] = 1.
        self.exts = np.concatenate([exts, exts_ones], axis=1)
        transform = np.array(self.cfg.get('transform', np.eye(4)))
        self.exts = (self.exts @ np.linalg.inv(transform)).astype(np.float32)  # align z to up
        self.exts_inv = np.linalg.inv(self.exts).astype(np.float32)
        self.Ds = np.array([cams[cam_id]['D'] for cam_id in self.basenames]).reshape(cam_len, 5).astype(np.float32)
        if self.ignore_dist_k3: self.Ds[..., 4] = 0.

        # process ixts according to crop and resize
        special_views = self.cfg.get('special_views', [])
        special_views = get_sample(special_views)
        self.rs = [self.cfg.get('special_resize_ratio') if i in special_views else self.cfg.get('resize_ratio') for i in range(len(self.ixts))]
        self.ixts[:, :2] *= np.array(self.rs)[:, None, None]
        if self.crop_h_w[0] != -1:
            frame_id = self.start_frame
            img_paths = [join(self.data_root, self.img_dir, basename, self.img_frame_format.format(frame_id)) for basename in self.basenames]
            imgs = parallel_execution(img_paths, action=load_img, num_processes=32,print_progress=True, sequential=False, async_return=False, desc='loading images for compute ixts')
            tar_h, tar_w = self.crop_h_w
            new_ixts = []
            for idx, (img, ixt) in enumerate(zip(imgs, self.ixts)):
                h, w = img.shape[:2]
                if self.rs[idx] != 1.: h, w = int(h*self.rs[idx]), int(w*self.rs[idx])
                crop_h, crop_w = (h - tar_h)//2, (w - tar_w)//2
                ixt[0, 2] -= crop_w
                ixt[1, 2] -= crop_h
                new_ixts.append(ixt)
            self.ixts = new_ixts
        self.ixts = np.array(self.ixts).astype(np.float32)
        self.ixts_inv = np.linalg.inv(self.ixts).astype(np.float32)

    def prepare_render_meta(self):
        # mem bank
        self.meta_imgs_path = []
        self.meta_msks_path = []
        self.meta_viewids = []
        # index mem bank: given frame_id, and view_id, return the index in meta_imgs_path
        self.meta_idx = {}
        # render_meta
        self.meta = []

        render_views = get_sample(self.cfg.get('render_view_sample'))
        input_views = get_sample(self.cfg.get('input_view_sample'))
        test_views = self.cfg.get('test_views', [-1])
        if -1 not in test_views:
            render_views = test_views
            input_views = [view for view in input_views if view not in render_views]
            if self.split == 'train': render_views = input_views
        logger.info('Dataset split: {}, render views: {}'.format(self.split, str(render_views)))
        self.get_metas(render_views, input_views)

    def load_pointcloud(self):
        if 'scene_pcd' not in self.cfg: self.point_cloud = None; return
        pcd_path = join(self.data_root, 'nerfstudio/exports/pcd/{}'.format(self.cfg.scene_pcd))
        if os.path.exists(pcd_path):
            pcd = np.asarray(trimesh.load(pcd_path).vertices)
        else:
            pcd_path = join(self.data_root, self.cfg.scene_pcd)
            pcd = np.asarray(trimesh.load(pcd_path).vertices)
        if 'nerfstudio' in pcd_path:
            logger.debug('Transform point cloud from opengl to opencv')
            pcd = transform_points_opengl_to_opencv(pcd)
        num_max_pcd = 10000
        if pcd.shape[0] > num_max_pcd:
            np.random.seed(0)
            rand_msk = np.random.random(pcd.shape[0])
            pcd = pcd[rand_msk < (num_max_pcd/pcd.shape[0])]
        self.point_cloud = pcd

    def precompute(self, **kwargs):
        MAX_IMAGESIZE = 5120
        H, W = 5120, 5120  # MAX_IMGAESIZE
        self.coords = np.stack(np.meshgrid(np.linspace(
            0, W-1, W), np.linspace(0, H-1, H)), -1).astype(np.int32)
        self.coords = np.concatenate(
            [self.coords, np.ones_like(self.coords[..., :1])], axis=-1)

    def load_bboxes(self):
        # priority:
        # 1. per_frame bbox
        # 2. dataset bbox
        # 3. global bbox
        # TODO: implement priority
        frames = self.cfg.frame_sample
        self.corners = {}
        self.bounds = {}

        transform = np.array(self.cfg.get('transform', np.eye(4)))
        for frame_id in tqdm(np.arange(*frames), desc='loading bboxs'):
            if self.bbox_type == 'RENBODY': bounds = data_utils.read_bbox_renbody(self.data_root, frame_id, transform, cfg.bound_padding)
            elif self.bbox_type == 'NHR': bounds = data_utils.read_bbox_nhr(self.data_root, frame_id, transform, cfg.bound_padding)
            elif self.bbox_type == 'GENERAL': import ipdb; ipdb.set_trace()
            else: bounds = np.asarray(cfg.bounds).astype(np.float32)
            self.bounds[frame_id] = bounds
            self.corners[frame_id] = np.array(get_bound_corners(self.bounds[frame_id])).astype(np.float32)

        self.bound = np.array([np.array([self.bounds[k] for k in self.bounds]).min(axis=(0, 1)),
                               np.array([self.bounds[k] for k in self.bounds]).max(axis=(0, 1))]).astype(np.float32)
        self.corner = get_bound_corners(self.bound).astype(np.float32)
        print('Dataset bound: ', self.bound)

    def preload_img(self, frame_id, view_id, **kwargs):
        img_path = join(self.data_root, self.img_dir, self.basenames[view_id], self.img_frame_format.format(frame_id))
        if img_path in self.meta_imgs_path:
            return
        if self.have_msk:
            msk_path = join(self.data_root, self.msk_dir, self.basenames[view_id], self.msk_frame_format.format(frame_id))
            self.meta_msks_path.append(msk_path)
        self.meta_imgs_path.append(img_path)
        self.meta_viewids.append(view_id)
        if frame_id not in self.meta_idx: self.meta_idx[frame_id] = {}
        self.meta_idx[frame_id][view_id] = len(self.meta_imgs_path) - 1

    def get_render_path_meta(self, input_views, **kwargs):
        if self.bbox_type == 'LLFF':
            path_c2ws = np.load(join(self.data_root, cfg.path_name))
            bottom = np.array([[0, 0, 0, 1.]])[None].repeat(cfg.num_circle, 0)
            path_c2ws = np.concatenate([path_c2ws, bottom], axis=1)
            path_w2cs = np.linalg.inv(path_c2ws)
        else:
            path_w2cs = rend_utils.create_center_radius(
                np.array(cfg.render_center),
                radius=cfg.render_radius,
                up=cfg.render_axis,
                ranges=cfg.render_view_range + [cfg.num_circle],
                angle_x=cfg.render_angle)  # path 1
            bottom = np.array([[0, 0, 0, 1.]])[None].repeat(cfg.num_circle, 0)
            path_w2cs = np.concatenate([path_w2cs, bottom], axis=1)
            path_w2cs = path_w2cs.astype(np.float32)
            path_c2ws = np.linalg.inv(path_w2cs)
        # DEBUG:
        # np.save('cams.npy', {'input': self.exts_inv, 'render': self.path_c2ws})
        # self.render_ids, render_id = [], 0

        self.render_exts = []
        self.render_ixts = []
        self.render_ids = []
        render_id = 0 # index
        for frame_id in tqdm(range(*cfg.render_frames), desc='loading render meta'):
            tar_ext = path_w2cs[render_id % cfg.num_circle]
            self.preload_img(frame_id, 0)
            self.meta += [(frame_id, -1)]
            self.render_exts.append(tar_ext.astype(np.float32))
            self.render_ixts.append(self.ixts[0].astype(np.float32))
            self.render_ids.append(render_id)
            render_id += 1

            if frame_id in cfg.render_stop_frame:
                for i in range(cfg.render_stop_time[cfg.render_stop_frame.index(frame_id)]):
                    tar_ext = path_w2cs[render_id % cfg.num_circle]
                    self.meta += [(frame_id, -1)]
                    self.render_exts.append(tar_ext.astype(np.float32))
                    self.render_ixts.append(self.ixts[0].astype(np.float32))
                    self.render_ids.append(render_id)
                    render_id += 1

        if len(cfg.render_ixts_start) >= 1:
            ixts_ratios = np.ones_like(np.array(self.render_ids)).astype(np.float32)
            prev = 1.
            for idx in range(len(cfg.render_ixts_start)):
                start = cfg.render_ixts_start[idx]
                end = cfg.render_ixts_end[idx]
                if idx >= 1: ixts_ratios[cfg.render_ixts_end[idx-1]:start] = prev
                ixts_ratios[start:end] = get_ratios(end-start, prev, cfg.render_ixts_ratio[idx], cfg.render_smooth)
                prev = cfg.render_ixts_ratio[idx]
            self.render_ixts = np.array(self.render_ixts).astype(np.float32)
            self.render_ixts[:, :2, :2] *= ixts_ratios[:, None, None]
            
        self.render_h_w = cfg.render_h_w
        self.render_exts = np.array(self.render_exts).astype(np.float32)
        self.render_ixts = np.array(self.render_ixts).astype(np.float32)
        self.render_exts_inv = np.linalg.inv(self.render_exts).astype(np.float32)
        self.render_ixts_inv = np.linalg.inv(self.render_ixts).astype(np.float32)

        # render_meta = {'render_exts': np.array(self.render_exts),
        #                'render_ixts': np.array(self.render_ixts),
        #                'frame_ids': [meta[0] for meta in self.meta],
        #                'src_views': [meta[2] for meta in self.meta]}
        # os.system('mkdir -p {}'.format(cfg.result_dir))
        # print(cfg.result_dir)
        # np.save(join(cfg.result_dir, 'render_meta.npy'), render_meta)
        # cfg_meta = {
        #     'render_center': cfg.render_center,
        #     'render_radius': cfg.render_radius,
        #     'render_axis': cfg.render_axis,
        #     'render_view_range': cfg.render_view_range,
        #     'render_angle': cfg.render_angle,
        #     'num_circle': cfg.num_circle,
        #     'render_frames': cfg.render_frames,
        #     'render_stop_frame': cfg.render_stop_frame,
        #     'render_stop_time': cfg.render_stop_time,
        #     'render_ixts_start': cfg.render_ixts_start,
        #     'render_ixts_end': cfg.render_ixts_end,
        #     'render_ixts_ratio': cfg.render_ixts_ratio,
        #     'render_smooth': cfg.render_smooth}
        # import json
        # json.dump(cfg_meta, open(
        #     join(cfg.result_dir, 'cfg.json'), 'w'), indent=4)

    def get_metas(self, render_views, input_views):
        if self.render_path:
            self.get_render_path_meta(input_views)
            return
        for frame_id in tqdm(range(*self.cfg.frame_sample), desc='loading meta'):
            for view_id in render_views:
                self.preload_img(frame_id, view_id)
                self.meta += [(frame_id, view_id)]

    def load_cameras(self, **kwargs):
        intri_path = join(self.data_root, 'intri.yml')
        extri_path = join(self.data_root, 'extri.yml')
        cams = easy_utils.read_camera(intri_path, extri_path)
        self.basenames = sorted([k for k in cams])
        cam_len = len(self.basenames)
        self.ixts = np.array([cams[cam_id]['K'] for cam_id in self.basenames]).reshape(
            cam_len, 3, 3).astype(np.float32)
        self.orig_ixts = self.ixts.copy()
        exts = np.array([cams[cam_id]['RT'] for cam_id in self.basenames]).reshape(
            cam_len, 3, 4).astype(np.float32)
        exts_ones = np.zeros_like(exts[:, :1, :])
        exts_ones[..., 3] = 1.
        self.exts = np.concatenate([exts, exts_ones], axis=1)
        transform = cfg.get('transform', np.eye(4))
        self.exts = (self.exts @ np.linalg.inv(transform)
                     ).astype(np.float32)  # align z to up
        self.exts_inv = np.linalg.inv(self.exts).astype(np.float32)
        self.Ds = np.array([cams[cam_id]['D'] for cam_id in self.basenames]).reshape(
            cam_len, 5).astype(np.float32)
        if kwargs.get('ignore_dist_k3', False):
            self.Ds[..., 4] = 0

    def preload_data_func(self):
        # preloading
        self.meta_imgs_buffer = parallel_execution(self.meta_imgs_path,
                                                   [self.orig_ixts[view_id]
                                                       for view_id in self.meta_viewids],
                                                   [self.Ds[view_id]
                                                       for view_id in self.meta_viewids],
                                                   [self.crop_h_w for i in range(
                                                       len(self.meta_viewids))],
                                                   # [self.ccms[view_id] for view_id in (self.meta_viewids)],
                                                   [self.rs[view_id]
                                                       for view_id in (self.meta_viewids)],
                                                   '.png',
                                                   action=load_resize_undist_im_bytes, num_processes=64, print_progress=True, sequential=False, async_return=False, desc='preload bytes for images')
        if self.have_msk:
            self.meta_msks_buffer = parallel_execution(self.meta_msks_path,
                                                   [self.orig_ixts[view_id]
                                                       for view_id in self.meta_viewids],
                                                   [self.Ds[view_id]
                                                       for view_id in self.meta_viewids],
                                                   [self.crop_h_w for i in range(
                                                       len(self.meta_viewids))],
                                                   [self.rs[view_id]
                                                       for view_id in (self.meta_viewids)],
                                                   '.png',
                                                   action=load_resize_undist_im_bytes, num_processes=64, print_progress=True, sequential=False, async_return=False, desc='preload bytes for masks')

    def read_img(self, frame_id, view_id):
        meta_idx = self.meta_idx[frame_id][view_id]
        if self.preload_data:
            img = load_image_from_bytes(
                self.meta_imgs_buffer[meta_idx], decode_flag=cv2.IMREAD_COLOR)[..., [2, 1, 0]]
            if self.have_msk:
                msk = load_image_from_bytes(
                self.meta_msks_buffer[meta_idx], decode_flag=cv2.IMREAD_GRAYSCALE)
            else: msk  = (np.ones_like(img[..., 0]) * 255).astype(np.uint8)
        else:
            img_path = self.meta_imgs_path[meta_idx]
            img = cv2.imread(img_path, cv2.IMREAD_COLOR)
            K, D, ratio = self.orig_ixts[view_id], self.Ds[view_id], self.rs[view_id]
            if np.linalg.norm(D) != 0.: img = cv2.undistort(img, K, D)
            if ratio != 1.: img = cv2.resize(img, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA)
            if self.have_msk:
                msk_path = self.meta_msks_path[meta_idx]
                msk = cv2.imread(msk_path, cv2.IMREAD_GRAYSCALE)
                if np.linalg.norm(D) != 0.: msk = cv2.undistort(msk, K, D)
                if ratio != 1.: msk = cv2.resize(msk, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_NEAREST)
            else:
                msk = (np.ones_like(img[..., 0]) * 255).astype(np.uint8)
            if self.crop_h_w[0] != -1 and self.crop_h_w[0] != img.shape[0]:
                assert(self.crop_h_w[0] < img.shape[0])
                crop = (img.shape[0] - self.crop_h_w[0]) // 2
                img = img[crop:crop+self.crop_h_w[0], ...]
                msk = msk[crop:crop+self.crop_h_w[0]]
            if self.crop_h_w[1] != -1 and self.crop_h_w[1] != img.shape[1]:
                assert(self.crop_h_w[1] < img.shape[1])
                crop = (img.shape[1] - self.crop_h_w[1]) // 2
                img = img[:, crop:crop+self.crop_h_w[1], ...]
                msk = msk[:, crop:crop+self.crop_h_w[1]]
            img = img[..., [2, 1, 0]]
        ###################
        # import matplotlib.pyplot as plt
        # plt.imsave('msk_orig.png', msk)
        # msk = (msk > 0.5).astype(np.uint8)
        # kernel = np.ones((5, 5), np.uint8)
        # plt.imsave('msk_01.png', msk)
        # msk = cv2.erode(msk, kernel)
        # plt.imsave('msk_erode.png', msk)
        # msk = cv2.dilate(msk, kernel)
        # plt.imsave('msk_dilate.png', msk)
        # msk = msk.astype(np.float32)
        # plt.imsave('msk_final.png', msk)
        ###################
        if 'schp' not in self.msk_dir: msk = (msk / 255.).astype(np.float32)
        else: msk = (msk != 0).astype(np.float32)
        img = (img / 255.).astype(np.float32)
        img = img * msk[..., None]
        if cfg.white_bkgd:
            img = img + 1 * (1 - msk[..., None])
        return img, msk

    def __getitem__(self, index):
        if self.split == 'train':
            batch_ids = np.random.choice(
                np.arange(len(self.meta)), self.imgs_per_batch, replace=False)
            ret = [self.get_batch(index) for index in batch_ids]
        else:
            ret = [self.get_batch(index)]
        return default_collate(ret)

    def sample_pixels(self, msk, frame_id, tar_view, h, w):
        num_pixels = cfg.num_pixels
        if self.bbox_type in ['RENBODY', 'NHR']:
            sample_msk = base_utils.get_bound_2d_mask(self.bounds[frame_id], self.ixts[tar_view], self.exts[tar_view][:3], h, w)
            select_inds = np.random.choice(sample_msk.reshape(-1).nonzero()[0], size=[num_pixels], replace=True)
            if cfg.get('num_patches', 0) >= 1: 
                patch_ratio = np.random.choice(cfg.patch_ratio)
                select_inds_ = PixelSampler.sample_patches(cfg.num_patches, cfg.patch_size, patch_ratio, h, w, (msk!=0).astype(np.uint8))
                select_inds_ = (select_inds_[1] * w + select_inds_[0]).astype(np.int32)
                select_inds = np.concatenate([select_inds, select_inds_], axis=0)
        else:
            select_inds = np.random.choice(h*w, size=[cfg.num_pixels], replace=True)
            sample_msk = None
            if cfg.get('num_patches', 0) >= 1: import ipdb; ipdb.set_trace()
        return select_inds, sample_msk

    def get_near_far(self, ext, ixt, h, w, frame_id):
        near_far = self.near_far.copy()
        # logger.debug('orig near_far: ' + str(near_far))
        if self.point_cloud is not None:
            near_far_ = get_near_far(self.point_cloud, np.linalg.inv(ext), ixt, h, w)
            near_far_[1] *= self.cfg.get('far_plane_ratio', 1.4)
            near_far_[0] *= self.cfg.get('near_plane_ratio', 0.8)
            near_far[0] = max(near_far[0], near_far_[0])
            near_far[1] = min(near_far[1], near_far_[1])
        # logger.debug('pcd near_far: ' + str(near_far))
        view_corners = self.corners[frame_id] @ ext[:3, :3].T + ext[:3, 3:].T
        near_far_ = np.array([view_corners[:, 2].min(), view_corners[:, 2].max()])
        near_far[0] = max(near_far[0], near_far_[0])
        near_far[1] = min(near_far[1], near_far_[1])
        # logger.debug('final near_far: ' + str(near_far))
        return near_far

    def get_batch(self, index):
        index = index % len(self.meta)
        frame_id, tar_view = self.meta[index][0], self.meta[index][1]
        if cfg.render_path: h, w = self.render_h_w
        else:
            rgb, msk = self.read_img(frame_id, tar_view)
            h, w = rgb.shape[:2]
        coords = self.coords[:h, :w].reshape(-1, 3).copy()
        if self.split == 'train':
            select_inds, sample_msk = self.sample_pixels(msk, frame_id, tar_view, h, w)
            coords = coords[select_inds]
            rgb, msk = rgb[coords[:, 1], coords[:, 0]
                           ], msk[coords[:, 1], coords[:, 0]]
        if cfg.render_path:
            rgb, msk = np.ones_like(coords[..., :3]), np.ones_like(coords[..., 0])
        else:
            rgb, msk = rgb.reshape(-1, 3), msk.reshape(-1)

        t = (frame_id - self.start_frame) / max((self.frame_len - 1),1)
        render_ext = self.render_exts[index] if cfg.render_path else self.exts[tar_view]
        render_ixt = self.render_ixts[index] if cfg.render_path else self.ixts[tar_view]
        render_ext_inv = self.render_exts_inv[index] if cfg.render_path else self.exts_inv[tar_view]
        render_ixt_inv = self.render_ixts_inv[index] if cfg.render_path else self.ixts_inv[tar_view]
        near_far = self.get_near_far(render_ext, render_ixt, h, w, frame_id)

        ret = {}
        if self.have_msk and 'enerf' in cfg.task and self.split == 'train': # compute tight cost volume
            ret.update({'mask_at_box': sample_msk.astype(np.uint8)})
        ret.update({'coords': coords, 'rgb': rgb, 'msk': msk})
        ret.update({'tar_ext': render_ext.astype(np.float32),
                    'tar_ixt': render_ixt.astype(np.float32),
                    'tar_ext_inv': render_ext_inv.astype(np.float32),
                    'tar_ixt_inv': render_ixt_inv.astype(np.float32),
                    'time': np.array(t).astype(np.float32),
                    'h_w': np.array([h, w]).astype(np.int32),
                    'near_far': near_far.astype(np.float32)})
        if self.have_msk: ret.update({'bounds': self.bounds[frame_id].astype(np.float32)})
        if self.have_msk and 'renbody' in self.data_root and self.cfg.get('use_global_bbox', True): ret.update({'bounds': self.bound.astype(np.float32)}) #
        ret.update({'meta': {'H': h, 'W': w, 'idx': index, 'frame_id': frame_id, 'view_id': tar_view}})
        return ret

    def __len__(self):
        return len(self.meta)

lib/datasets/volcap/base_dataset.py
zju3dv/im4d
import numpy as np
import cv2
import os
from os.path import join
from tqdm import tqdm
from torch.utils.data.dataloader import default_collate

from lib.config import cfg, logger
from lib.datasets.volcap.base_dataset import Dataset as BaseDataset
from lib.utils.data_utils import load_path_pose


class Dataset(BaseDataset):
    
    def get_render_path_meta(self, input_views):
        super().get_render_path_meta(input_views)
        for index, meta in enumerate(self.meta):
            render_ext = self.render_exts[index]
            near_views = self.get_near_views(np.linalg.inv(render_ext), input_views)
            self.meta[index] = (meta[0], -1, near_views)
            for idx in near_views:
                self.preload_img(meta[0], idx)

    def get_near_views(self, c2w, input_views, remove_view=-1):
        input_cam_xyz = np.array([self.exts_inv[view_id][:3, 3]
                                 for view_id in input_views])
        render_cam_xyz = c2w[:3, 3]
        spatial_distances = np.linalg.norm(
            input_cam_xyz - render_cam_xyz[None], axis=-1)
        
        input_cam_rot = np.array([self.exts_inv[view_id][:3, :3] for view_id in input_views])
        render_cam_rot = c2w[:3, :3]
        rot_distances = np.linalg.norm(input_cam_rot - render_cam_rot[None], axis=(-2, -1))
        argsorts = np.argsort(spatial_distances)
        argsorts_rot = np.argsort(rot_distances)
        max_input_views = self.cfg.get('train_input_views')[-1] + 2 if self.split == 'train' else self.cfg.get('test_input_views')
        rot_num = max_input_views // 2
        spatial_num = max_input_views - rot_num
        if remove_view in input_views:
            return [input_views[i] for i in argsorts[1:max_input_views+1]]
        else:
            return [input_views[i] for i in argsorts[1:max_input_views+1]]
            ret_views = [input_views[i] for i in argsorts[:spatial_num]] + [input_views[i] for i in argsorts_rot[:rot_num]]
            # # may repeat, then remove the repeated and add new ones acoording to rot_distances
            # ret_views = list(set(ret_views))
            # while len(ret_views) < max_input_views:
            #     for i in range(len(input_views)):
            #         if input_views[i] not in ret_views:
            #             ret_views.append(input_views[i])
            #             break
            # print(ret_views)
            # return ret_views
                

    def get_metas(self, render_views, input_views):
        logger.info('input views: ' + str(input_views))
        if self.render_path:
            self.get_render_path_meta(input_views)
            return
        for frame_id in tqdm(range(*self.cfg.get('frame_sample')), desc='loading meta'):
            for view_id in render_views:
                self.preload_img(frame_id, view_id)
                near_views = self.get_near_views(
                    self.exts_inv[view_id], input_views, remove_view=view_id)
                for idx in near_views:
                    self.preload_img(frame_id, idx)
                self.meta += [(frame_id, view_id, near_views)]

    def __getitem__(self, index):
        if self.split == 'train':
            batch_ids = np.random.choice(
                np.arange(len(self.meta)), self.imgs_per_batch, replace=False)
            num_input_views = np.random.choice(self.cfg.get(
                'train_input_views'), p=self.cfg.get('train_input_views_prob'))
            ret = [self.get_batch(index, num_input_views=num_input_views)
                   for index in batch_ids]
            ret = self.padding_srcinps(ret)
        else:
            ret = [self.get_batch(
                index, num_input_views=self.cfg.get('test_input_views'))]
        return default_collate(ret)

    def padding_srcinps(self, ret):
        hs = [batch['src_inps'].shape[2] for batch in ret]
        ws = [batch['src_inps'].shape[3] for batch in ret]
        H, W = np.max(hs), np.max(ws)
        for idx, batch in enumerate(ret):
            h, w = hs[idx], ws[idx]
            if h != H or w != W:
                src_inps = np.ones((batch['src_inps'].shape[0], batch['src_inps'].shape[1], H, W), dtype=np.float32) if cfg.white_bkgd else np.zeros(
                    (batch['src_inps'].shape[0], batch['src_inps'].shape[1], H, W), dtype=np.float32) - 1.
                crop_h, crop_w = (H - h)//2, (W - w)//2
                src_inps[:, :, crop_h:crop_h+h,
                         crop_w:crop_w+w] = batch['src_inps']
                batch['src_inps'] = src_inps
                batch['src_ixts'][:, 0, 2] += crop_w
                batch['src_ixts'][:, 1, 2] += crop_h
        return ret

    def get_batch(self, index, num_input_views):
        ret = super().get_batch(index)
        frame_id, tar_view, near_views = self.meta[index]
        if self.split == 'train':
            if np.random.random() <= 0.05:
                near_views = [tar_view] + near_views
            src_views = np.random.choice(
                near_views[:num_input_views+2], num_input_views, replace=False)
        else:
            src_views = near_views[:num_input_views]
        src_inps, src_exts, src_ixts = self.read_srcs(frame_id, src_views)
        ret.update({'src_inps': src_inps, 'src_exts': src_exts, 'src_ixts': src_ixts})
        return ret

    def read_srcs(self, frame_id, src_views):
        if not self.cfg.get('crop_srcinps', False):
            src_inps = (np.asarray([self.read_img(frame_id, view_id)[
                        0] for view_id in src_views]).transpose(0, 3, 1, 2)).astype(np.float32) * 2. - 1.
            src_exts = self.exts[src_views]
            src_ixts = self.ixts[src_views]
            return src_inps, src_exts, src_ixts
        src_imgs, src_ixts = [], []
        ws, hs = [], []
        for view_id in src_views:
            img, msk = self.read_img(frame_id, view_id)
            x, y, w, h = cv2.boundingRect(msk.astype(np.uint8))
            ws.append(w)
            hs.append(h)
            src_imgs.append((img[y:y+h, x:x+w]).astype(np.float32))
            src_ixt = self.ixts[view_id].copy()
            src_ixt[0, 2] -= x
            src_ixt[1, 2] -= y
            src_ixts.append(src_ixt)
        src_exts = self.exts[src_views]
        src_ixts = np.array(src_ixts)
        W, H = (np.max(ws) // 8 + 1) * 8, (np.max(hs) // 8 + 1) * 8
        # padding src_imgs to the same shape (H, W)
        src_inps = []
        for idx, src_img in enumerate(src_imgs):
            src_inp = np.ones((H, W, 3), dtype=np.float32) if cfg.white_bkgd else np.zeros(
                (H, W, 3), dtype=np.float32)
            h, w = src_img.shape[:2]
            crop_h, crop_w = (H - h)//2, (W - w)//2
            src_inp[crop_h:crop_h+h, crop_w:crop_w+w] = src_img
            src_ixts[idx][0, 2] += crop_w
            src_ixts[idx][1, 2] += crop_h
            src_inps.append(src_inp)
        src_inps = np.array(src_inps).transpose(0, 3, 1, 2) * 2 - 1
        return src_inps, src_exts, src_ixts

lib/datasets/volcap/ibr_dataset.py
zju3dv/im4d
import numpy as np
from lib.config import cfg, logger
import os
import json
from os.path import join
from lib.utils.data_utils import save_img, export_dir_to_video, save_depth, save_srcinps, dpt_convert_to_img, rgb_convert_to_img, export_depth_to_pcd
from termcolor import colored

class BaseEvaluator:
    def __init__(self,):
        self.step = 999999
        if cfg.clear_result: os.system('rm -rf ' + join(cfg.result_dir, 'step{:08d}'.format(self.step), '*'))
        os.system('mkdir -p ' + cfg.result_dir)
        # meta information
        if 'num_levels' in cfg.network:
            self.num_levels = cfg.network.num_levels
        else:
            self.num_levels = len(cfg.samplers)
        self.reset_metrics()
        self.evaluate_metrics = cfg.evaluate.metrics
        
    def reset_metrics(self,):
        self.metrics = {}
        for level in range(self.num_levels):
            self.metrics.update({k+f'_{level}':[] for k in cfg.evaluate.metrics})
        
    def update_meta(self, output, batch):
        if 'step' in batch:
            self.step = batch['step']
            
    def update_metrics(self, output, batch, img_set):
        metric_stats = img_set['meta'] + ' '
        for level in range(self.num_levels):
            for k in self.evaluate_metrics:
                metric_item = getattr(self, 'metric_' + k)(output, batch, img_set, level=level)
                if metric_item is not None: self.metrics[k + f'_{level}'].append(metric_item)
                if len(self.metrics[k + f'_{level}']) == 0: continue
                metric_stats += '{}: {:.3f} '.format(k + f'_{level}', self.metrics[k + f'_{level}'][-1])
        logger.info(metric_stats)
    
    def save_results(self, output, batch, img_set):
        if not cfg.save_result:
            return
        meta = img_set['meta']
        for k in img_set:
            for kk in cfg.evaluate.save_keys:
                if kk in k:
                    acc_k = k.replace(kk, 'acc')
                    if cfg.evaluate.get('save_with_alpha', True) and acc_k in img_set: acc = img_set[acc_k]
                    else: acc = None
                    img_path = join(cfg.result_dir, 'step{:08d}/{}/{}'.format(self.step, k, meta))
                    if 'depth' in k:
                        save_depth(img_path, img_set[k], acc_map=acc)
                    else:
                        save_img(img_path, img_set[k], acc_map=acc)
        if 'src_inps' in cfg.evaluate.save_keys:
            img_path = join(cfg.result_dir, 'step{:08d}/{}/{}'.format(self.step, 'src_inps', meta))
            save_srcinps(img_path, batch['src_inps'].detach().cpu().numpy()[0])
            
    def parse_imgs(self, output, batch):
        pass
        
    def evaluate(self, output, batch):
        self.update_meta(output, batch)
        img_set = self.parse_imgs(output, batch)
        # export_depth_to_pcd(img_set['depth_0'] / np.clip(img_set['acc_0'], 0.01, 1.), batch['tar_ixt'][0].detach().cpu().numpy(), 'debug_depth0.ply')
        # export_depth_to_pcd(img_set['depth_1'] / np.clip(img_set['acc_1'], 0.01, 1.), batch['tar_ixt'][0].detach().cpu().numpy(), 'debug_depth1.ply')
        self.update_metrics(output, batch, img_set)
        self.save_results(output, batch, img_set)
        return self.to_image_stats(img_set)
    
    def to_image_stats(self, img_set):
        image_stats = {}
        
        for k in img_set:
            for kk in cfg.evaluate.save_keys:
                if kk in k:
                    acc_k = k.replace(kk, 'acc')
                    if cfg.evaluate.get('save_with_alpha', True) and acc_k in img_set: acc = img_set[acc_k]
                    else: acc = None
                    if 'depth' in k:
                        image_stats[k + '_' + img_set['meta']] = dpt_convert_to_img(img_set[k], acc).transpose(2, 0, 1)
                    elif 'rgb' in k:
                        image_stats[k + '_' + img_set['meta']] = rgb_convert_to_img(img_set[k], acc).transpose(2, 0, 1)
                    else:
                        pass
        return image_stats
        
    def ret_metrics(self,):
        ret = {k:float(np.mean(self.metrics[k])) for k in self.metrics if len(self.metrics[k]) > 0}
        self.reset_metrics()
        print('\n', '=' * 80, '\n', ret, '\n', '=' * 80)
        if cfg.save_result: print(colored('Results saved to {}'.format(cfg.result_dir), color='green'))
        os.makedirs(join(cfg.result_dir, 'step{:08d}'.format(self.step)), exist_ok=True)
        json.dump(ret, open(os.path.join(cfg.result_dir, 'step{:08d}'.format(self.step), 'metrics.json'), 'w'))
        return ret
        
    def write_videos(self,):
        if not cfg.get('write_video', False):
            return
        save_dir = join(cfg.result_dir, 'step{:08d}'.format(self.step))
        # 获取save_dir下所有目录
        dirs = [item for item in os.scandir(save_dir) if item.is_dir()]
        for dir in dirs:
            for k in cfg.evaluate.save_keys:
                if k in dir.name:
                    image_dir = dir.path
                    video_path = join(save_dir, 'mp4s/{}.mp4'.format(dir.name))
                    os.system('mkdir -p {}'.format(os.path.dirname(video_path)))
                    export_dir_to_video(image_dir, video_path, fps=cfg.fps)
        print(colored('Videos saved to {}'.format(join(save_dir, 'mp4s')), color='green'))
            
    def summarize(self):
        # eval and save metrics
        ret = self.ret_metrics()
        self.write_videos()
        return ret

lib/evaluators/base_evaluator.py
zju3dv/im4d
from lib.evaluators.nerf import Evaluator as BaseEvaluator
from lib.utils.data_utils import save_img, export_dir_to_video, save_depth, save_srcinps, dpt_convert_to_img, rgb_convert_to_img
from os.path import join
from lib.config import cfg

class Evaluator(BaseEvaluator):
    
    def save_results(self, output, batch, img_set):
        super().save_results(output, batch, img_set)
        if not cfg.save_enerf:
            return
        t_mean_std_min_max = output['batch_share_info']['t_mean_std_min_max'].permute(0, 2, 3, 1)
        meta = img_set['meta']
        ks = ['mean', 'std', 'min', 'max'] 
        idx = 0
        for k in ks:
            img_path = join(cfg.result_dir, 'step{:08d}/{}/{}'.format(self.step, k, meta))
            save_img(img_path, t_mean_std_min_max[0, :, :, idx].detach().cpu().numpy())
            idx += 1
        
lib/evaluators/enerf.py
zju3dv/im4d
import imp
import os

def _evaluator_factory(cfg):
    module = cfg.evaluator_module
    path = cfg.evaluator_path
    evaluator = imp.load_source(module, path).Evaluator()
    return evaluator


def make_evaluator(cfg):
    if cfg.skip_eval:
        return None
    else:
        return _evaluator_factory(cfg)

lib/evaluators/make_evaluator.py
zju3dv/im4d
import numpy as np
from lib.config import cfg
import cv2
from os.path import join
from lib.evaluators.base_evaluator import BaseEvaluator
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr
import lpips
import torch

class Evaluator(BaseEvaluator):
    def __init__(self,):
        super(Evaluator, self).__init__()
        if 'lpips' in cfg.evaluate.metrics:
            import warnings
            # warnings.filterwarnings('ignore')
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=UserWarning)
                self.loss_fn_vgg = lpips.LPIPS(net='vgg', verbose=False).cuda()
            # warnings.resetwarnings()
            # del warnings
            
    def parse_imgs_msk(self, output, batch, level=-1):
        suffix = '' if level == -1 else f'_{level}'
        h, w = batch['meta']['H'].item(), batch['meta']['W'].item()
        ret = {}
        mask_at_box = batch['mask_at_box'].detach().cpu().numpy().reshape(h, w)
        if level == -1: ret['eval_mask'] = mask_at_box 
        if level == -1: gt = np.ones((h, w, 3)) if cfg.white_bkgd else np.zeros((h, w, 3))
        rgb = np.ones((h, w, 3)) if cfg.white_bkgd else np.zeros((h, w, 3))
        depth = np.zeros((h, w))
        acc = np.zeros((h, w))
        if level == -1: gt[mask_at_box==1] = batch['rgb'].detach().cpu().numpy()[0]
        if f'rgb_map{suffix}' in output: 
            rgb[mask_at_box==1] = output[f'rgb_map{suffix}'].detach().cpu().numpy()[0]; 
            ret.update({f'rgb{suffix}': rgb.astype(np.float32)})
        if f'depth_map{suffix}' in output: 
            depth[mask_at_box==1] = output[f'depth_map{suffix}'].detach().cpu().numpy()[0]; 
            ret.update({f'depth{suffix}': depth.astype(np.float32)})
        if f'acc_map{suffix}' in output: 
            acc[mask_at_box==1] = output[f'acc_map{suffix}'].detach().cpu().numpy()[0]; 
            ret.update({f'acc{suffix}': acc.astype(np.float32)})
        if level == -1: 
            ret['gt'] = gt.astype(np.float32);
            view_id = batch['meta']['view_id'].item()
            view_id = batch['meta']['idx'].item() if view_id == -1 else view_id
            ret['meta'] = 'frame{:04d}_view{:04d}.png'.format(batch['meta']['frame_id'].item(), view_id) 
        return ret
        
    def parse_imgs(self, output, batch, level=-1):
        
        if 'mask_at_box' in batch:
            ret = {}
            ret.update(self.parse_imgs_msk(output, batch))
            for i in range(self.num_levels):
                ret.update(self.parse_imgs_msk(output, batch, level=i))
            return ret
        
        ret = {}
        h, w = batch['meta']['H'].item(), batch['meta']['W'].item()
        if level == -1:
            gt = batch['rgb'].detach().cpu().numpy().reshape(h, w, 3)
            view_id = batch['meta']['view_id'].item()
            view_id = batch['meta']['idx'].item() if view_id == -1 else view_id
            ret['meta'] = 'frame{:04d}_view{:04d}.png'.format(batch['meta']['frame_id'].item(), view_id)
            ret['eval_mask'] = np.ones_like(gt[..., 0]).astype(np.uint8)
            ret['gt'] = gt.astype(np.float32)
            for i in range(self.num_levels):
                ret.update(self.parse_imgs(output, batch, level=i))
        
        suffix = '' if level == -1 else f'_{level}'
        
        if f'rgb_map{suffix}' in output:  ret.update({f'rgb{suffix}': output[f'rgb_map{suffix}'].detach().cpu().numpy().reshape(h, w, 3).astype(np.float32)})
        if f'depth_map{suffix}' in output: ret.update({f'depth{suffix}': output[f'depth_map{suffix}'].detach().cpu().numpy().reshape(h, w).astype(np.float32)})
        if f'acc_map{suffix}' in output:   ret.update({f'acc{suffix}': output[f'acc_map{suffix}'].detach().cpu().numpy().reshape(h, w).astype(np.float32)})

        return ret
    
    def metric_psnr(self, output, batch, img_set, level=-1):
        suffix = '' if level == -1 else f'_{level}'
        if f'rgb{suffix}' not in img_set: return None
        psnr_item = psnr(img_set['gt'][img_set['eval_mask']==1], img_set[f'rgb{suffix}'][img_set['eval_mask']==1], data_range=1.)
        return psnr_item
        
    def metric_ssim(self, output, batch, img_set, level=-1):
        suffix = '' if level == -1 else f'_{level}'
        if f'rgb{suffix}' not in img_set: return None
        x, y, w, h = cv2.boundingRect(img_set['eval_mask'])
        ssim_item = ssim(img_set['gt'][y:y+h, x:x+w], img_set[f'rgb{suffix}'][y:y+h, x:x+w], data_range=2., channel_axis=-1)
        return ssim_item
        
    def metric_lpips(self, output, batch, img_set, level=-1):
        suffix = '' if level == -1 else f'_{level}'
        if f'rgb{suffix}' not in img_set: return None
        x, y, w, h = cv2.boundingRect(img_set['eval_mask'])
        gt, pred = img_set['gt'][y:y+h, x:x+w], img_set[f'rgb{suffix}'][y:y+h, x:x+w]
        gt, pred = torch.from_numpy(gt).cuda(), torch.from_numpy(pred).cuda()
        gt, pred = gt.permute(2, 0, 1).unsqueeze(0), pred.permute(2, 0, 1).unsqueeze(0)
        lpips_item = self.loss_fn_vgg(gt * 2. -  1., pred * 2 - 1.).item()
        return lpips_item

lib/evaluators/nerf.py
zju3dv/im4d
import os
import imp


def make_network(cfg):
    module = cfg.network_module
    path = cfg.network_path
    network = imp.load_source(module, path).Network(cfg.network)
    return network

lib/networks/make_network.py
zju3dv/im4d
import torch.nn as nn
activations = {
    'relu': nn.ReLU,
    'none': None
}

def get_activation(name):
    return activations[name]
lib/networks/nerf/activations/__init__.py
zju3dv/im4d
from lib.networks.nerf.encoders.frequency_encoder import FrequencyEncoder
from lib.networks.nerf.encoders.none_encoder import NoneEncoder
from lib.networks.nerf.encoders.imgfeat_ibrnet import ImgfeatIbrnet
from lib.networks.nerf.encoders.tcnn_wrapper import TCNN
from lib.networks.nerf.encoders.kplanes_encoder import KplanesEncoder
from lib.networks.nerf.encoders.im4d_encoder import Im4DEncoder

encoding_dict = {
    'frequency': FrequencyEncoder,
    'none': NoneEncoder,
    'imgfeat_ibrnet': ImgfeatIbrnet,
    'tcnn': TCNN,
    'kplanes': KplanesEncoder,
    'im4d': Im4DEncoder,
}

def get_encoder(encoding_cfg):
    return encoding_dict[encoding_cfg.type](**encoding_cfg)

lib/networks/nerf/encoders/__init__.py
zju3dv/im4d
import torch
import torch.nn as nn
from typing import Optional

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class FrequencyEncoder(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        in_dim: int, 
        num_frequencies: int, 
        include_input: bool = False,
        min_freq_exp: Optional[int] = None,
        max_freq_exp: Optional[int] = None,
        **kwargs
    ) -> None:
        super(FrequencyEncoder, self).__init__()
        self.in_dim = in_dim
        
        self.num_frequencies = num_frequencies
        self.min_freq = min_freq_exp or 0
        self.max_freq = max_freq_exp or self.num_frequencies - 1
        
        self.freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(device)
        
        self.include_input = include_input

    def get_out_dim(self) -> int:
        if self.in_dim is None:
            raise ValueError("Input dimension has not been set")
        out_dim = self.in_dim * self.num_frequencies * 2
        if self.include_input:
            out_dim += self.in_dim
        return out_dim
    
    def forward(
        self, in_tensor, covs = None, **kwargs):
        
        scaled_inputs = in_tensor[..., None] * self.freqs        
        scaled_inputs = scaled_inputs.view(*scaled_inputs.shape[:-2], -1)  # [..., "input_dim" * "num_scales"]

        # sinx cosx sin2x cos2x sin4x cos4x ...
        if covs is None:
            encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1))
        else:
            input_var = torch.diagonal(covs, dim1=-2, dim2=-1)[..., :, None] * self.freqs[None, :] ** 2
            input_var = input_var.reshape((*input_var.shape[:-2], -1))
            encoded_inputs = expected_sin(
                torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1)
            )
        
        # reorder input to nerf order
        reorder_encoded_inputs = torch.zeros_like(encoded_inputs).view(*in_tensor.shape[:-1], -1, self.in_dim)
        encoded_inputs = encoded_inputs.view(*in_tensor.shape[:-1], 2, self.in_dim, self.num_frequencies).transpose(-2, -1) # [..., 2, "num_scales", "input_dim"]
        reorder_encoded_inputs[..., ::2, :] = encoded_inputs[..., 0, :, :]
        reorder_encoded_inputs[..., 1::2, :] = encoded_inputs[..., 1, :, :]
        encoded_inputs = reorder_encoded_inputs.reshape(*in_tensor.shape[:-1], -1) 
        
        if self.include_input:
            encoded_inputs = torch.cat([in_tensor, encoded_inputs], dim=-1)
        return encoded_inputs
lib/networks/nerf/encoders/frequency_encoder.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
from lib.config import cfg

from lib.networks.nerf.encoders.imgfeat_ibrnet import ImgfeatIbrnet
from lib.networks.nerf.encoders.kplanes_encoder import KplanesEncoder


class Im4DEncoder(KplanesEncoder):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        **kwargs
    ) -> None:
        super(Im4DEncoder, self).__init__(**kwargs)
        self.feature_projection = ImgfeatIbrnet(**kwargs)
        
    
    def forward(
        self, in_tensor, only_geo=False, **kwargs):
        geo_feat = super().forward(in_tensor, **kwargs)
        if only_geo: return geo_feat
        # if self.training and kwargs['batch']['step'] >= cfg.joint_iters and kwargs['batch']['step'] % cfg.finetune_iters_per != 0:
            # with torch.no_grad():
                # ibr_feat = self.feature_projection(in_tensor, **kwargs)
        # else: ibr_feat = self.feature_projection(in_tensor, **kwargs)
        ibr_feat = self.feature_projection(in_tensor, **kwargs)
        return (geo_feat, ibr_feat)
lib/networks/nerf/encoders/im4d_encoder.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
from lib.config import cfg
from lib.utils.im4d.im4d_utils import Im4DUtils

class ImgfeatIbrnet(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        feat_scale = 0.25,
        feat_dim = 32,
        **kwargs
    ) -> None:
        super(ImgfeatIbrnet, self).__init__()
        self.feat_dim = feat_dim
        self.feat_scale = feat_scale

    def get_out_dim(self) -> int:
        return 3 + self.feat_dim + 4 + 1
    
    def forward(
        self, xyz, viewdir=None, **kwargs):
        # xyz: BxN_raysxN_samplesx3
        if cfg.get('ndc', False): xyz = Im4DUtils.ndc2world(xyz, cfg.ndc_FOCAL, cfg.ndc_H, cfg.ndc_W)
        sample_level, batch = kwargs['sample_level'], kwargs['batch']
        if xyz.shape[-1] == 3: viewdir = viewdir if viewdir is not None else kwargs['viewdir']
        else: xyz, viewdir = torch.split(xyz, [3, 3], dim=-1)
        img_feat = kwargs['batch_share_info']['img_feat'][sample_level]
        src_inps, src_exts, src_ixts = batch['src_inps'], batch['src_exts'], batch['src_ixts']
        B, S, C, H, W = src_inps.shape
        src_exts = src_exts.reshape(B * S, 4, 4)
        src_ixts = src_ixts.reshape(B * S, 3, 3)
        src_inps = src_inps.reshape(B * S, C, H, W) * 0.5 + 0.5
        F_C, F_H, F_W = img_feat.shape[2:]
        src_feat = img_feat.reshape(B*S, F_C, F_H, F_W)
        
        N_rays, N_samples = xyz.shape[1], xyz.shape[2]
        view_xyz = xyz[:, None].repeat(1, S, 1, 1, 1).reshape(B * S, N_rays * N_samples, 3) # B*S x N_rays*N_samples x 3
        if len(viewdir.shape) == 4: viewdir = viewdir[:, None].repeat(1, S, 1, 1, 1).reshape(B * S, N_rays * N_samples, 3)
        else: viewdir = viewdir[:, None, :, None].repeat(1, S, 1, N_samples, 1).reshape(B * S, N_rays * N_samples, 3)# B*S x N_rays*N_samples x 3
        
        point_src_diff = src_exts.inverse()[:, :3, 3:].transpose(-1, -2) - view_xyz
        point_src_diff = point_src_diff / (point_src_diff.norm(dim=-1, keepdim=True) + 1e-6)
        point_tar_diff = -viewdir
        point_diff_dot = (point_tar_diff * point_src_diff).sum(dim=-1, keepdim=True)
        point_diff = point_tar_diff - point_src_diff
        point_diff = point_diff / (point_diff.norm(dim=-1, keepdim=True) + 1e-6)
        point_dir = torch.cat([point_diff, point_diff_dot], dim=-1)
        
        view_xyz = view_xyz @ src_exts[:, :3, :3].transpose(-1, -2) + src_exts[:, :3, 3:].transpose(-1, -2)
        view_xyz = view_xyz @ src_ixts[:, :3, :3].transpose(-1, -2)
        
        # view_xyz_ = view_xyz[..., :2] / torch.clamp_min(view_xyz[..., 2:], 1e-6)
        # view_xyz_[..., 0] /= (W-1)
        # view_xyz_[..., 1] /= (H-1)
        # xyz_depth = view_xyz
        # view_xyz = view_xyz_
        grid_xyz = view_xyz[..., :2] / torch.clamp_min(view_xyz[..., 2:], 1e-6)
        grid_xyz[..., 0] /= (W-1)
        grid_xyz[..., 1] /= (H-1)
        
        point_msk = ((grid_xyz[..., 0] >= 0) & (grid_xyz[..., 0] <= 1) & (grid_xyz[..., 1] >= 0) & (grid_xyz[..., 1] <= 1) & (view_xyz[..., 2] >= 1e-6)).float()
        point_rgb = F.grid_sample(src_inps, grid_xyz[..., None, :2]*2-1, align_corners=True, padding_mode='border')[..., 0].transpose(-1, -2)
        point_feat = F.grid_sample(src_feat, grid_xyz[..., None, :2]*2-1, align_corners=True, padding_mode='border')[..., 0].transpose(-1, -2)
        
        point_feat = torch.cat([point_rgb, point_feat, point_dir, point_msk[..., None]], dim=-1)
        
        return point_feat.reshape(B, S, N_rays, N_samples, self.get_out_dim()) # BxSxN_raysxN_samplesx(out_dim)
lib/networks/nerf/encoders/imgfeat_ibrnet.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F


class KplanesEncoder(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        **kwargs
    ) -> None:
        super().__init__()
        spatial_res, temporal_res = kwargs['spatial_res'], kwargs['temporal_res']
        assert(len(spatial_res) == len(temporal_res))
        num_feat_ch = kwargs['num_feat_ch']
        self.num_res, self.num_feat_ch, self.include_input = len(spatial_res), num_feat_ch, kwargs['include_input']
        self.spatial_embedding = nn.ParameterList([nn.Parameter(torch.zeros(3, num_feat_ch, spatial_res[i], spatial_res[i])) for i in range(self.num_res)])
        self.temporal_embedding = nn.ParameterList([nn.Parameter(torch.zeros(3, num_feat_ch, temporal_res[i], spatial_res[i])) for i in range(self.num_res)])
        
        self.res_composite = kwargs['res_composite']
        self.planes_composite = kwargs['planes_composite'] # concat, product, sum
        std = 1e-1
        for data in self.spatial_embedding: data.data.uniform_(-std, std)
        for data in self.temporal_embedding: data.data.uniform_(-std, std)
        if self.planes_composite == 'product':
            for data in self.spatial_embedding: nn.init.uniform_(data, a=0.1, b=0.15)
            for data in self.temporal_embedding: nn.init.ones_(data)
        
    def get_out_dim(self) -> int:
        return self.num_res * (self.num_feat_ch * 6 if self.planes_composite == 'concat' else self.num_feat_ch) + (4 if self.include_input else 0)
    
    def forward(
        self, in_tensor, **kwargs):
        # TODO: spatial
        sh = in_tensor.shape
        bbox = kwargs.get('bbox')
        xyz = (in_tensor.reshape(-1, 3) - bbox[:1]) / (bbox[1:] - bbox[:1])
        xyz = torch.clip(xyz, 0, 1)
        time = kwargs['time'].reshape(-1, 1)
        
        xy, yz, xz = xyz[:, :2], xyz[:, 1:], xyz[:, [0, 2]]
        spatial_coord = torch.stack([xy, yz, xz], dim=0) 
        xt, yt, zt = torch.cat([xyz[:, :1], time], dim=-1), torch.cat([xyz[:, 1:2], time], dim=-1), torch.cat([xyz[:, 2:], time], dim=-1)
        temporal_coord = torch.stack([xt, yt, zt], dim=0)
        
        output = []
        for i in range(self.num_res):
            spatial_embedding = F.grid_sample(self.spatial_embedding[i], spatial_coord[:, None]*2-1, align_corners=True, mode='bilinear')[:, :, 0].transpose(-1, -2)
            temporal_embedding = F.grid_sample(self.temporal_embedding[i], temporal_coord[:, None]*2-1, align_corners=True, mode='bilinear')[:, :, 0].transpose(-1, -2)
            embedding = torch.cat([spatial_embedding, temporal_embedding], dim=0)
            
            if self.planes_composite == 'concat': embedding = torch.cat([embedding[i] for i in range(len(embedding))], dim=-1)
            elif self.planes_composite == 'product': embedding = torch.prod(embedding, dim=0)
            elif self.planes_composite == 'sum': embedding = torch.sum(embedding, dim=0)
            else: import ipdb; ipdb.set_trace()
            
            output.append(embedding)
        if self.res_composite == 'concat': output = torch.cat(output, dim=-1)
        else: import ipdb; ipdb.set_trace()
        
        return output.reshape(sh[:-1] + (self.get_out_dim(),))
lib/networks/nerf/encoders/kplanes_encoder.py
zju3dv/im4d
import torch
import torch.nn as nn


class NoneEncoder(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        **kwargs
    ) -> None:
        super(NoneEncoder, self).__init__()
        
    def get_out_dim(self) -> int:
        return 0
    
    def forward(
        self, in_tensor, covs = None):
        return None
        
lib/networks/nerf/encoders/none_encoder.py
zju3dv/im4d
import torch
import torch.nn as nn
try:
    import tinycudann as tcnn
except:
    pass

class TCNN(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self, 
        **kwargs
    ) -> None:
        super().__init__()
        del kwargs['type']
        self.in_dim = kwargs['in_dim']
        self.encoder = tcnn.Encoding(kwargs['in_dim'], kwargs)
        self.include_input = kwargs['include_input']
        
    def get_out_dim(self) -> int:
        return self.encoder.n_output_dims + self.in_dim if self.include_input else self.encoder.n_output_dims
    
    def forward(
        self, in_tensor, **kwargs):
        # TODO: spatial
        sh = in_tensor.shape
        bbox = kwargs.get('bbox')
        xyz = (in_tensor.reshape(-1, 3) - bbox[:1]) / (bbox[1:] - bbox[:1])
        xyz = torch.clip(xyz, 0, 1)
        xyz_embedding = self.encoder(xyz)
        if self.include_input: xyz_embedding = torch.cat([xyz, xyz_embedding], dim=-1)
        return xyz_embedding.reshape(sh[:-1] + (self.get_out_dim(),))
        
lib/networks/nerf/encoders/tcnn_wrapper.py
zju3dv/im4d
from lib.networks.nerf.fields.nerf_net import NerfNet
from lib.networks.nerf.fields.enerf_net import ENeRFNet
from lib.networks.nerf.fields.mini_ibr_net import MiniIBRNet
from lib.networks.nerf.fields.im4d_net import Im4DNet

nets_dict = {
    'nerf': NerfNet,
    'enerf': ENeRFNet,
    'mini_ibr': MiniIBRNet,
    'im4d': Im4DNet,
}
def get_net(net_cfg, **kwargs):
    return nets_dict[net_cfg.type](**kwargs)
lib/networks/nerf/fields/__init__.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
from lib.config import cfg

class ENeRFNet(nn.Module):
    def __init__(self, hid_n=64, feat_ch=8+3, only_rgb=False, **kwargs):
        """
        """
        super(ENeRFNet, self).__init__()
        self.only_rgb = only_rgb
        self.hid_n = hid_n
        self.agg = Agg(feat_ch)
        self.lr0 = nn.Sequential(nn.Linear(16, hid_n),
                                 nn.ReLU())
        self.lrs = nn.ModuleList([
            nn.Sequential(nn.Linear(hid_n, hid_n), nn.ReLU()) for i in range(0)
        ])
        self.sigma = nn.Sequential(nn.Linear(hid_n, 1), nn.ReLU())
        self.color = nn.Sequential(
                nn.Linear(64+16+feat_ch+4, hid_n),
                nn.ReLU(),
                nn.Linear(hid_n, 1),
                nn.ReLU())
        self.lr0.apply(weights_init)
        self.lrs.apply(weights_init)
        self.sigma.apply(weights_init)
        self.color.apply(weights_init)

    # def forward(self, vox_feat, img_feat_rgb_dir):
    def forward(self, xyz_encoding, dir_encoding):
        img_feat_rgb_dir_msk = xyz_encoding
        img_feat_rgb_dir = xyz_encoding[..., :-1] # RGB 3, FEAT 8, DIR 4
        B, N_views, N_points = img_feat_rgb_dir.shape[:3]
        
        img_feat = self.agg(img_feat_rgb_dir)
        S = img_feat_rgb_dir.shape[1]
        
        vox_img_feat = img_feat
        x = self.lr0(vox_img_feat)
        for i in range(len(self.lrs)):
            x = self.lrs[i](x)
        sigma = None if self.only_rgb else self.sigma(x) 
        x = torch.cat((x, vox_img_feat), dim=-1)
        x = x[:, None].repeat(1, S, 1, 1, 1)
        x = torch.cat((x, img_feat_rgb_dir), dim=-1)
        color_weight = F.softmax(self.color(x), dim=1)
        color = torch.sum((img_feat_rgb_dir[..., :3] * color_weight), dim=1)
        return color, sigma

class Agg(nn.Module):
    def __init__(self, feat_ch):
        """
        """
        super(Agg, self).__init__()
        self.feat_ch = feat_ch
        if False:
            self.view_fc = nn.Sequential(
                    nn.Linear(4, feat_ch),
                    nn.ReLU(),
                    )
            self.view_fc.apply(weights_init)
        self.global_fc = nn.Sequential(
                nn.Linear(feat_ch*3, 32),
                nn.ReLU(),
                )

        self.agg_w_fc = nn.Sequential(
                nn.Linear(32, 1),
                nn.ReLU(),
                )
        self.fc = nn.Sequential(
                nn.Linear(32, 16),
                nn.ReLU(),
                )
        self.global_fc.apply(weights_init)
        self.agg_w_fc.apply(weights_init)
        self.fc.apply(weights_init)

    def forward(self, img_feat_rgb_dir):
        B, S, N_points, N_samples = img_feat_rgb_dir.shape[:4]
        if False:
            view_feat = self.view_fc(img_feat_rgb_dir[..., -4:])
            img_feat_rgb =  img_feat_rgb_dir[..., :-4] + view_feat
        else:
            img_feat_rgb =  img_feat_rgb_dir[..., :-4]
            
        var_feat = torch.var(img_feat_rgb, dim=1).view(B, 1, N_points, N_samples, self.feat_ch).repeat(1, S, 1, 1, 1)
        avg_feat = torch.mean(img_feat_rgb, dim=1).view(B, 1, N_points, N_samples, self.feat_ch).repeat(1, S, 1, 1, 1)

        feat = torch.cat([img_feat_rgb, var_feat, avg_feat], dim=-1)
        global_feat = self.global_fc(feat)
        agg_w = F.softmax(self.agg_w_fc(global_feat), dim=1)
        im_feat = (global_feat * agg_w).sum(dim=1)
        return self.fc(im_feat)

def weights_init(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight.data)
        if m.bias is not None:
            nn.init.zeros_(m.bias.data)
lib/networks/nerf/fields/enerf_net.py
zju3dv/im4d
from lib.networks.nerf.fields.nerf_net import NerfNet
from lib.networks.nerf.fields.enerf_net import ENeRFNet
import torch

class Im4DNet(NerfNet):
    def __init__(self, **kwargs):
        super(Im4DNet, self).__init__(**kwargs)
        self.enerf_net = ENeRFNet(**kwargs)
        self.extra_color = kwargs.get('extra_color', False)
        # print('Extra color: ', self.extra_color)
        
    def forward(self, xyz_encoding, view_encoding, only_geo=False):
        color_, sigma = super().forward(xyz_encoding[0], view_encoding, only_geo)
        if only_geo: return None, sigma
        # if self.training and kwargs['batch']['step'] >= cfg.joint_iters and kwargs['batch']['step'] % cfg.finetune_iters_per != 0:
        color, _ = self.enerf_net(xyz_encoding[1], view_encoding[..., :3])
        # if self.training and self.extra_color: 
        #     color = torch.cat([color, color_], dim=-1)
        return color, sigma
lib/networks/nerf/fields/im4d_net.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
from lib.config import cfg

class MiniIBRNet(nn.Module):
    def __init__(self, hid_n=64, feat_ch=8+3, only_rgb=False, **kwargs):
        """
        """
        super(MiniIBRNet, self).__init__()

    # def forward(self, vox_feat, img_feat_rgb_dir):
    def forward(self, xyz_encoding, dir_encoding):
        img_feat_rgb_dir_msk = xyz_encoding
        img_feat_rgb_dir = xyz_encoding[..., :-1] # RGB 3, FEAT 8, DIR 4
        B, N_views, N_points = img_feat_rgb_dir.shape[:3]
        rgb = img_feat_rgb_dir[..., :3]
        weight = img_feat_rgb_dir[..., 3:4]
        weight = F.softmax(F.softplus(weight-1.), dim=1)
        color = (weight * rgb).sum(dim=1)
        return color, None
        # img_feat = self.agg(img_feat_rgb_dir)
        # S = img_feat_rgb_dir.shape[1]
        
        # vox_img_feat = img_feat
        # x = self.lr0(vox_img_feat)
        # for i in range(len(self.lrs)):
            # x = self.lrs[i](x)
        # sigma = None if self.only_rgb else self.sigma(x) 
        # x = torch.cat((x, vox_img_feat), dim=-1)
        # x = x[:, None].repeat(1, S, 1, 1, 1)
        # x = torch.cat((x, img_feat_rgb_dir), dim=-1)
        # color_weight = F.softmax(self.color(x), dim=1)
        # color = torch.sum((img_feat_rgb_dir[..., :3] * color_weight), dim=1)
        # return color, sigma

class Agg(nn.Module):
    def __init__(self, feat_ch):
        """
        """
        super(Agg, self).__init__()
        self.feat_ch = feat_ch
        if False:
            self.view_fc = nn.Sequential(
                    nn.Linear(4, feat_ch),
                    nn.ReLU(),
                    )
            self.view_fc.apply(weights_init)
        self.global_fc = nn.Sequential(
                nn.Linear(feat_ch*3, 32),
                nn.ReLU(),
                )

        self.agg_w_fc = nn.Sequential(
                nn.Linear(32, 1),
                nn.ReLU(),
                )
        self.fc = nn.Sequential(
                nn.Linear(32, 16),
                nn.ReLU(),
                )
        self.global_fc.apply(weights_init)
        self.agg_w_fc.apply(weights_init)
        self.fc.apply(weights_init)

    def forward(self, img_feat_rgb_dir):
        B, S, N_points, N_samples = img_feat_rgb_dir.shape[:4]
        if False:
            view_feat = self.view_fc(img_feat_rgb_dir[..., -4:])
            img_feat_rgb =  img_feat_rgb_dir[..., :-4] + view_feat
        else:
            img_feat_rgb =  img_feat_rgb_dir[..., :-4]
            
        var_feat = torch.var(img_feat_rgb, dim=1).view(B, 1, N_points, N_samples, self.feat_ch).repeat(1, S, 1, 1, 1)
        avg_feat = torch.mean(img_feat_rgb, dim=1).view(B, 1, N_points, N_samples, self.feat_ch).repeat(1, S, 1, 1, 1)

        feat = torch.cat([img_feat_rgb, var_feat, avg_feat], dim=-1)
        global_feat = self.global_fc(feat)
        agg_w = F.softmax(self.agg_w_fc(global_feat), dim=1)
        im_feat = (global_feat * agg_w).sum(dim=1)
        return self.fc(im_feat)

def weights_init(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight.data)
        if m.bias is not None:
            nn.init.zeros_(m.bias.data)
lib/networks/nerf/fields/mini_ibr_net.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional
from lib.networks.nerf.activations import get_activation

class NerfNet(nn.Module):
    def __init__(self,
                 xyz_in_dim: int,
                 view_in_dim: Optional[int] = None,
                 width: int = 256,
                 depth: int = 8,
                 view_depth: int = 1,
                 view_width: int = 128,
                 activation: str = 'relu',
                 skips: List[int] = [4],
                 **kwargs,
                 ):
        super(NerfNet, self).__init__()
        self.skips = skips
        self.backbone_nets = nn.ModuleList()
        segments = [0] + skips + [depth-1]
        for i in range(len(segments)-1):
            in_ch = xyz_in_dim if i == 0 else width + xyz_in_dim
            segment_net = nn.Sequential(*(
                [nn.Linear(in_ch, width), get_activation(activation)(inplace=True)] +
                [nn.Sequential(nn.Linear(width, width), get_activation(activation)(inplace=True)) for item in range(segments[i], segments[i+1])]
            ))
            self.backbone_nets.append(segment_net)
        self.sigma_net = nn.Linear(width, 1)
        self.rgb_net = nn.Sequential(*(
            [nn.Linear(view_in_dim + width, view_width), get_activation(activation)(inplace=True)] +
            [nn.Sequential(nn.Linear(view_width, view_width), get_activation(activation)(inplace=True)) for item in range(view_depth)] + 
            [nn.Linear(view_width, 3), nn.Sigmoid()]))

    def forward(self, xyz_encoding, view_encoding, only_geo=False):
        
        x = self.backbone_nets[0](xyz_encoding)
        for i, net in enumerate(self.backbone_nets[1:]):
            x = net(torch.cat([x, xyz_encoding], dim=-1))
        
        sigma = self.sigma_net(x)
        if only_geo: return None, sigma
        
        rgb = self.rgb_net(torch.cat([x, view_encoding], dim=-1))
        return rgb, sigma
lib/networks/nerf/fields/nerf_net.py
zju3dv/im4d
import torch
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from lib.config import cfg
from lib.networks.nerf.methods.mvibr import Network as MvibrNetwork
from lib.networks.nerf.methods.ibr.cost_volume_utils import CostVolumeUtils
from lib.networks.nerf.methods.ibr.cost_reg_net import CostRegNet
from lib.utils.base_utils import get_xywh_tensor

class Network(MvibrNetwork):

    def __init__(self, net_cfg, sample_level=0):
        super(Network, self).__init__(net_cfg, sample_level=sample_level)
        self.depth_inv = net_cfg.depth_inv
        self.num_planes = net_cfg.num_planes
        self.mvs_scale = net_cfg.mvs_scale
        self.cost_feat_input_ch = net_cfg.cost_feat_input_ch
        self.cost_feat_level = net_cfg.cost_feat_level
        self.cost_feat_scale = net_cfg.cost_feat_scale
        self.cost_reg_net = CostRegNet(net_cfg.cost_feat_input_ch, light_weight=net_cfg.cost_net_light_weight)

    def compute_cost_volume(self, batch, batch_share_info):
        render_h, render_w = batch['meta']['H'][0].item(), batch['meta']['W'][0].item()
        render_h_w = [render_h, render_w]
        cost_volume, t_vals, t2z_func = CostVolumeUtils.build_cost_volume(
                render_h_w,
                batch['near_far'],
                batch['tar_ext'],
                batch['tar_ixt'],
                batch['src_exts'],
                batch['src_ixts'],
                batch_share_info['img_feat'][self.cost_feat_level],
                self.depth_inv,
                self.num_planes,
                perturb=self.training,
                tar_scale=self.mvs_scale,
                src_scale=self.cost_feat_scale)
        prob_volume = self.cost_reg_net(cost_volume)
        batch_share_info.update({
            't2z_func': t2z_func,
            't_vals': t_vals,
            'prob_volume': prob_volume
        })

    def predict_depth(self, batch, batch_share_info):
        prob_volume, t_vals = batch_share_info['prob_volume'], batch_share_info['t_vals']
        prob_volume = F.softmax(F.relu(prob_volume), dim=1)
        t_mean = (t_vals * prob_volume).sum(dim=1)
        t_std =  (prob_volume * (t_vals - t_mean.unsqueeze(1))**2).sum(1)
        t_std = t_std.clamp_min(1e-10).sqrt()
        t_min = torch.clamp(t_mean-t_std, min=0, max=None)
        t_max = torch.clamp(t_mean+t_std, min=None, max=1.)
        batch_share_info.update({'t_mean_std_min_max': torch.cat([t_mean[:, None], t_std[:, None], t_min[:, None], t_max[:, None]], dim=1)})

    def forward_feat(self, batch, batch_share_info):
        super().forward_feat(batch, batch_share_info)
        self.compute_cost_volume(batch, batch_share_info)
        self.predict_depth(batch, batch_share_info) # BxCxnum_planesxHxW, Bx1xnum_planesxHxW

lib/networks/nerf/methods/enerf.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvBnReLU3D(nn.Module):
    def __init__(self, in_channels, out_channels,
                 kernel_size=3, stride=1, pad=1,
                 norm_act=nn.BatchNorm3d):
        super(ConvBnReLU3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels,
                              kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = norm_act(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

class CostRegNet(nn.Module):
    def __init__(self, 
                 in_channels=32, 
                 light_weight=False,
                 norm_act=nn.InstanceNorm3d):
        super(CostRegNet, self).__init__()
        self.light_weight = light_weight
        # light_weight is used if the resolution of the cost volume is H/8 x W/8
        self.conv0 = ConvBnReLU3D(in_channels, 8, norm_act=norm_act)
        self.conv1 = ConvBnReLU3D(8, 16, stride=2, norm_act=norm_act)
        self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act)

        self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act)
        self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act)

        self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act)
        self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act)

        self.conv7 = nn.Sequential(
            nn.ConvTranspose3d(64, 32, 3, padding=1, output_padding=1,
                               stride=2, bias=False),
            norm_act(32))

        self.conv9 = nn.Sequential(
            nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1,
                               stride=2, bias=False),
            norm_act(16))

        self.conv11 = nn.Sequential(
            nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1,
                               stride=2, bias=False),
            norm_act(8))
        self.depth_conv = nn.Sequential(nn.Conv3d(8, 1, 3, padding=1, bias=False))
        
    def add_up(self, x1, x2):
        # import ipdb; ipdb.set_trace()
        diffD = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        diffX = x2.size()[4] - x1.size()[4] 
        # x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2, diffD // 2, diffD - diffD // 2))
        if diffD !=0 or diffY != 0 or diffX != 0: x2 = F.interpolate(x2, size=x1.shape[-3:], mode='trilinear')
        return x1 + x2

    def forward(self, x):
        conv0 = self.conv0(x)
        conv2 = self.conv2(self.conv1(conv0))
        conv4 = self.conv4(self.conv3(conv2))
        if self.light_weight: x = conv4
        else: x = self.conv6(self.conv5(conv4)); x = conv4 + self.add_up(conv4, self.conv7(x))
        x = self.add_up(conv2, self.conv9(x))
        x = self.add_up(conv0, self.conv11(x))
        depth = self.depth_conv(x)
        return depth.squeeze(1)
lib/networks/nerf/methods/ibr/cost_reg_net.py
zju3dv/im4d
from typing import List
import torch
from kornia.utils import create_meshgrid
import torch.nn.functional as F
import numpy as np

class CostVolumeUtils:
    def homo_warp(src_feat, proj_mat, depth_values):
        B, D, H_T, W_T = depth_values.shape
        C, H_S, W_S = src_feat.shape[1:]
        device = src_feat.device

        R = proj_mat[:, :, :3] # (B, 3, 3)
        T = proj_mat[:, :, 3:] # (B, 3, 1)
        # create grid from the ref frame
        ref_grid = create_meshgrid(H_T, W_T, normalized_coordinates=False,
                               device=device) # (1, H, W, 2)
        ref_grid = ref_grid.permute(0, 3, 1, 2) # (1, 2, H, W)
        ref_grid = ref_grid.reshape(1, 2, H_T*W_T) # (1, 2, H*W)
        ref_grid = ref_grid.expand(B, -1, -1) # (B, 2, H*W)
        ref_grid = torch.cat((ref_grid, torch.ones_like(ref_grid[:,:1])), 1) # (B, 3, H*W)
        ref_grid_d = ref_grid.repeat(1, 1, D) # (B, 3, D*H*W)
        src_grid_d = R @ ref_grid_d + T/depth_values.view(B, 1, D*H_T*W_T)
        del ref_grid_d, ref_grid, proj_mat, R, T, depth_values # release (GPU) memory
        src_grid = src_grid_d[:, :2] / torch.clamp_min(src_grid_d[:, 2:], 1e-6) # divide by depth (B, 2, D*H*W)
        # del src_grid_d
        src_grid[:, 0] = (src_grid[:, 0])/((W_S - 1) / 2) - 1 # scale to -1~1
        src_grid[:, 1] = (src_grid[:, 1])/((H_S - 1) / 2) - 1 # scale to -1~1
        src_grid = src_grid.permute(0, 2, 1) # (B, D*H*W, 2)
        src_grid = src_grid.view(B, D, H_T*W_T, 2)
    
        warped_src_feat = F.grid_sample(src_feat, src_grid,
                                        mode='bilinear', padding_mode='border',
                                        align_corners=True) # (B, C, D, H*W)
        warped_src_feat = warped_src_feat.view(B, C, D, H_T, W_T)
        src_grid = src_grid.view(B, D, H_T, W_T, 2)
        if torch.isnan(warped_src_feat).isnan().any():
            __import__('ipdb').set_trace()
        return warped_src_feat, src_grid
    @staticmethod 
    def get_proj_mats(src_ext, src_ixt, tar_ext, tar_ixt, src_scale, tar_scale):
        B, S_V = src_ext.shape[:2]
        src_ixt = src_ixt.clone()
        src_ixt[:, :, :2] *= src_scale
        src_projs = src_ixt @ src_ext[:, :, :3]
    
        tar_ixt = tar_ixt.clone()
        tar_ixt[:, :2] *= tar_scale
        tar_projs = tar_ixt @ tar_ext[:, :3]
        tar_ones = torch.zeros((B, 1, 4)).to(tar_projs.device)
        tar_ones[:, :, 3] = 1
        tar_projs = torch.cat((tar_projs, tar_ones), dim=1)
        tar_projs_inv = torch.inverse(tar_projs)
    
        src_projs = src_projs.view(B, S_V, 3, 4)
        tar_projs_inv = tar_projs_inv.view(B, 1, 4, 4)
    
        proj_mats = src_projs @ tar_projs_inv
        return proj_mats
    
    @staticmethod
    def build_cost_volume(
        render_h_w: List[int], # [render_h, render_w]
        near_far: torch.Tensor, # Bx2
        tar_ext: torch.Tensor, # Bx4x4
        tar_ixt: torch.Tensor, # Bx3x3
        src_exts: torch.Tensor, # Bx4x4
        src_ixts: torch.Tensor, # Bx3x3
        img_feat: torch.Tensor, # BxSxCxHxW
        depth_inv: bool = False,
        num_planes: int = 64,
        tar_scale: float = 0.25,
        src_scale: float = 0.25,
        perturb: bool = False,
    ):
        B = len(near_far)
        h, w = render_h_w
        h, w = int(h * tar_scale), int(w * tar_scale)
        t_vals = torch.linspace(0., 1., steps=num_planes, device=near_far.device, dtype=torch.float32).view(1, -1, 1, 1).repeat(B, 1, h, w)
        if perturb:
            interval = 1/num_planes
            t_vals = np.random.random() * interval + t_vals 
        if depth_inv:
            t2z_func = lambda x, near_far: 1/(1/near_far[..., 0:1] + x*(1/near_far[..., 1:2] - 1/near_far[..., 0:1]))
        else:
            t2z_func = lambda x, near_far: near_far[..., 0:1] + x*(near_far[..., 1:2] - near_far[..., 0:1])
        B, D, H, W = t_vals.shape
        depth_values = t2z_func(t_vals.reshape(B, -1), near_far).reshape(B, D, H, W)
        proj_mats = CostVolumeUtils.get_proj_mats(src_ext=src_exts, 
                                                  src_ixt=src_ixts,
                                                  tar_ext=tar_ext,
                                                  tar_ixt=tar_ixt,
                                                  src_scale=src_scale, # ixt -> src_feat
                                                  tar_scale=tar_scale) # ixt -> depth_values)
        S = src_exts.shape[1]
        volume_sum = 0
        volume_sq_sum = 0
        for s in range(S):
            feature_s = img_feat[:, s]
            proj_mat = proj_mats[:, s]
            warped_volume, _ = CostVolumeUtils.homo_warp(feature_s, proj_mat, depth_values)
            volume_sum = volume_sum + warped_volume
            volume_sq_sum = volume_sq_sum + warped_volume ** 2
        volume_variance = volume_sq_sum.div_(S).sub_(volume_sum.div_(S).pow_(2))
        return volume_variance, t_vals, t2z_func
        
    
lib/networks/nerf/methods/ibr/cost_volume_utils.py
zju3dv/im4d
from lib.networks.nerf.methods.ibr.feature_nets.featurenet_unet import FeatureNet as ENeRFUnet

nets_dict = {
    'enerf_unet': ENeRFUnet
}
def get_feature_net(net_cfg, **kwargs):
    return nets_dict[net_cfg.type](**net_cfg)
lib/networks/nerf/methods/ibr/feature_nets/__init__.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
import importlib

class ConvBnReLU(nn.Module):
    def __init__(self, in_channels, out_channels,
                 kernel_size=3, stride=1, pad=1,
                 norm_act=nn.BatchNorm2d):
        super(ConvBnReLU, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels,
                              kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = norm_act(out_channels)
        self.relu = nn.ReLU(inplace=True)
    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

class ConvBnReLU3D(nn.Module):
    def __init__(self, in_channels, out_channels,
                 kernel_size=3, stride=1, pad=1,
                 norm_act=nn.BatchNorm3d):
        super(ConvBnReLU3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels,
                              kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = norm_act(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))


class FeatureNet(nn.Module):
    def __init__(self, input_ch=3, output_ch=None, norm_act=nn.InstanceNorm2d, **kwargs):
        super(FeatureNet, self).__init__()
        self.conv0 = nn.Sequential(
                        ConvBnReLU(input_ch, 8, 3, 1, 1, norm_act=norm_act),
                        ConvBnReLU(8, 8, 3, 1, 1, norm_act=norm_act))
        self.conv1 = nn.Sequential(
                        ConvBnReLU(8, 16, 5, 2, 2, norm_act=norm_act),
                        ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act))
        self.conv2 = nn.Sequential(
                        ConvBnReLU(16, 32, 5, 2, 2, norm_act=norm_act),
                        ConvBnReLU(32, 32, 3, 1, 1, norm_act=norm_act))
        self.toplayer = nn.Conv2d(32, 32, 1)
        self.lat1 = nn.Conv2d(16, 32, 1)
        self.lat0 = nn.Conv2d(8, 32, 1)
        self.smooth1 = nn.Conv2d(32, 16, 3, padding=1)
        self.smooth0 = nn.Conv2d(32, 8 if output_ch is None else output_ch, 3, padding=1)

    def _upsample_add(self, x, y):
        x_h, x_w, y_h, y_w = x.shape[2], x.shape[3], y.shape[2], y.shape[3]
        if 2*x_h == y_h and 2*x_w == y_w:
            return F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True) + y
        else:
            return F.interpolate(x, size=(y_h, y_w), mode='bilinear', align_corners=True) + y

    def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)
        feat2 = self.toplayer(conv2)
        feat1 = self._upsample_add(feat2, self.lat1(conv1))
        feat0 = self._upsample_add(feat1, self.lat0(conv0))
        feat1 = self.smooth1(feat1)
        feat0 = self.smooth0(feat0)
        return [feat0, feat1, feat2]
lib/networks/nerf/methods/ibr/feature_nets/featurenet_unet.py
zju3dv/im4d
import torch
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from lib.config import cfg
from lib.networks.nerf.network import Network as NeRFNetwork
from lib.networks.nerf.methods.ibr.feature_nets import get_feature_net

class Network(NeRFNetwork):
    def __init__(self, net_cfg, sample_level=0):
        super(Network, self).__init__(net_cfg, sample_level=sample_level)
        if sample_level == 0: self.feature_net = get_feature_net(net_cfg.feature_net)
        
    def forward_feat(self, batch, batch_share_info):
        B, S, C, H, W = batch['src_inps'].shape
        img_feat = self.feature_net(batch['src_inps'].reshape(B*S, C, H, W)) # list of feature maps, sample_level
        img_feat = [item.reshape(B, S, item.shape[-3], item.shape[-2], item.shape[-1]) for item in img_feat]
        batch_share_info['img_feat'] = img_feat
        
    def forward(self, batch):
        batch_share_info = {}
        # Im4D Sec. 3.3
        if self.training and batch['step'] >= cfg.joint_iters and batch['step'] % cfg.finetune_iters_per != 0:
            with torch.no_grad():
                self.forward_feat(batch, batch_share_info)
        else: self.forward_feat(batch, batch_share_info)
        return super().forward(batch, batch_share_info)

lib/networks/nerf/methods/mvibr.py
zju3dv/im4d
import torch
import torch.nn as nn
from torch.nn import functional as F
from os.path import join
import numpy as np
from lib.config import cfg, logger
if cfg.fast_render:
    import nerfacc
from lib.networks.nerf.encoders import get_encoder
from lib.networks.nerf.samplers import get_sampler
from lib.networks.nerf.fields import get_net
from lib.utils.ops.activations import init_density_activation
from lib.networks.nerf.activations import get_activation
from lib.utils.data_utils import get_ray_near_far_torch
from lib.utils.parallel_utils import chunkify, cat_dict, cat_tensor
from lib.utils.im4d.fastrend_utils import FastRendUtils
from lib.utils.im4d.im4d_utils import Im4DUtils
from lib.utils import data_utils

class Network(nn.Module):
    def __init__(self, net_cfg, sample_level=0):
        super(Network, self).__init__()
        # encoding
        self.net_cfg = net_cfg
        self.sample_level = sample_level
        self.only_geo = net_cfg.get('only_geo', False) # for some proposal network
        bbox = net_cfg.get('bounds') if 'bounds' in net_cfg else cfg.get('bounds')
        self.register_buffer('bbox', torch.tensor(bbox, dtype=torch.float32).reshape(2, 3))
        self.xyz_encoder = get_encoder(net_cfg.xyz_encoding)
        self.dir_encoder = get_encoder(net_cfg.dir_encoding)

        # proposal network
        self.sampler = get_sampler(cfg.samplers[sample_level], global_cfg=cfg, sample_level=sample_level)

        # radiance_network
        self.net = get_net(net_cfg.net,
                           xyz_in_dim = self.xyz_encoder.get_out_dim(),
                           view_in_dim = self.dir_encoder.get_out_dim() if net_cfg.get('dir_encoding', False) else None,
                           **net_cfg.net)
        # density activation
        self.density_act = init_density_activation(net_cfg.get('density_act', 'trunc_exp'))
        self.ignore_dist = net_cfg.get('ignore_dist', False) # for multi-view image-based rendering methods
        if cfg.fast_render and sample_level == 0: self.prepare_fast_render(net_cfg)
        
        logger.debug('network-level{}: xyz_encoder {}'.format(sample_level, str(self.xyz_encoder)))
        # logger.debug('network-level{}: dir_encoder {}'.format(sample_level, str(self.dir_encoder)))
        # logger.debug('network-level{}: decoder_net {}'.format(sample_level, str(self.net)))
        logger.debug('network-level{}: density_act {}'.format(sample_level, str(net_cfg.get('density_act', 'trunc_exp'))))
        logger.debug('network-level{}: ignore_dist {}'.format(sample_level, str(self.ignore_dist))) # Multi-view image-based rendering methods need to ignore the distance
        
    def prepare_fast_render(self, net_cfg):
        binarys = np.load(join(cfg.grid_dir, 'binarys.npz'))['arr_0']
        bounds = np.load(join(cfg.grid_dir, 'bounds.npz'))['arr_0']
        self.occ_grids, self.render_step_size = {}, {}
        for frame_id in range(*cfg.test_dataset.frame_sample):
            occ_grid = nerfacc.OccGridEstimator(torch.tensor(bounds[frame_id].reshape(-1)).float(), binarys[frame_id].shape)
            occ_grid.binaries = torch.tensor(binarys[frame_id]).bool()[None]
            del occ_grid.grid_coords
            del occ_grid.grid_indices
            del occ_grid.occs
            occ_grid.cuda()
            self.occ_grids[frame_id] = occ_grid
            self.render_step_size[frame_id] = np.linalg.norm(bounds[frame_id][1] - bounds[frame_id][0]) / cfg.num_max_samples

    @staticmethod
    def volume_rendering(density, z_vals, ray_dir, act, ignore_dist=False):
        '''
        density: BxN_raysxN_samplesx1
        z_vals: Bx(N_rays+1)xN_samples
        '''
        if density.shape[2] == 1: return F.softmax(density, dim=-1)
        if act is not None: density = act(density)

        if ignore_dist: delta_density = density
        else:
            delta = z_vals[..., 1:] - z_vals[..., :-1]
            delta = delta * ray_dir.norm(dim=-1, keepdim=True)
            delta_mask = delta > 1e-6
            delta = delta[delta_mask]

            delta_density = torch.zeros_like(density)
            delta_density[delta_mask] = delta * density[delta_mask]

        alpha = 1 - torch.exp(-delta_density)

        T = torch.cumprod(1 - alpha + 1e-10, dim=-1)[..., :-1]
        T = torch.cat([torch.ones_like(alpha[..., 0:1]), T], dim=-1)
        weights = alpha * T
        return weights
    
    @staticmethod
    def compute_ndc_rays(rays_o, rays_d, near_far, focal, H, W, near=1.):
        t = -(near + rays_o[...,2]) / rays_d[...,2]
        rays_o = rays_o + t[...,None] * rays_d

        o0 = -1./(W/(2.*focal)) * rays_o[...,0] / rays_o[...,2]
        o1 = -1./(H/(2.*focal)) * rays_o[...,1] / rays_o[...,2]
        o2 = 1. + 2. * near / rays_o[...,2]
        
        d0 = -1./(W/(2.*focal)) * (rays_d[...,0]/rays_d[...,2] - rays_o[...,0]/rays_o[...,2])
        d1 = -1./(H/(2.*focal)) * (rays_d[...,1]/rays_d[...,2] - rays_o[...,1]/rays_o[...,2])
        d2 = -2. * near / rays_o[...,2]

        rays_o = torch.stack([o0,o1,o2], -1)
        rays_d = torch.stack([d0,d1,d2], -1)
        # rays_d = rays_d / np.linalg.norm(rays_d, axis=-1)[..., None]
        rays_d = rays_d / 2.
        near_far[..., 0] = 0. 
        near_far[..., 1] = 1.99 
        return rays_o, rays_d, near_far
        
    @staticmethod
    def compute_rays(coords, ext_inv, ixt_inv, time, h_w, near_far):
        rays_o = ext_inv[:, :3, 3][:, None, :].repeat(1, coords.shape[1], 1)
        rays_d = coords.float() @ ixt_inv.transpose(-1, -2) @ ext_inv[:, :3, :3].transpose(-1, -2)
        near_far = near_far[:, None, :].repeat(1, coords.shape[1], 1)
        rays_t = time[:, None, None].repeat(1, coords.shape[1], 1)
        uv = coords[:, :, :2] / (h_w[..., [1, 0]][:, None, :] - 1)
        if cfg.get('ndc', False): rays_o, rays_d, near_far = Network.compute_ndc_rays(rays_o, rays_d, near_far, cfg.ndc_FOCAL, cfg.ndc_H, cfg.ndc_W)
        return torch.cat([rays_o, rays_d, near_far, rays_t, uv], dim=-1)

    @staticmethod
    def prefilter_rays(batch, rays):
        assert(len(rays) == 1)
        _near, _far, mask_at_box = get_ray_near_far_torch(batch['bounds'][0], rays[0, :, :3], rays[0, :, 3:6])
        rays = rays[0][mask_at_box][None]
        batch.update({'mask_at_box': mask_at_box[None].to(torch.uint8),
                      'rgb': batch['rgb'][0][mask_at_box][None],
                      'msk': batch['msk'][0][mask_at_box][None]})
        return rays

    def render_rays(self, rays, **kwargs):
        '''
        rays: BxN_raysx8
        '''
        ray_dir = rays[..., 3:6]
        points, z_vals, ret_info = self.sampler(rays, **kwargs)
        # BxN_raysxN_samplesxC

        viewdir = ray_dir / ray_dir.norm(dim=-1, keepdim=True)
        time = rays[..., 8:9][:, :, None].repeat(1, 1, points.shape[-2], 1)
        xyz_encoding = self.xyz_encoder(points, viewdir=viewdir, sample_level=self.sample_level, bbox=self.bbox, time=time, **kwargs) # feature encoding may need viewdir

        if self.only_geo:  _, sigma = self.net(xyz_encoding, None, only_geo=True)
        else:
            dir_encoding = self.dir_encoder(viewdir)
            dir_encoding = dir_encoding[..., None, :].repeat(1, 1, points.shape[-2], 1) if dir_encoding is not None else None
            rgb, sigma = self.net(xyz_encoding, dir_encoding)

        vr_weights = self.volume_rendering(density=sigma[..., 0],
                                           z_vals=z_vals,
                                           ray_dir=ray_dir,
                                           act=self.density_act,
                                           ignore_dist=self.ignore_dist)

        output = {}
        for k in ret_info:
            if 'weights' in k and not self.training: continue # do not need weights in testing
            output.update({k: ret_info[k]})

        acc_map = vr_weights.sum(dim=-1)
        if cfg.depth_method == 'expected': depth_map = (vr_weights * (z_vals[..., :-1] + z_vals[..., 1:]) * 0.5).sum(dim=-1) # / (acc_map + 1e-6)
        elif cfg.depth_method == 'median': 
            depth_map = Im4DUtils.render_median_depth(vr_weights, z_vals)
            pts = rays[..., :3] + depth_map[..., None] * ray_dir
            output.update({'pts_{}'.format(self.sample_level): pts})
        else: import ipdb; ipdb.set_trace()
        # TODO: median depth
        output.update({
            'depth_map_{}'.format(self.sample_level): depth_map,
            'acc_map_{}'.format(self.sample_level): acc_map,
            'z_vals_{}'.format(self.sample_level): z_vals,
        })
        if not self.only_geo:
            # if rgb.shape[-1] == 6: rgb_map = torch.cat([(vr_weights[..., None] * rgb[..., :3]).sum(dim=-2), (vr_weights[..., None] * rgb[..., 3:]).sum(dim=-2)], dim=-1)
            # else: rgb_map = (vr_weights[..., None] * rgb).sum(dim=-2)
            rgb_map = (vr_weights[..., None] * rgb).sum(dim=-2)
            if cfg.white_bkgd: rgb_map = rgb_map + (1-acc_map[..., None])
            output.update({'rgb_map_{}'.format(self.sample_level): rgb_map})

        if self.training or self.sample_level > 0: output.update({'weights_{}'.format(self.sample_level): vr_weights})
        # include the weights for training or sampling at a higher level
        return output

    def render_rays_fast(self, rays, **kwargs):
        # only support for Im4D
        assert cfg.task == 'im4d'
        rays_o, rays_d = rays[0, :, :3], rays[0, :, 3:6]
        viewdir = rays_d / rays_d.norm(dim=-1, keepdim=True)
        n_rays = len(rays_o)
        time = rays[0, :, 8:9]
        frame_id = kwargs['batch']['meta']['frame_id'].item() # batch_size = 1

        # raymarching
        ray_indices, t_starts, t_ends = self.occ_grids[frame_id].sampling(
            rays_o,
            rays_d,
            render_step_size=self.render_step_size[frame_id]
        )
        if len(ray_indices) == 0:
            acc_map = torch.zeros_like(rays[..., 0])
            depth_map = torch.zeros_like(rays[..., 0])
            rgb_map = torch.zeros_like(rays[..., :3])
            if cfg.white_bkgd: rgb_map = rgb_map + (1-acc_map[..., None])
            output = {}
            output.update({
                'depth_map_{}'.format(self.sample_level): depth_map,
                'acc_map_{}'.format(self.sample_level): acc_map,
                'rgb_map_{}'.format(self.sample_level): rgb_map
            })
            return output

        # compute vr weights
        t_mid = (t_starts + t_ends) / 2.0
        points = rays_o[ray_indices] + t_mid[:, None] * rays_d[ray_indices]
        v = viewdir[ray_indices]
        t = time[0].item() * torch.ones_like(v[..., :1])
        xyz_encoding = self.xyz_encoder(points[None, :, None], only_geo=True, sample_level=self.sample_level, bbox=self.bbox, time=t[None, :, None], **kwargs)
        _, density = self.net([xyz_encoding], None, only_geo=True)
        density = self.density_act(density)
        delta_density = density * self.render_step_size[frame_id] * rays_d[ray_indices].norm(dim=-1)[None, :, None, None]
        alpha = 1 - torch.exp(-delta_density)
        weights, trans = nerfacc.render_weight_from_alpha(alpha[0, :, 0, 0], ray_indices=ray_indices, n_rays=n_rays)

        # filter weight_threshold
        rgb_msk = weights >= cfg.weight_thresh # only compute rgb for weighted samples
        if rgb_msk.sum() == 0:
            acc_map = torch.zeros_like(rays[..., 0])
            depth_map = torch.zeros_like(rays[..., 0])
            rgb_map = torch.zeros_like(rays[..., :3])
            if cfg.white_bkgd: rgb_map = rgb_map + (1-acc_map[..., None])
            output = {}
            output.update({
                'depth_map_{}'.format(self.sample_level): depth_map,
                'acc_map_{}'.format(self.sample_level): acc_map,
                'rgb_map_{}'.format(self.sample_level): rgb_map
            })
            return output

        # compute rgb fields
        ibr_feat = self.xyz_encoder.feature_projection(points[rgb_msk][None, :, None], viewdir=v[rgb_msk][None, :, None], time=t[rgb_msk][None, :, None], sample_level=self.sample_level, **kwargs)
        rgb, _ = self.net.enerf_net(ibr_feat, v[rgb_msk][None, :, None][..., :3])
        rgb_map = nerfacc.accumulate_along_rays(
            weights=weights[rgb_msk],
            values=rgb[0, :, 0, :],
            ray_indices=ray_indices[rgb_msk],
            n_rays=n_rays,
        )[None, :]
        acc_map = nerfacc.accumulate_along_rays(
            weights=weights[rgb_msk],
            values=None,
            ray_indices=ray_indices[rgb_msk],
            n_rays=n_rays,
        )[None, :, 0]
        depth_map = nerfacc.accumulate_along_rays(
            weights=weights[rgb_msk],
            values=t_mid[rgb_msk][:, None],
            ray_indices=ray_indices[rgb_msk],
            n_rays=n_rays,
        )[None, :, 0]

        if cfg.white_bkgd: rgb_map = rgb_map + (1-acc_map[..., None])
        output = {}
        output.update({
            'depth_map_{}'.format(self.sample_level): depth_map,
            'acc_map_{}'.format(self.sample_level): acc_map,
            'rgb_map_{}'.format(self.sample_level): rgb_map
        })
        return output

    def forward(self, batch, batch_share_info={}):
        rays = self.compute_rays(batch['coords'], batch['tar_ext_inv'], batch['tar_ixt_inv'], batch['time'], batch['h_w'], batch['near_far'])
        if not self.training and 'bounds' in batch: rays = self.prefilter_rays(batch, rays)
        render_rays_func = self.render_rays_fast if cfg.fast_render else self.render_rays
        ret = chunkify(render_rays_func, cat_dict, [rays], chunk_dim=1, chunk_size=cfg.chunk_size, batch=batch, batch_share_info=batch_share_info)
        ret.update({'model': self}) # for regularization
        return ret

    def inference_density(self, wpts, **kwargs):
        wpts = wpts[None, :, None]
        time = torch.ones_like(wpts[..., :1]) * kwargs['batch']['time'][0]
        xyz_encoding = self.xyz_encoder(wpts, only_geo=True, sample_level=self.sample_level, bbox=self.bbox, time=time, **kwargs)
        _, sigma = self.net([xyz_encoding], None, only_geo=True)
        return sigma.reshape(-1, 1)

    def cache_grid(self, batch, save_ply=False):
        assert(len(batch['bounds']) == 1)
        wpts = FastRendUtils.prepare_wpts(batch['bounds'][0], cfg.grid_resolution, cfg.subgrid_resolution)
        # wpts = data_utils.get_space_points_torch(batch['bounds'][0], cfg.grid_resolution)
        shape = wpts.shape
        sigma = chunkify(self.inference_density, cat_tensor, [wpts.reshape(-1, 3)], chunk_size=cfg.chunk_size, chunk_dim=0, batch=batch)
        sigma = sigma.reshape(shape[:-1])
        binary = (sigma > cfg.sigma_thresh).float().sum(dim=-1) > 0
        if cfg.padding_grid:
            structuring_element = torch.ones(1, 1, 3, 3, 3, device=binary.device)
            dilated_tensor = F.conv3d(binary[None, None].float(), structuring_element.float(), padding=1) > 0
            binary = dilated_tensor[0, 0].float()
        if save_ply: data_utils.save_mesh_with_extracted_fields(join(cfg.grid_dir, 'meshes', '{:06d}.ply'.format(batch['meta']['frame_id'][0].item())), binary.cpu().numpy(), batch['bounds'][0].cpu().numpy(), grid_resolution=cfg.grid_resolution)
        return binary.cpu().numpy().astype(np.bool), batch['bounds'][0].cpu().numpy()


lib/networks/nerf/network.py
zju3dv/im4d
from lib.networks.nerf.samplers.cascade_sampler import CascadeSampler
from lib.networks.nerf.samplers.uniform_sampler import UniformSampler
from lib.networks.nerf.samplers.dispuniform_sampler import DispUniformSampler
from lib.networks.nerf.samplers.mvs_sampler import MvsSampler

samplers_dict = {
    'mvs_sampler': MvsSampler,
    'cascade_sampler': CascadeSampler,
    'uniform_sampler': UniformSampler,
    'dispuniform_sampler': DispUniformSampler,
}

def get_sampler(sampler_cfg, global_cfg=None, sample_level=0):
    return samplers_dict[sampler_cfg['type']](global_cfg=global_cfg, sample_level=sample_level+1, **sampler_cfg)
lib/networks/nerf/samplers/__init__.py
zju3dv/im4d
import torch
import torch.nn as nn
import imp
import numpy as np
from lib.networks.nerf.samplers.utils import pdf_sample
from lib.networks.nerf.samplers.uniform_sampler import UniformSampler

from lib.config.yacs import CfgNode as CN


class CascadeSampler(UniformSampler):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self,
        global_cfg=None,
        **kwargs,
    ) -> None:
        super(CascadeSampler, self).__init__(**kwargs)
        module = global_cfg.network_module
        path = global_cfg.network_path
        self.network = imp.load_source(module, path).Network(CN(kwargs), sample_level=kwargs['sample_level'])
        self.sample_level = kwargs['sample_level']
        self.num_samples = kwargs['num_samples']
        
        self.N_anneal_steps = kwargs['N_anneal_steps']
        self.N_anneal_slope = kwargs['N_anneal_slope']
        self.bias = lambda x, b: (b * x) / ((b - 1) * x + 1)
        self.include_last_level = kwargs.get('include_input', False)
        
    def compute_t_vals(self, rays, **kwargs):
        ret_info = self.network.render_rays(rays, **kwargs)
        pdf_weights = ret_info['weights_{}'.format(self.sample_level)]
        if self.training:
            assert('step' in kwargs['batch'])
            step = kwargs['batch']['step']
            train_frac = np.clip(step / self.N_anneal_steps, 0, 1)
            anneal = self.bias(train_frac, self.N_anneal_slope)
            pdf_weights = torch.pow(pdf_weights, anneal)
        with torch.no_grad():
            t_vals = pdf_sample(ret_info[f't_vals_{self.sample_level}'], pdf_weights, self.training, self.num_samples)
        if self.include_last_level:
            t_vals = torch.cat([t_vals, ret_info[f't_vals_{self.sample_level}'][..., 1:]], dim=-1)
            t_vals = torch.sort(t_vals, dim=-1)[0]
        return t_vals, ret_info
    
    def compute_z_vals(self, t_vals, rays, share_info, **kwargs):
        near, far = rays[..., 6:7], rays[..., 7:8]
        t2z_func = share_info[f't2z_func']
        return t2z_func(t_vals, near, far)
lib/networks/nerf/samplers/cascade_sampler.py
zju3dv/im4d
import torch
import torch.nn as nn
import imp
from lib.networks.nerf.samplers.uniform_sampler import UniformSampler

from lib.config.yacs import CfgNode as CN
from lib.config import logger


class DispUniformSampler(UniformSampler):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self,
        global_cfg=None,
        **kwargs,
    ) -> None:
        super(DispUniformSampler, self).__init__(**kwargs)
        self.num_samples = kwargs['num_samples']
        self.sample_level = kwargs['sample_level']
    
    def compute_z_vals(self, t_vals, rays, share_info, **kwargs):
        near, far = rays[..., 6:7], rays[..., 7:8]
        z_vals_inv = 1/near * (1. - t_vals) + 1/far * t_vals
        z_vals = 1/z_vals_inv
        t2z_func = lambda t_vals, near, far: 1/(1/near * (1. - t_vals) + 1/far * t_vals)
        share_info.update({'t2z_func': t2z_func})
        return z_vals
lib/networks/nerf/samplers/dispuniform_sampler.py
zju3dv/im4d
import torch
import torch.nn as nn
import torch.nn.functional as F
import imp
import numpy as np
from lib.networks.nerf.samplers.utils import pdf_sample
from lib.networks.nerf.samplers.uniform_sampler import UniformSampler

from lib.config.yacs import CfgNode as CN


class MvsSampler(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self,
        global_cfg=None,
        **kwargs,
    ) -> None:
        super(MvsSampler, self).__init__()
        self.num_samples = kwargs['num_samples']
        
        
    def compute_t_vals(self, rays, **kwargs):
        sh = rays.shape[:-1]
        t_vals = torch.linspace(0., 1., steps=self.num_samples+1, device=rays.device)
        t_vals = t_vals.expand(sh + (self.num_samples+1,))
        return t_vals, {}
    
    def compute_z_vals(self, t_vals, rays, share_info, **kwargs):
        near, far = rays[..., 6:7], rays[..., 7:8]
        z_vals = near * (1. - t_vals) + far * t_vals
        return z_vals
        
    def forward(self, rays, **kwargs):
        t_mean_std_min_max = kwargs['batch_share_info']['t_mean_std_min_max']
        t2z_func = kwargs['batch_share_info']['t2z_func']
        uv = rays[..., -2:]
        t_h_w = t_mean_std_min_max.shape[2:]
        render_h_w = [kwargs['batch']['meta']['H'][0].item(), kwargs['batch']['meta']['W'][0].item()]
        if t_h_w[0] == render_h_w[0] and t_h_w[1] == render_h_w[1]:
            x, y = (uv[..., 0] * (t_h_w[1] - 1)).long(), (uv[..., 1] * (t_h_w[0] -1 )).long()
            ray_near_far = t_mean_std_min_max[:, :, y, x][:, :, 0].transpose(-1, -2)
        else:
            ray_near_far = F.grid_sample(t_mean_std_min_max, uv[:, None]*2-1., align_corners=True, mode='bilinear')[:, :, 0].transpose(-1, -2)
        uniform_vals, share_info = self.compute_t_vals(rays, **kwargs)
        t_vals = ray_near_far[..., 2:3] * (1 - uniform_vals) + ray_near_far[..., 3:4] * uniform_vals
        B, num_rays, num_samples = t_vals.shape
        z_vals = t2z_func(t_vals.reshape(B, -1), kwargs['batch']['near_far']).reshape(B, num_rays, num_samples)
        z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
        rays_o, rays_d = rays[..., :3], rays[..., 3:6]
        points = rays_o[..., None, :] + rays_d[..., None, :] * z_vals_mid[..., :, None]
        # t2z_func = lambda t_vals, near, far: near * (1. - t_vals) + far * t_vals
        share_info.update({'t_vals': t_vals})
        return points, z_vals, share_info
lib/networks/nerf/samplers/mvs_sampler.py
zju3dv/im4d
import torch
import torch.nn as nn
import imp

from lib.config.yacs import CfgNode as CN
from lib.config import logger


class UniformSampler(nn.Module):
    """Multi-scale sinusoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.

    Args:
        in_dim: Input dimension of tensor
        num_frequencies: Number of encoded frequencies per axis
        min_freq_exp: Minimum frequency exponent
        max_freq_exp: Maximum frequency exponent
        include_input: Append the input coordinate to the encoding
    """

    def __init__(
        self,
        global_cfg=None,
        **kwargs,
    ) -> None:
        super(UniformSampler, self).__init__()
        self.num_samples = kwargs['num_samples']
        self.sample_level = kwargs['sample_level']
        
    
    def compute_t_vals(self, rays, **kwargs):
        sh = rays.shape[:-1]
        t_vals = torch.linspace(0., 1., steps=self.num_samples+1, device=rays.device)
        t_vals = t_vals.expand(sh + (self.num_samples+1,))
        if self.training:
            interval = 1/self.num_samples
            t_vals = torch.rand_like(t_vals) * interval + t_vals 
        return t_vals, {}
    
    def compute_z_vals(self, t_vals, rays, share_info, **kwargs):
        near, far = rays[..., 6:7], rays[..., 7:8]
        z_vals = near * (1. - t_vals) + far * t_vals
        return z_vals
        
    def forward(self, rays, **kwargs):
        rays_o, rays_d, near, far = rays[..., :3], rays[..., 3:6], rays[..., 6:7], rays[..., 7:8]
        t_vals, share_info = self.compute_t_vals(rays, **kwargs)
        z_vals = self.compute_z_vals(t_vals, rays, share_info, **kwargs)
        z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
        points = rays_o[..., None, :] + rays_d[..., None, :] * z_vals_mid[..., :, None]
        if 't2z_func' not in share_info:
            t2z_func = lambda t_vals, near, far: near * (1. - t_vals) + far * t_vals
            share_info.update({'t2z_func': t2z_func})
        share_info.update({'t_vals_{}'.format(self.sample_level-1): t_vals})# , 't2z_func': t2z_func})
        return points, z_vals, share_info
lib/networks/nerf/samplers/uniform_sampler.py
zju3dv/im4d
import torch

def pdf_sample(spacing_vals, weights, training, num_samples, hist_padding=1e-2, eps=1e-5):
    num_bins = num_samples + 1

    weights = weights + hist_padding
    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)

    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 training:
        rand = torch.rand((*cdf.shape[:-1], num_samples + 1), device=cdf.device) / num_bins
        u = u + rand
    u = u.contiguous()

    inds = torch.searchsorted(cdf, u, right=True)
    existing_bins = spacing_vals
    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)
    return bins

lib/networks/nerf/samplers/utils.py
zju3dv/im4d
import torch
import torch.nn as nn
from lib.utils import net_utils
from lib.config import cfg
import math
EPS = 1.0e-7
from lib.train.losses.vgg_perceptual_loss import VGGPerceptualLoss

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

    # t0 = torch.cat([t0_starts, t0_ends[..., -1:]], dim=-1)
    # t1 = torch.cat([t1_starts, t1_ends[..., -1:]], dim=-1)
    # i = torch.arange(t1.shape[-1], device=t1.device)
    # t0_ge_t1 = t0[..., None, :] >= t1[..., :, None]
    # idx_lo = torch.where(t0_ge_t1, i[..., :, None], i[..., :1, None]).max(dim=-2).values
    # idx_hi = torch.where(~t0_ge_t1, i[..., :, None], i[..., -1:, None]).min(dim=-2).values
    # cy1_lo = torch.take_along_dim(cy1, idx_lo, dim=-1)
    # cy1_hi = torch.take_along_dim(cy1, idx_hi, dim=-1)
    # y0_outer = cy1_hi[..., 1:] - cy1_lo[..., :-1]

    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)
    """
    # c: t, w, cp, wp
    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)

class NetworkWrapper(nn.Module):
    def __init__(self, net):
        super(NetworkWrapper, self).__init__()
        self.device = torch.device('cuda:{}'.format(cfg.local_rank))
        self.net = net
        self.color_crit = nn.MSELoss(reduction='mean')
        if 'num_levels' in cfg.network: self.level = cfg.network.num_levels
        else: self.level = len(cfg.samplers)
        if cfg.loss.get('perc_loss', 0.) > 0: self.perc_loss = VGGPerceptualLoss().to(self.device)
    
    def compute_color_loss(self, output, batch, scalar_stats, suffix):
        rgb_map = output['rgb_map{}'.format(suffix)]
        if cfg.loss.get('rand_bkgd', False): 
            assert(not cfg.white_bkgd); 
            if rgb_map.shape[-1] == 3:
                rgb_map = rgb_map + torch.rand_like(rgb_map) * (1. - output['acc_map{}'.format(suffix)][..., None])
            elif rgb_map.shape[-1] == 6:
                rgb_map[..., 3:6] = rgb_map[..., 3:6] + torch.rand_like(rgb_map[..., 3:6]) * (1. - output['acc_map{}'.format(suffix)][..., None])
                
        if rgb_map.shape[-1] == 3:
            color_mse = self.color_crit(rgb_map, batch['rgb'])
            psnr = -10. * torch.log(color_mse.detach()) / math.log(10)
            scalar_stats.update({'psnr{}'.format(suffix): psnr})
            scalar_stats.update({'color_mse{}'.format(suffix): color_mse})
            return color_mse
        elif rgb_map.shape[-1] == 6:
            color_mse = self.color_crit(rgb_map[..., :3], batch['rgb'])
            color_mse_ = color_mse
            psnr = -10. * torch.log(color_mse.detach()) / math.log(10)
            scalar_stats.update({'psnr{}'.format(suffix): psnr})
            scalar_stats.update({'color_mse{}'.format(suffix): color_mse})
            color_mse = self.color_crit(rgb_map[..., 3:], batch['rgb'])
            psnr = -10. * torch.log(color_mse.detach()) / math.log(10)
            scalar_stats.update({'psnr{}_sup'.format(suffix): psnr})
            scalar_stats.update({'color_mse{}_sup'.format(suffix): color_mse})
            return color_mse + color_mse_
    
    def compute_msk_loss(self, output, batch, scalar_stats, suffix):
        msk_mse = self.color_crit(output['acc_map{}'.format(suffix)], batch['msk'])
        scalar_stats.update({'msk_mse{}'.format(suffix): msk_mse})
        return cfg.loss.get('msk_weight') * msk_mse
        
    def compute_prop_loss(self, output, batch, scalar_stats, suffix):
        c = output['t_vals_0']
        w = output['weights_0']
        cp = output['t_vals{}'.format(suffix)]
        wp = output['weights{}'.format(suffix)]
        loss = torch.mean(lossfun_outer(c, w, cp, wp))
        scalar_stats.update({'prop_loss{}'.format(suffix): loss})
        if loss.isnan().any(): loss = 0
        return cfg.loss.get('prop_weight') * loss
    
    def compute_planetv_loss(self, output, batch, scalar_stats, suffix, level):
        model = output['model']
        if level == 0: encoder = model.xyz_encoder
        elif level == 1: encoder = model.sampler.network.xyz_encoder
        elif level == 2: encoder = model.sampler.network.sampler.network.xyz_encoder
        elif level == 3: encoder = model.sampler.network.sampler.network.sampler.network.xyz_encoder
        else: import ipdb; ipdb.set_trace()
        
        loss = 0. 
        spatial_embedding, temporal_embedding = encoder.spatial_embedding, encoder.temporal_embedding
        for i in range(encoder.num_res):
            loss += ( spatial_embedding[i][:, :, 1:, :] -  spatial_embedding[i][:, :, :-1, :]).pow(2).mean()
            loss += ( spatial_embedding[i][:, :, :, 1:] -  spatial_embedding[i][:, :, :, :-1]).pow(2).mean()
            loss += (temporal_embedding[i][:, :, 1:, :] - temporal_embedding[i][:, :, :-1, :]).pow(2).mean()
            loss += (temporal_embedding[i][:, :, :, 1:] - temporal_embedding[i][:, :, :, :-1]).pow(2).mean()
        scalar_stats.update({'plane_tv_loss{}'.format(suffix): loss})
        return cfg.loss.get('plane_tv') * loss
    
    def compute_timesmooth_loss(self, output, batch, scalar_stats, suffix, level):
        model = output['model']
        if level == 0: encoder = model.xyz_encoder
        elif level == 1: encoder = model.sampler.network.xyz_encoder
        elif level == 2: encoder = model.sampler.network.sampler.network.xyz_encoder
        elif level == 3: encoder = model.sampler.network.sampler.network.sampler.network.xyz_encoder
        else: import ipdb; ipdb.set_trace()
        
        loss = 0. 
        temporal_embedding = encoder.temporal_embedding
        for i in range(encoder.num_res):
            first_difference = (temporal_embedding[i][:, :, 1:, :] - temporal_embedding[i][:, :, :-1, :])
            second_difference = first_difference[:, :, 1:, :] - first_difference[:, :, :-1, :]
            loss += second_difference.pow(2).mean()
        scalar_stats.update({'time_smooth{}'.format(suffix): loss})
        return cfg.loss.get('time_smooth')[level] * loss
    
    def compute_distortion_loss(self, output, batch, scalar_stats, suffix):
        # TODO: only support Uniform samples: ENeRF, IBRNet
        weights = output['weights{}'.format(suffix)]
        t_vals = torch.linspace(0, 1., weights.shape[-1], device=weights.device, dtype=weights.dtype)[None, None, :].repeat(weights.shape[0], weights.shape[1], 1)
        t_matrix = t_vals[..., None] - t_vals[..., None, :]
        w_matrix = weights[..., None] * weights[..., None, :]
        loss_matrix = w_matrix * t_matrix.abs()
        loss = loss_matrix.sum(dim=-1).mean()
        scalar_stats.update({'distortion{}'.format(suffix): loss})
        if loss.isnan().any(): loss = 0
        return loss * cfg.loss.get('distortion')
    
    def compute_perc_loss(self, output, batch, scalar_stats, suffix):
        B = len(batch['rgb'])
        pred = output['rgb_map{}'.format(suffix)][:, cfg.num_pixels:].reshape(B*cfg.num_patches, cfg.patch_size, cfg.patch_size, 3).permute(0, 3, 1, 2)
        gt = batch['rgb'][:, cfg.num_pixels:].reshape(B*cfg.num_patches, cfg.patch_size, cfg.patch_size, 3).permute(0, 3, 1, 2) 
        loss = self.perc_loss(pred, gt)
        scalar_stats.update({'perc_loss{}'.format(suffix): loss})
        return loss * cfg.loss.get('perc_loss')
    
    def compute_l1_timeplanes_loss(self, output, batch, scalar_stats, suffix, level):
        model = output['model']
        if level == 0: encoder = model.xyz_encoder
        elif level == 1: encoder = model.sampler.network.xyz_encoder
        elif level == 2: encoder = model.sampler.network.sampler.network.xyz_encoder
        elif level == 3: encoder = model.sampler.network.sampler.network.sampler.network.xyz_encoder
        else: import ipdb; ipdb.set_trace()
        
        loss = 0. 
        temporal_embedding = encoder.temporal_embedding
        for plane in temporal_embedding:
            loss += (1 - plane).pow(2).mean()
        scalar_stats.update({'l1_timeplanes_loss{}'.format(suffix): loss})
        return loss * cfg.loss.get('l1_time_planes')
    
    def compute_loss_level(self, output, batch, scalar_stats, level=-1):
        suffix = '' if level == -1 else '_{}'.format(level)
        loss = 0 
        if 'rgb_map{}'.format(suffix) in output: loss += self.compute_color_loss(output, batch, scalar_stats, suffix) 
        if self.training and 'rgb_map{}'.format(suffix) in output and cfg.get('num_patches', 0) > 0 and cfg.loss.get('perc_loss', 0.) > 0.: loss += self.compute_perc_loss(output, batch, scalar_stats, suffix)
        if 'acc_map{}'.format(suffix) in output and 'msk' in batch and cfg.loss.get('msk_weight', 0.) > 0.: loss += self.compute_msk_loss(output, batch, scalar_stats, suffix) 
        if level >= 1 and cfg.loss.get('prop_weight', 0.) > 0. and 'weights{}'.format(suffix) in output: loss += self.compute_prop_loss(output, batch, scalar_stats, suffix) 
        if cfg.loss.get('plane_tv', 0.) > 0. and batch['step'] % cfg.loss.get('plane_tv_step', 1) == 0: loss += self.compute_planetv_loss(output, batch, scalar_stats, suffix, level=level) 
        if cfg.loss.get('time_smooth', [0., 0., 0., 0., 0.])[level] > 0. and batch['step'] % cfg.loss.get('time_smooth_step', 1) == 0: loss += self.compute_timesmooth_loss(output, batch, scalar_stats, suffix, level=level) 
        if cfg.loss.get('l1_time_planes', 0.) > 0. and batch['step'] % cfg.loss.get('l1_time_planes_step', 1) == 0: loss += self.compute_l1_timeplanes_loss(output, batch, scalar_stats, suffix, level=level) 
        if cfg.loss.get('distortion', 0.) > 0. and 'weights{}'.format(suffix) in output: loss += self.compute_distortion_loss(output, batch, scalar_stats, suffix) 
        return loss

    def forward(self, batch):
        output = self.net(batch)
        loss = 0
        scalar_stats = {} 
        for i in range(self.level):
            loss += self.compute_loss_level(output, batch, scalar_stats, level=i) 
        scalar_stats.update({'loss': loss})
        image_stats = {}
        return output, loss, scalar_stats, image_stats
lib/train/losses/nerf.py
zju3dv/im4d
import torch
import torchvision

class VGGPerceptualLoss(torch.nn.Module):
    def __init__(self, resize=False):
        super(VGGPerceptualLoss, self).__init__()
        blocks = []
        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())
        for bl in blocks:
            for p in bl.parameters():
                p.requires_grad = False
        self.blocks = torch.nn.ModuleList(blocks)
        self.transform = torch.nn.functional.interpolate
        self.resize = resize
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, input, target, feature_layers=[0, 1, 2, 3], style_layers=[]):
        if input.shape[1] != 3:
            input = input.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        input = (input-self.mean) / self.std
        target = (target-self.mean) / self.std
        if self.resize:
            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
        loss = 0.0
        x = input
        y = target
        for i, block in enumerate(self.blocks):
            x = block(x)
            y = block(y)
            if i in feature_layers:
                loss += torch.nn.functional.l1_loss(x, y)
            if i in style_layers:
                act_x = x.reshape(x.shape[0], x.shape[1], -1)
                act_y = y.reshape(y.shape[0], y.shape[1], -1)
                gram_x = act_x @ act_x.permute(0, 2, 1)
                gram_y = act_y @ act_y.permute(0, 2, 1)
                loss += torch.nn.functional.l1_loss(gram_x, gram_y)
        return loss

lib/train/losses/vgg_perceptual_loss.py
zju3dv/im4d
import torch
from lib.utils.optimizer.radam import RAdam
from lib.config import logger


_optimizer_factory = {
    'adam': torch.optim.Adam,
    'radam': RAdam,
    'sgd': torch.optim.SGD
}


def get_lr(lr_table, key):
    splits = key.split('.')
    for split in splits:
        if split in lr_table:
            return lr_table[split]
    return -1.

def make_optimizer(cfg, net):
    params = []
    lr = cfg.train.lr
    weight_decay = cfg.train.weight_decay
    lr_table = cfg.train.lr_table
    eps = cfg.train.eps

    for key, value in net.named_parameters():
        if not value.requires_grad:
            continue
        key_lr = get_lr(lr_table, key)
        if key_lr == 0.:
            logger.info('Skip learning for: ', key)
            continue
        elif key_lr > 0.:
            logger.info('Learning rate for {}: {}'.format(key, key_lr))
            params += [{"params": [value], "lr": key_lr, "weight_decay": weight_decay, "eps": eps}]
            continue
        params += [{"params": [value], "lr": lr, "weight_decay": weight_decay, "eps": eps}]
        # print('Learning rate for {}: {}'.format(key, lr))

    if 'adam' in cfg.train.optim:
        optimizer = _optimizer_factory[cfg.train.optim](params, lr, weight_decay=weight_decay, eps=eps)
    else:
        optimizer = _optimizer_factory[cfg.train.optim](params, lr, momentum=0.9)

    return optimizer

lib/train/optimizer.py
zju3dv/im4d
from collections import deque, defaultdict
import torch
from tensorboardX import SummaryWriter
import os
from lib.config.config import cfg

from termcolor import colored


class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20):
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0

    def update(self, value):
        self.deque.append(value)
        self.count += 1
        self.total += value

    @property
    def median(self):
        d = torch.tensor(list(self.deque))
        return d.median().item()

    @property
    def avg(self):
        d = torch.tensor(list(self.deque))
        return d.mean().item()

    @property
    def global_avg(self):
        return self.total / self.count


def process_volsdf(image_stats):
    for k, v in image_stats.items():
        image_stats[k] = torch.clamp(v[0].permute(2, 0, 1), min=0., max=1.)
    return image_stats

process_neus = process_volsdf

class Recorder(object):
    def __init__(self, cfg):
        if cfg.local_rank > 0:
            return

        log_dir = cfg.record_dir
        if not cfg.resume:
            print(colored('remove contents of recording directory %s' % log_dir, 'red'))
            assert(log_dir != '/')
            os.system('rm -rf {}'.format(os.path.join(log_dir, '*')))
        self.writer = SummaryWriter(log_dir=log_dir)

        # scalars
        self.epoch = 0
        self.step = 0
        self.loss_stats = defaultdict(SmoothedValue)
        self.batch_time = SmoothedValue()
        self.data_time = SmoothedValue()

        # images
        self.image_stats = defaultdict(object)
        # if 'process_' + cfg.task in globals():
            # self.processor = globals()['process_' + cfg.task]
        # else:
            # self.processor = None

    def update_loss_stats(self, loss_dict):
        if cfg.local_rank > 0:
            return
        for k, v in loss_dict.items():
            self.loss_stats[k].update(v.detach().cpu())

    def update_image_stats(self, image_stats):
        if cfg.local_rank > 0:
            return
        # if self.processor is None:
            # return
        # image_stats = self.processor(image_stats)
        if image_stats is None:
            return
        for k, v in image_stats.items():
            self.image_stats[k] = v.detach().cpu()

    def record(self, prefix, step=-1, loss_stats=None, image_stats=None):
        if cfg.local_rank > 0:
            return

        pattern = prefix + '/{}'
        step = step if step >= 0 else self.step
        loss_stats = loss_stats if loss_stats else self.loss_stats

        for k, v in loss_stats.items():
            if isinstance(v, SmoothedValue):
                self.writer.add_scalar(pattern.format(k), v.median, step)
            else:
                self.writer.add_scalar(pattern.format(k), v, step)

        if image_stats is None:
            return
        # if self.processor is None:
            # return
        # image_stats = self.processor(image_stats) if image_stats else self.image_stats
        for k, v in image_stats.items():
            self.writer.add_image(pattern.format(k), v, step)

    def state_dict(self):
        if cfg.local_rank > 0:
            return
        scalar_dict = {}
        scalar_dict['step'] = self.step
        return scalar_dict

    def load_state_dict(self, scalar_dict):
        if cfg.local_rank > 0:
            return
        self.step = scalar_dict['step']

    def __str__(self):
        if cfg.local_rank > 0:
            return
        loss_state = []
        for k, v in self.loss_stats.items():
            loss_state.append('{}: {:.4f}'.format(k, v.avg))
        loss_state = '  '.join(loss_state)

        recording_state = '  '.join(['epoch: {}', 'step: {}', '{}', 'data: {:.4f}', 'batch: {:.4f}'])
        return recording_state.format(self.epoch, self.step, loss_state, self.data_time.avg, self.batch_time.avg)


def make_recorder(cfg):
    return Recorder(cfg)

lib/train/recorder.py
zju3dv/im4d
from collections import Counter
from lib.utils.optimizer.lr_scheduler import WarmupMultiStepLR, MultiStepLR, ExponentialLR
from lib.utils.kplanes.lr_scheduling import get_cosine_schedule_with_warmup
from lib.config import cfg as glo_cfg

def make_lr_scheduler(cfg, optimizer):
    cfg_scheduler = cfg.train.scheduler
    if cfg_scheduler.type == 'multi_step':
        scheduler = MultiStepLR(optimizer,
                                milestones=cfg_scheduler.milestones,
                                gamma=cfg_scheduler.gamma)
    elif cfg_scheduler.type == 'exponential':
        scheduler = ExponentialLR(optimizer,
                                  decay_epochs=cfg_scheduler.decay_epochs,
                                  gamma=cfg_scheduler.gamma)
    elif cfg_scheduler.type == 'warmup_cosine':
        scheduler = get_cosine_schedule_with_warmup(
                optimizer,
                num_warmup_steps=128,
                num_training_steps=glo_cfg.train.epoch * glo_cfg.ep_iter
                )
    return scheduler


def set_lr_scheduler(cfg, scheduler):
    cfg_scheduler = cfg.train.scheduler
    if cfg_scheduler.type == 'multi_step':
        scheduler.milestones = Counter(cfg_scheduler.milestones)
    elif cfg_scheduler.type == 'exponential':
        scheduler.decay_epochs = cfg_scheduler.decay_epochs
    scheduler.gamma = cfg_scheduler.gamma

lib/train/scheduler.py
zju3dv/im4d
from .trainer import Trainer
import imp


def _wrapper_factory(cfg, network, train_loader=None):
    module = cfg.loss_module
    path = cfg.loss_path
    network_wrapper = imp.load_source(module, path).NetworkWrapper(network)
    return network_wrapper


def make_trainer(cfg, network, train_loader=None):
    network = _wrapper_factory(cfg, network)
    return Trainer(network)

lib/train/trainers/make_trainer.py
zju3dv/im4d
import time
import datetime
import torch
import tqdm
from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel
from lib.config import cfg, logger, args
from lib.utils.data_utils import to_cuda


class Trainer(object):
    def __init__(self, network):
        device = torch.device('cuda:{}'.format(cfg.local_rank))
        network = network.to(device)
        if cfg.distributed:
            network = torch.nn.SyncBatchNorm.convert_sync_batchnorm(network)
            network = DistributedDataParallel(
                network,
                device_ids=[cfg.local_rank],
                output_device=cfg.local_rank,
                find_unused_parameters=True
           )
        self.network = network
        self.local_rank = cfg.local_rank
        self.device = device
        self.global_step = 0
        self.gscaler = torch.cuda.amp.GradScaler(enabled=cfg.train_fp16)

    def reduce_loss_stats(self, loss_stats):
        reduced_losses = {k: torch.mean(v) for k, v in loss_stats.items()}
        return reduced_losses

    def to_cuda(self, batch):
        for k in batch:
            if isinstance(batch[k], tuple) or isinstance(batch[k], list):
                #batch[k] = [b.cuda() for b in batch[k]]
                batch[k] = [b.to(self.device) for b in batch[k]]
            elif isinstance(batch[k], dict):
                batch[k] = {key: self.to_cuda(batch[k][key]) for key in batch[k]}
            else:
                # batch[k] = batch[k].cuda()
                batch[k] = batch[k].to(self.device)
        return batch

    def train(self, epoch, data_loader, optimizer, recorder, scheduler):
        max_iter = len(data_loader)
        self.network.train()
        end = time.time()
        if self.global_step == 0:
            self.global_step = cfg.ep_iter * epoch
        iterator = enumerate(data_loader) if cfg.log_level in ['INFO', 'DEBUG'] else tqdm.tqdm(enumerate(data_loader), desc='Epoch {:03d}/{:03d}: '.format(epoch, cfg.train.epoch), total=len(data_loader))
        
        for iteration, batch in iterator:
            data_time = time.time() - end
            iteration = iteration + 1

            batch = to_cuda(batch, self.device)
            batch['step'] = self.global_step
            with torch.cuda.amp.autocast(enabled=cfg.train_fp16):
                output, loss, loss_stats, image_stats = self.network(batch)

                loss = loss.mean()

            # training stage: loss; optimizer; scheduler
            optimizer.zero_grad(set_to_none=True)
            self.gscaler.scale(loss).backward()
            # loss.backward()
            # torch.nn.utils.clip_grad_value_(self.network.parameters(), 10)
            self.gscaler.step(optimizer)
            scale = self.gscaler.get_scale()

            self.gscaler.update()
            scheduler.step()

            # if scale > self.gscaler.get_scale(): logger.warning('ATTETION')

            self.global_step += 1
            if cfg.local_rank > 0:
                continue

            # data recording stage: loss_stats, time, image_stats
            recorder.step += 1

            loss_stats = self.reduce_loss_stats(loss_stats)
            recorder.update_loss_stats(loss_stats)

            batch_time = time.time() - end
            end = time.time()
            recorder.batch_time.update(batch_time)
            recorder.data_time.update(data_time)

            if iteration % cfg.log_interval == 0 or iteration == (max_iter - 1):
                # print training state
                eta_seconds = recorder.batch_time.global_avg * (max_iter - iteration)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                lr = optimizer.param_groups[0]['lr']
                memory = torch.cuda.max_memory_allocated() / 1024.0 / 1024.0

                training_state = '  '.join(['eta: {}', '{}', 'lr: {:.6f}', 'max_mem: {:.0f}'])
                training_state = training_state.format(eta_string, str(recorder), lr, memory)
                logger.info(training_state)

                # record loss_stats and image_dict
                recorder.update_image_stats(image_stats)
                recorder.record('train')

    def val(self, epoch, data_loader, evaluator=None, recorder=None):
        self.network.eval()
        torch.cuda.empty_cache()
        val_loss_stats = {}
        image_stats = {}
        data_size = len(data_loader)
        for batch in tqdm.tqdm(data_loader):
            batch = to_cuda(batch, self.device)
            batch['step'] = recorder.step
            with torch.no_grad():
                with torch.cuda.amp.autocast(enabled=cfg.eval_fp16):
                    output, loss, loss_stats, _ = self.network(batch)
                if evaluator is not None:
                    image_stats_ = evaluator.evaluate(output, batch)
                    if image_stats_ is not None:
                        image_stats.update(image_stats_)

            loss_stats = self.reduce_loss_stats(loss_stats)
            for k, v in loss_stats.items():
                val_loss_stats.setdefault(k, 0)
                val_loss_stats[k] += v

        loss_state = []
        for k in val_loss_stats.keys():
            val_loss_stats[k] /= data_size
            loss_state.append('{}: {:.4f}'.format(k, val_loss_stats[k]))
        logger.info(loss_state)

        if evaluator is not None:
            result = evaluator.summarize()
            val_loss_stats.update(result)

        if recorder:
            recorder.record('val', epoch, val_loss_stats, image_stats)

lib/train/trainers/trainer.py
zju3dv/im4d
import pickle
import os
import numpy as np
import cv2
import time
from termcolor import colored
import importlib
import torch.distributed as dist
import math
import torch

class perf_timer:
    def __init__(self, msg="Elapsed time: {}s", logf=lambda x: print(colored(x, 'yellow')), sync_cuda=True, use_ms=False, disabled=False):
        self.logf = logf
        self.msg = msg
        self.sync_cuda = sync_cuda
        self.use_ms = use_ms
        self.disabled = disabled

        self.loggedtime = None

    def __enter__(self,):
        if self.sync_cuda:
            torch.cuda.synchronize()
        self.start = time.perf_counter()

    def __exit__(self, exc_type, exc_value, traceback):
        if self.sync_cuda:
            torch.cuda.synchronize()
        self.logtime(self.msg)

    def logtime(self, msg=None, logf=None):
        if self.disabled:
            return
        # SAME CLASS, DIFFERENT FUNCTIONALITY, is this good?
        # call the logger for timing code sections
        if self.sync_cuda:
            torch.cuda.synchronize()

        # always remember current time
        prev = self.loggedtime
        self.loggedtime = time.perf_counter()

        # print it if we've remembered previous time
        if prev is not None and msg:
            logf = logf or self.logf
            diff = self.loggedtime-prev
            diff *= 1000 if self.use_ms else 1
            logf(msg.format(diff))

        return self.loggedtime

def read_pickle(pkl_path):
    with open(pkl_path, 'rb') as f:
        return pickle.load(f)


def save_pickle(data, pkl_path):
    os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
    with open(pkl_path, 'wb') as f:
        pickle.dump(data, f)


def project(xyz, K, RT):
    """
    xyz: [N, 3]
    K: [3, 3]
    RT: [3, 4]
    """
    xyz = np.dot(xyz, RT[:, :3].T) + RT[:, 3:].T
    xyz = np.dot(xyz, K.T)
    xy = xyz[:, :2] / xyz[:, 2:]
    return xy

def get_bbox_2d(bbox, K, RT):
    pts = np.array([[bbox[0, 0], bbox[0, 1], bbox[0, 2]],
                    [bbox[0, 0], bbox[0, 1], bbox[1, 2]],
                    [bbox[0, 0], bbox[1, 1], bbox[0, 2]],
                    [bbox[0, 0], bbox[1, 1], bbox[1, 2]],
                    [bbox[1, 0], bbox[0, 1], bbox[0, 2]],
                    [bbox[1, 0], bbox[0, 1], bbox[1, 2]],
                    [bbox[1, 0], bbox[1, 1], bbox[0, 2]],
                    [bbox[1, 0], bbox[1, 1], bbox[1, 2]],
                    ])
    pts_2d = project(pts, K, RT)
    return [pts_2d[:, 0].min(), pts_2d[:, 1].min(), pts_2d[:, 0].max(), pts_2d[:, 1].max()]


def get_bound_corners(bounds):
    min_x, min_y, min_z = bounds[0]
    max_x, max_y, max_z = bounds[1]
    corners_3d = np.array([
        [min_x, min_y, min_z],
        [min_x, min_y, max_z],
        [min_x, max_y, min_z],
        [min_x, max_y, max_z],
        [max_x, min_y, min_z],
        [max_x, min_y, max_z],
        [max_x, max_y, min_z],
        [max_x, max_y, max_z],
    ])
    return corners_3d

def get_bound_2d_mask(bounds, K, pose, H, W):
    corners_3d = get_bound_corners(bounds)
    corners_2d = project(corners_3d, K, pose)
    corners_2d = np.round(corners_2d).astype(int)
    mask = np.zeros((H, W), dtype=np.uint8)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 3, 2, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[4, 5, 7, 6, 5]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 5, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[2, 3, 7, 6, 2]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 2, 6, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[1, 3, 7, 5, 1]]], 1)
    return mask

def get_bound_2d_mask_padding(bounds, K, pose, H, W, padding=0.1):
    bounds = bounds.copy()
    bounds[:1] -= padding
    bounds[1:] += padding
    corners_3d = get_bound_corners(bounds)
    corners_2d = project(corners_3d, K, pose)
    corners_2d = np.round(corners_2d).astype(int)
    mask = np.zeros((H, W), dtype=np.uint8)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 3, 2, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[4, 5, 7, 6, 5]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 5, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[2, 3, 7, 6, 2]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 2, 6, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[1, 3, 7, 5, 1]]], 1)
    return mask



def load_object(module_name, module_args, **extra_args):
    module_path = '.'.join(module_name.split('.')[:-1])
    module = importlib.import_module(module_path)
    name = module_name.split('.')[-1]
    obj = getattr(module, name)(**extra_args, **module_args)
    return obj



def get_indices(length):
    num_replicas = dist.get_world_size()
    rank = dist.get_rank()
    num_samples = int(math.ceil(length * 1.0 / num_replicas))
    total_size = num_samples * num_replicas
    indices = np.arange(length).tolist()
    indices += indices[: (total_size - len(indices))]
    offset = num_samples * rank
    indices = indices[offset:offset+num_samples]
    return indices


class DotDict(dict):
    """
    a dictionary that supports dot notation
    as well as dictionary access notation
    usage: d = DotDict() or d = DotDict({'val1':'first'})
    set attributes: d.val2 = 'second' or d['val2'] = 'second'
    get attributes: d.val2 or d['val2']
    """
    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

    def __init__(self, dct=None):
        if dct is not None:
            for key, value in dct.items():
                if hasattr(value, 'keys'):
                    value = DotDict(value)
                # self[key] = value

def get_sample(input):
    if len(input) == 4 and input[-1] == -1:
        return np.arange(input[0], input[1])[::input[2]].tolist()
    elif -1 in input and len(input==2):
        return input[:1]
    elif -1 not in input:
        return input
    else:
        import ipdb; ipdb.set_trace()
        

def get_xywh_tensor(mask):
    xywh = []
    for b in range(len(mask)):
        nonzero = mask[b].nonzero()
        x1, y1, x2, y2 = nonzero[:, 1].min(), nonzero[:, 0].min(), nonzero[:, 1].max(), nonzero[:, 0].max()
        xywh.append([x1, y1, x2-x1+1, y2-y1+1])
    return torch.as_tensor(xywh)
lib/utils/base_utils.py
zju3dv/im4d
import cv2
import numpy as np
import os
from os.path import join
class FileStorage(object):
    def __init__(self, filename, isWrite=False):
        version = cv2.__version__
        self.major_version = int(version.split('.')[0])
        self.second_version = int(version.split('.')[1])

        if isWrite:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            self.fs = open(filename, 'w')
            self.fs.write('%YAML:1.0\r\n')
            self.fs.write('---\r\n')
        else:
            assert os.path.exists(filename), filename
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_READ)
        self.isWrite = isWrite

    def __del__(self):
        if self.isWrite:
            self.fs.close()
        else:
            cv2.FileStorage.release(self.fs)

    def _write(self, out):
        self.fs.write(out+'\r\n')

    def write(self, key, value, dt='mat'):
        if dt == 'mat':
            self._write('{}: !!opencv-matrix'.format(key))
            self._write('  rows: {}'.format(value.shape[0]))
            self._write('  cols: {}'.format(value.shape[1]))
            self._write('  dt: d')
            self._write('  data: [{}]'.format(', '.join(['{:.6f}'.format(i) for i in value.reshape(-1)])))
        elif dt == 'list':
            self._write('{}:'.format(key))
            for elem in value:
                self._write('  - "{}"'.format(elem))
        elif dt == 'int':
            self._write('{}: {}'.format(key, value))

    def read(self, key, dt='mat'):
        if dt == 'mat':
            output = self.fs.getNode(key).mat()
        elif dt == 'list':
            results = []
            n = self.fs.getNode(key)
            for i in range(n.size()):
                val = n.at(i).string()
                if val == '':
                    val = str(int(n.at(i).real()))
                if val != 'none':
                    results.append(val)
            output = results
        elif dt == 'int':
            output = int(self.fs.getNode(key).real())
        else:
            raise NotImplementedError
        return output

    def close(self):
        self.__del__(self)

def read_intri(intri_name):
    assert os.path.exists(intri_name), intri_name
    intri = FileStorage(intri_name)
    camnames = intri.read('names', dt='list')
    cameras = {}
    for key in camnames:
        cam = {}
        cam['K'] = intri.read('K_{}'.format(key))
        cam['invK'] = np.linalg.inv(cam['K'])
        cam['dist'] = intri.read('dist_{}'.format(key))
        cameras[key] = cam
    return cameras

def write_intri(intri_name, cameras):
    if not os.path.exists(os.path.dirname(intri_name)):
        os.makedirs(os.path.dirname(intri_name))
    intri = FileStorage(intri_name, True)
    results = {}
    camnames = list(cameras.keys())
    intri.write('names', camnames, 'list')
    for key_, val in cameras.items():
        key = key_.split('.')[0]
        K, dist = val['K'], val['dist']
        assert K.shape == (3, 3), K.shape
        assert dist.shape == (1, 5) or dist.shape == (5, 1) or dist.shape == (1, 4) or dist.shape == (4, 1), dist.shape
        intri.write('K_{}'.format(key), K)
        intri.write('dist_{}'.format(key), dist.flatten()[None])

def write_extri(extri_name, cameras):
    if not os.path.exists(os.path.dirname(extri_name)):
        os.makedirs(os.path.dirname(extri_name))
    extri = FileStorage(extri_name, True)
    results = {}
    camnames = list(cameras.keys())
    extri.write('names', camnames, 'list')
    for key_, val in cameras.items():
        key = key_.split('.')[0]
        extri.write('R_{}'.format(key), val['Rvec'])
        extri.write('Rot_{}'.format(key), val['R'])
        extri.write('T_{}'.format(key), val['T'])
    return 0

def read_camera(intri_name, extri_name, cam_names=[]):
    assert os.path.exists(intri_name), intri_name
    assert os.path.exists(extri_name), extri_name

    intri = FileStorage(intri_name)
    extri = FileStorage(extri_name)
    cams, P = {}, {}
    cam_names = intri.read('names', dt='list')
    for cam in cam_names:
        # 内参只读子码流的
        cams[cam] = {}
        cams[cam]['K'] = intri.read('K_{}'.format( cam))
        cams[cam]['invK'] = np.linalg.inv(cams[cam]['K'])
        H = intri.read('H_{}'.format(cam), dt='int')
        W = intri.read('W_{}'.format(cam), dt='int')
        if H is None or W is None:
            print('[camera] no H or W for {}'.format(cam))
            H, W = -1, -1
        cams[cam]['H'] = H
        cams[cam]['W'] = W
        Rvec = extri.read('R_{}'.format(cam))
        Tvec = extri.read('T_{}'.format(cam))
        assert Rvec is not None, cam
        R = cv2.Rodrigues(Rvec)[0]
        RT = np.hstack((R, Tvec))

        cams[cam]['RT'] = RT
        cams[cam]['R'] = R
        cams[cam]['Rvec'] = Rvec
        cams[cam]['T'] = Tvec
        cams[cam]['center'] = - Rvec.T @ Tvec
        P[cam] = cams[cam]['K'] @ cams[cam]['RT']
        cams[cam]['P'] = P[cam]

        cams[cam]['dist'] = intri.read('dist_{}'.format(cam))
        if cams[cam]['dist'] is None:
            cams[cam]['dist'] = intri.read('D_{}'.format(cam))
            if cams[cam]['dist'] is None:
                print('[camera] no dist for {}'.format(cam))
    cams['basenames'] = cam_names
    return cams

def read_cameras(path, intri='intri.yml', extri='extri.yml', subs=[]):
    cameras = read_camera(join(path, intri), join(path, extri))
    cameras.pop('basenames')
    if len(subs) > 0:
        cameras = {key:cameras[key].astype(np.float32) for key in subs}
    return cameras

def write_camera(camera, path, suffix=None):
    from os.path import join
    intri_name = join(path, 'intri.yml') if suffix is None else join(path, 'intri_{}.yml'.format(suffix))
    extri_name = join(path, 'extri.yml') if suffix is None else join(path, 'extri_{}.yml'.format(suffix))
    intri = FileStorage(intri_name, True)
    extri = FileStorage(extri_name, True)
    results = {}
    camnames = [key_.split('.')[0] for key_ in camera.keys()]
    intri.write('names', camnames, 'list')
    extri.write('names', camnames, 'list')
    for key_, val in camera.items():
        if key_ == 'basenames':
            continue
        key = key_.split('.')[0]
        intri.write('K_{}'.format(key), val['K'])
        intri.write('dist_{}'.format(key), val['dist'])
        if 'H' in val.keys() and 'W' in val.keys():
            intri.write('H_{}'.format(key), val['H'], dt='int')
            intri.write('W_{}'.format(key), val['W'], dt='int')
        if 'Rvec' not in val.keys():
            val['Rvec'] = cv2.Rodrigues(val['R'])[0]
        extri.write('R_{}'.format(key), val['Rvec'])
        extri.write('Rot_{}'.format(key), val['R'])
        extri.write('T_{}'.format(key), val['T'])

def camera_from_img(img):
    height, width = img.shape[0], img.shape[1]
    # focal = 1.2*max(height, width) # as colmap
    focal = 1.2*min(height, width) # as colmap
    K = np.array([focal, 0., width/2, 0., focal, height/2, 0. ,0., 1.]).reshape(3, 3)
    camera = {'K':K ,'R': np.eye(3), 'T': np.zeros((3, 1)), 'dist': np.zeros((1, 5))}
    camera['invK'] = np.linalg.inv(camera['K'])
    camera['P'] = camera['K'] @ np.hstack((camera['R'], camera['T']))
    return camera

class Undistort:
    distortMap = {}
    @classmethod
    def image(cls, frame, K, dist, sub=None, interp=cv2.INTER_NEAREST):
        if sub is None:
            return cv2.undistort(frame, K, dist, None)
        else:
            if sub not in cls.distortMap.keys():
                h,  w = frame.shape[:2]
                mapx, mapy = cv2.initUndistortRectifyMap(K, dist, None, K, (w,h), 5)
                cls.distortMap[sub] = (mapx, mapy)
            mapx, mapy = cls.distortMap[sub]
            img = cv2.remap(frame, mapx, mapy, interp)
            return img

    @staticmethod
    def points(keypoints, K, dist):
        # keypoints: (N, 3)
        assert len(keypoints.shape) == 2, keypoints.shape
        kpts = keypoints[:, None, :2]
        kpts = np.ascontiguousarray(kpts)
        kpts = cv2.undistortPoints(kpts, K, dist, P=K)
        keypoints = np.hstack([kpts[:, 0], keypoints[:, 2:]])
        return keypoints
    
    @staticmethod
    def bbox(bbox, K, dist):
        keypoints = np.array([[bbox[0], bbox[1], 1], [bbox[2], bbox[3], 1]])
        kpts = Undistort.points(keypoints, K, dist)
        bbox = np.array([kpts[0, 0], kpts[0, 1], kpts[1, 0], kpts[1, 1], bbox[4]])
        return bbox

class Distort:
    @staticmethod
    def points(keypoints, K, dist):
        pass

    @staticmethod
    def bbox(bbox, K, dist):
        keypoints = np.array([[bbox[0], bbox[1]], [bbox[2], bbox[3]]], dtype=np.float32)
        k3d = cv2.convertPointsToHomogeneous(keypoints)
        k3d = (np.linalg.inv(K) @ k3d[:, 0].T).T[:, None]
        k2d, _ = cv2.projectPoints(k3d, np.zeros((3,)), np.zeros((3,)), K, dist)
        k2d = k2d[:, 0]
        bbox = np.array([k2d[0,0], k2d[0,1], k2d[1, 0], k2d[1, 1], bbox[-1]])
        return bbox

def unproj(kpts, invK):
    homo = np.hstack([kpts[:, :2], np.ones_like(kpts[:, :1])])
    homo = homo @ invK.T
    return np.hstack([homo[:, :2], kpts[:, 2:]])
class UndistortFisheye:
    @staticmethod
    def image(frame, K, dist):
        Knew = K.copy()
        frame = cv2.fisheye.undistortImage(frame, K, dist, Knew=Knew)
        return frame, Knew

    @staticmethod
    def points(keypoints, K, dist, Knew):
        # keypoints: (N, 3)
        assert len(keypoints.shape) == 2, keypoints.shape
        kpts = keypoints[:, None, :2]
        kpts = np.ascontiguousarray(kpts)
        kpts = cv2.fisheye.undistortPoints(kpts, K, dist, P=Knew)
        keypoints = np.hstack([kpts[:, 0], keypoints[:, 2:]])
        return keypoints
    
    @staticmethod
    def bbox(bbox, K, dist, Knew):
        keypoints = np.array([[bbox[0], bbox[1], 1], [bbox[2], bbox[3], 1]])
        kpts = UndistortFisheye.points(keypoints, K, dist, Knew)
        bbox = np.array([kpts[0, 0], kpts[0, 1], kpts[1, 0], kpts[1, 1], bbox[4]])
        return bbox


def get_Pall(cameras, camnames):
    Pall = np.stack([cameras[cam]['K'] @ np.hstack((cameras[cam]['R'], cameras[cam]['T'])) for cam in camnames])
    return Pall

def get_fundamental_matrix(cameras, basenames):
    skew_op = lambda x: np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
    fundamental_op = lambda K_0, R_0, T_0, K_1, R_1, T_1: np.linalg.inv(K_0).T @ (
            R_0 @ R_1.T) @ K_1.T @ skew_op(K_1 @ R_1 @ R_0.T @ (T_0 - R_0 @ R_1.T @ T_1))
    fundamental_RT_op = lambda K_0, RT_0, K_1, RT_1: fundamental_op (K_0, RT_0[:, :3], RT_0[:, 3], K_1,
                                                                          RT_1[:, :3], RT_1[:, 3] )
    F = np.zeros((len(basenames), len(basenames), 3, 3))  # N x N x 3 x 3 matrix
    F = {(icam, jcam): np.zeros((3, 3)) for jcam in basenames for icam in basenames}
    for icam in basenames:
        for jcam in basenames:
            F[(icam, jcam)] += fundamental_RT_op(cameras[icam]['K'], cameras[icam]['RT'], cameras[jcam]['K'], cameras[jcam]['RT'])
            if F[(icam, jcam)].sum() == 0:
                F[(icam, jcam)] += 1e-12  # to avoid nan
    return F

def interp_cameras(cameras, keys, step=20, loop=True, allstep=-1, **kwargs):
    from scipy.spatial.transform import Rotation as R
    from scipy.spatial.transform import Slerp
    if allstep != -1:
        tall = np.linspace(0., 1., allstep+1)[:-1].reshape(-1, 1, 1)
    elif allstep == -1 and loop:
        tall = np.linspace(0., 1., 1+step*len(keys))[:-1].reshape(-1, 1, 1)
    elif allstep == -1 and not loop:
        tall = np.linspace(0., 1., 1+step*(len(keys)-1))[:-1].reshape(-1, 1, 1)
    cameras_new = {}
    for ik in range(len(keys)):
        if ik == len(keys) -1 and not loop:
            break
        if loop:
            start, end = (ik * tall.shape[0])//len(keys),     int((ik+1)*tall.shape[0])//len(keys)
            print(ik, start, end, tall.shape)
        else:
            start, end = (ik * tall.shape[0])//(len(keys)-1), int((ik+1)*tall.shape[0])//(len(keys)-1)
        t = tall[start:end].copy()
        t = (t-t.min())/(t.max()-t.min())
        left, right = keys[ik], keys[0 if ik == len(keys)-1 else ik + 1]
        camera_left = cameras[left]
        camera_right = cameras[right]
        # 插值相机中心: center = - R.T @ T
        center_l = - camera_left['R'].T @ camera_left['T']
        center_r = - camera_right['R'].T @ camera_right['T']
        center_l, center_r = center_l[None], center_r[None]
        if False:
            centers = center_l * (1-t) + center_r * t
        else:
            # 球面插值
            norm_l, norm_r = np.linalg.norm(center_l), np.linalg.norm(center_r)
            center_l, center_r = center_l/norm_l, center_r/norm_r
            costheta = (center_l*center_r).sum()
            sintheta = np.sqrt(1. - costheta**2)
            theta = np.arctan2(sintheta, costheta)
            centers = (np.sin(theta*(1-t)) * center_l + np.sin(theta * t) * center_r)/sintheta
            norm = norm_l * (1-t) + norm_r * t
            centers = centers * norm
        key_rots = R.from_matrix(np.stack([camera_left['R'], camera_right['R']]))
        key_times = [0, 1]
        slerp = Slerp(key_times, key_rots)
        interp_rots = slerp(t.squeeze()).as_matrix()
        # 计算相机T RX + T = 0 => T = - R @ X
        T = - np.einsum('bmn,bno->bmo', interp_rots, centers)
        K = camera_left['K'] * (1-t) + camera_right['K'] * t
        for i in range(T.shape[0]):
            cameras_new['{}-{}-{}'.format(left, right, i)] = \
                {
                    'K': K[i],
                    'dist': np.zeros((1, 5)),
                    'R': interp_rots[i],
                    'T': T[i]
                }
    return cameras_new
lib/utils/camera_utils.py
zju3dv/im4d
import numpy as np
import cv2
import random
import json
from torch import nn
import torch
from os.path import join
import trimesh
# from imgaug import augmenters as iaa
# from lib.config import cfg
from plyfile import PlyData
# from lib.utils import data_config
import re
import os
from io import BytesIO
from typing import Union, List, Tuple
import imageio
import torch.nn.functional as F
from lib.utils.parallel_utils import async_call
# from lib.utils.nerfstudio import colormaps

def read_bbox_from_dir(dir_path, frames, format='{:05d}.ply', padding_size=0.1, frame_start_number=1):
    if os.path.exists(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1]))):
        print('Load bbox from {}'.format(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1]))))
        return np.loadtxt(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1]))).astype(np.float32)
    import trimesh 
    bounds = []
    for frame_id in range(frames[0], frames[1]):
        if format[-3:] == 'npy':
            points = np.asarray(np.load(join(dir_path, format.format(frame_id + frame_start_number))))
        else:
            points = np.asarray(trimesh.load(join(dir_path, format.format(frame_id + frame_start_number))).vertices)
        # transform = np.array(cfg.transform)
        # points = points @ transform[:3, :3].T + transform[:3, 3:].T
        bound = np.array([points.min(0)-padding_size, points.max(0)+padding_size]).astype(np.float32)
        bounds.append(bound)
    bound = np.array([np.array(bounds).min(axis=(0, 1)), np.array(bounds).max(axis=(0, 1))]).astype(np.float32)
    try:
        print('Save bbox to {}'.format(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1]))))
        np.savetxt(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1])), bound)
        return np.loadtxt(join(dir_path, 'bbox_{}_{}_{}.txt'.format(padding_size, frames[0], frames[1]))).astype(np.float32)
    except:
        return bound
        
def export_dir_to_video(img_dir, video_path, fps=30, format='.png', quality=8.5, crf=8):
    import shutil
    if shutil.which('ffmpeg') is not None:
        cmd = 'ffmpeg -loglevel error -y -framerate {} -f image2 -pattern_type glob -nostdin -y -r {} -i "{}/*{}" -c:v libx264 -crf {} -pix_fmt yuv420p {}'.format(fps, fps, img_dir, format, crf, video_path)
        # cmd = 'ffmpeg -loglevel error -y -framerate {} -f image2 -pattern_type glob -nostdin -y -r {} -i "{}/*{}" -c:v libx264 -crf {} -pix_fmt yuv444p {}'.format(fps, fps, img_dir, format, crf, video_path)
        os.system(cmd)
    else:
        img_paths = [item.path for item in os.scandir(img_dir) if item.name.endswith(format)]
        img_paths = sorted(img_paths)
        imgs = []
        for img_path in img_paths:
            img = imageio.imread(img_path)
            if (img.shape) == 3:
                img = img[..., :3]
            else:
                img = img
            imgs.append(img)
        imageio.mimwrite(video_path, imgs, fps=fps, quality=quality)
            

def rgb_convert_to_img(img, acc=None):
    if img.dtype == np.float32 or img.dtype == np.float64:
        img = (img * 255.).astype(np.uint8)
    if acc is not None:
        acc = (acc * 255.).astype(np.uint8)
        img = np.concatenate([img, acc[..., None]], axis=-1)
    return img

def dpt_convert_to_img(dpt, acc=None):
    min_val, max_val = dpt.min(), dpt.max()
    if acc is not None and (acc > 0.9).sum() > 0:
        min_val = dpt[acc>0.9].min()
    dpt = (dpt - min_val) / (max_val - min_val)
    dpt = np.clip(dpt, 0., 1.)
    dpt = dpt * 0.6 + 0.2
    
    dpt = cv2.applyColorMap((dpt*255.).astype(np.uint8), cv2.COLORMAP_JET)
    if acc is not None:
        dpt = np.concatenate([dpt, (acc * 255.).astype(np.uint8)[..., None]], axis=-1)
    return dpt


def load_img(img_path):
    return np.asarray(imageio.imread(img_path)).astype(np.uint8)

    
@async_call
def save_img(img_path, img, acc_map=None):
    if acc_map is not None:
        img = np.concatenate([img.astype(np.float32), acc_map.astype(np.float32)[..., None]], axis=-1)
    if img.dtype == np.float32 or img.dtype == np.float64:
        img = (img * 255.).astype(np.uint8)
    os.system('mkdir -p {}'.format(os.path.dirname(img_path)))
    imageio.imwrite(img_path, img)
    
@async_call
def save_srcinps(img_path, src_inps):
    src_inp = np.concatenate([ src_inp.transpose(1, 2, 0) * 0.5 + 0.5 for src_inp in src_inps], axis=1)
    save_img(img_path, src_inp)
    
    
def gen_possible_depth(dpt, start=0, end=1.):
    dpt = dpt.copy()
    dpt = dpt * (end - start) + start
    return dpt

def gen_possible_depths(depth, acc_map):
    lengths = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.]
    for start in np.linspace(0, 0.9, 10):
        for length in lengths:
            if start + length <= 1.:
                dpt = gen_possible_depth(depth, start, start + length)
                img = cv2.applyColorMap((dpt*255.).astype(np.uint8), cv2.COLORMAP_JET)
                print(start, start + length, dpt.min(), dpt.max())
                cv2.imwrite('depths/{}_{}.png'.format(start, start + length), np.concatenate([img, (acc_map[..., None] * 255.).astype(np.uint8)], axis=-1))
    
@async_call
def save_depth(img_path, img, acc_map=None, min_val=None, save_orig=False):
    if save_orig:
        np.savez_compressed(img_path.replace('.jpg', '.npz'), img=img, acc_map=acc_map)
    img = img.copy()
    if min_val is not None:
        img = np.clip(img, min_val, None)
    elif acc_map is not None:
        try:
            min_val = img[acc_map > 0.9].min()
        except:
            print('Warning: acc_map is all zero, use min_val=0.')
            min_val = 0.
        img = np.clip(img, min_val, None)
        # dpt = colormaps.apply_depth_colormap(torch.from_numpy(img)[..., None], torch.from_numpy(acc_map)[..., None]).numpy()
        # imageio.imwrite(img_path, (dpt * 255.).astype(np.uint8))
    img -= img.min()
    img /= (img.max() + 1e-6)
    img = img * 0.6 + 0.2 # 0.2 == 0.8
    # import ipdb; ipdb.set_trace()
    # gen_possible_depths(img, acc_map)
    # img = 1 - img
    # img = (img * 255.).astype(np.uint8)
    # save_img(img_path, img)
    # img0 = cv2.applyColorMap((img*255.).astype(np.uint8), cv2.COLORMAP_PLASMA)
    img = cv2.applyColorMap((img*255.).astype(np.uint8), cv2.COLORMAP_JET)
    # img0_path = img_path.replace('.jpg', '_plasma.jpg')
    # img_path = img_path
    if acc_map is not None:
        # img0 = np.concatenate([img0, (acc_map[..., None] * 255.).astype(np.uint8)], axis=-1)
        img = np.concatenate([img, (acc_map[..., None] * 255.).astype(np.uint8)], axis=-1)
        # img0_path = img0_path.replace('.jpg', '.png')
        img_path = img_path.replace('.jpg', '.png')
    # cv2.imwrite(img0_path, img0)
    os.system('mkdir -p {}'.format(os.path.dirname(img_path)))
    cv2.imwrite(img_path, img)
    # save_img(img_path, img)
    
def calib_color_img(img_ori, bgr_sol, transpose=True):
    rs = []
    for i in range(3):
        channel = np.array(img_ori[:,:,i],dtype=np.double)
        X = np.stack([channel**2,channel ,np.ones_like(channel)])
        if transpose:
            y = np.dot(bgr_sol[2-i].reshape(1,3),X.reshape(3,-1)).reshape(channel.shape)
        else:
            y = np.dot(bgr_sol[i].reshape(1,3),X.reshape(3,-1)).reshape(channel.shape)
        rs.append(y)
    rs_img = np.stack(rs,axis=2)
    rs_i = cv2.normalize(rs_img, None, alpha=rs_img.min(), beta=rs_img.max(), norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return rs_i

def load_resize_undist_calib_im_bytes(imp: str, K: np.ndarray, D: np.ndarray, center_h_w: Union[int, List[int]], scale: Union[float, List[int]], ccm: np.ndarray,
                                encode_ext='.jpg', decode_flag=cv2.IMREAD_UNCHANGED, dist_opt_K: bool = False):
    # Load image -> resize -> undistort -> save to bytes (jpeg)
    img = load_image_from_bytes(load_image_bytes(imp), decode_flag=decode_flag)  # cv2 decoding (fast)
    img = calib_color_img(img, ccm)
    if np.linalg.norm(D) != 0.:
        img = cv2.undistort(img, K, D)
    oH, oW = img.shape[:2]
    if isinstance(scale, float): H, W = int(oH * scale), int(oW * scale)
    else: H, W = scale  # ratio is actually the target image size
    # rH, rW = H / oH, W / oW
    if oH != H or oW != W:
        if imp.split('.')[-1] == 'jpg':
            img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)  # H, W, 3, uint8
        else:
            img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)  # H, W, 3, uint8
        
    if isinstance(center_h_w, int):
        tar_h, tar_w = center_h_w, center_h_w
    else:
        tar_h, tar_w = center_h_w
    if tar_h != -1:
        crop_h, crop_w = (H - tar_h)//2, (W - tar_w)//2
        crop_h_, crop_w_ = H - tar_h - crop_h, W - tar_w - crop_w
        img = img[crop_h:-crop_h_, crop_w:-crop_w_]
    _, buffer = cv2.imencode(encode_ext if imp.split('.')[-1] == 'jpg' else '.png', img)
    buffer: np.ndarray
    buffer = BytesIO(buffer)  # is this slow? tobytes and BytesIO
    return buffer


def load_resize_undist_schp_im_bytes(imp: str, K: np.ndarray, D: np.ndarray, center_h_w: Union[int, List[int]], scale: Union[float, List[int]],
                                encode_ext='.jpg', decode_flag=cv2.IMREAD_UNCHANGED, dist_opt_K: bool = False):
    # Load image -> resize -> undistort -> save to bytes (jpeg)
    img = load_image_from_bytes(load_image_bytes(imp), decode_flag=decode_flag)  # cv2 decoding (fast)
    # for schp, we need to firstly convert it into binary mask, and then undistort
    msk = img.sum(axis=-1) != 0
    img = (msk * 255).astype(np.uint8)
    if np.linalg.norm(D) != 0.:
        img = cv2.undistort(img, K, D)
    oH, oW = img.shape[:2]
    if isinstance(scale, float): H, W = int(oH * scale), int(oW * scale)
    else: H, W = scale  # ratio is actually the target image size
    # rH, rW = H / oH, W / oW
    img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)  # H, W, 3, uint8
    
        
    if isinstance(center_h_w, int):
        tar_h, tar_w = center_h_w, center_h_w
    else:
        tar_h, tar_w = center_h_w
    if tar_h != -1:
        crop_h, crop_w = (H - tar_h)//2, (W - tar_w)//2
        crop_h_, crop_w_ = H - tar_h - crop_h, W - tar_w - crop_w
        img = img[crop_h:-crop_h_, crop_w:-crop_w_]
    _, buffer = cv2.imencode(encode_ext, img)
    buffer: np.ndarray
    buffer = BytesIO(buffer)  # is this slow? tobytes and BytesIO
    return buffer

def load_resize_undist_im_bytes(imp: str, K: np.ndarray, D: np.ndarray, center_h_w: Union[int, List[int]], scale: Union[float, List[int]],
                                encode_ext='.jpg', decode_flag=cv2.IMREAD_UNCHANGED, dist_opt_K: bool = False):
    # Load image -> resize -> undistort -> save to bytes (jpeg)
    img = load_image_from_bytes(load_image_bytes(imp), decode_flag=decode_flag)  # cv2 decoding (fast)
    if np.linalg.norm(D) != 0.:
        img = cv2.undistort(img, K, D)
    oH, oW = img.shape[:2]
    if isinstance(scale, float): H, W = int(oH * scale), int(oW * scale)
    else: H, W = scale  # ratio is actually the target image size
    # rH, rW = H / oH, W / oW
    if imp.split('.')[-1] == 'jpg':
        img = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)  # H, W, 3, uint8
    else:
        img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)  # H, W, 3, uint8
    
        
    if isinstance(center_h_w, int):
        tar_h, tar_w = center_h_w, center_h_w
    else:
        tar_h, tar_w = center_h_w
    if tar_h != -1:
        crop_h, crop_w = (H - tar_h)//2, (W - tar_w)//2
        img = img[crop_h:crop_h+tar_h, crop_w:crop_w+tar_w]
    if encode_ext is not None:
        _, buffer = cv2.imencode(encode_ext if imp.split('.')[-1] == 'jpg' else '.png', img)
        buffer: np.ndarray
        buffer = BytesIO(buffer)  # is this slow? tobytes and BytesIO
        return buffer
    else:
        return img

def load_image_from_bytes(buffer: np.ndarray, decode_flag=cv2.IMREAD_UNCHANGED, debug=False):
    if isinstance(buffer, BytesIO):
        buffer = buffer.getvalue()  # slow? copy?
    if isinstance(buffer, bytes):
        image: np.ndarray = cv2.imdecode(np.frombuffer(buffer, np.uint8), decode_flag)  # MARK: 10-15ms
    else:
        image = buffer
    return image

def load_image_bytes(im: str):
    with open(im, "rb") as fh:
        buffer = BytesIO(fh.read())
    return buffer

def parse_cameras_genebody(cams):
    basenames, c2ws, ixts, Ds = [], [], [], []
    for k in cams:
        basenames.append(k)
        c2ws.append(cams[k]['c2w'])
        ixts.append(cams[k]['K'])
        Ds.append(cams[k]['D'])
    exts = np.linalg.inv(np.array(c2ws)).astype(np.float32)
    return basenames, \
           np.array(exts).astype(np.float32), \
           np.array(c2ws).astype(np.float32), \
           np.array(ixts).astype(np.float32), \
           np.array(Ds).astype(np.float32)

def add_batch(batch):
    if isinstance(batch, tuple) or isinstance(batch, list):
        batch = [add_batch(b) for b in batch]
    elif isinstance(batch, dict):
        batch_ = {}
        for key in batch:
            batch_[key] = add_batch(batch[key])
        batch = batch_
    elif isinstance(batch, torch.Tensor) or isinstance(batch, np.ndarray):
        batch = batch[None]
    else:
        batch = torch.tensor(batch)[None]
    return batch

def get_bound_2d_mask(bounds, K, H, W):
    points = bounds[:, :3] @ K.T
    corners_2d = points[..., :2] / points[..., 2:]
    corners_2d = np.round(corners_2d).astype(int)
    mask = np.zeros((H, W), dtype=np.uint8)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 3, 2, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[4, 5, 7, 6, 5]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 5, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[2, 3, 7, 6, 2]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 2, 6, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[1, 3, 7, 5, 1]]], 1)
    return mask

def get_bound_2d_mask_padding(bounds, K, H, W, padding=0.1):
    bounds = bounds.coppy()
    points = bounds[:, :3] @ K.T
    corners_2d = points[..., :2] / points[..., 2:]
    corners_2d = np.round(corners_2d).astype(int)
    mask = np.zeros((H, W), dtype=np.uint8)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 3, 2, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[4, 5, 7, 6, 5]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 1, 5, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[2, 3, 7, 6, 2]]], 1)
    cv2.fillPoly(mask, [corners_2d[[0, 2, 6, 4, 0]]], 1)
    cv2.fillPoly(mask, [corners_2d[[1, 3, 7, 5, 1]]], 1)
    return mask

def carve_vhull(
    H: torch.Tensor, W: torch.Tensor, K: torch.Tensor, R: torch.Tensor, T: torch.Tensor,
    bounds: torch.Tensor, msk: torch.Tensor, voxel_size: float = 0.10, padding: float = 0.10,
    vhull_thresh: float = 0.75, count_thresh: int = 1,
):
    # Project these points onto the camera
    xyz1 = torch.cat([xyz, torch.ones_like(xyz[..., :1])], dim=-1)[None]  # homo
    cam = xyz1 @ affine_padding(torch.cat([R, T], dim=-1)).mT  # N, P, 4 @ N, 3, 4 -> N, P, 4, homo camera coordinates
    cam = cam[..., :-1]  # last dim will always be one for valid affine transforms
    pix = cam @ K.mT  # N, P, 3 @ N, 3, 3, homo pixel coords
    pix = pix[..., :-1] / pix[..., -1:]  # N, P, 2, pixel coords

    # Use max size, instead of the one passed in before sampling
    pixel_uv_range = pix / torch.stack([W, H], dim=-1)[..., None, :] * 2 - 1  # N, P, 2 to sample the msk
    should_count_camera = ((pixel_uv_range > -1.0) & (pixel_uv_range < 1.0)).all(dim=-1)  # N, P
    vhull_camera_count = should_count_camera.sum(dim=0)  # P,

    H, W = msk.shape[-3:-1]  # sampling size (the largest of all images)
    # pix = pixel_uv_range
    pix = pix / msk.new_tensor([W, H]) * 2 - 1  # N, P, 2 to sample the msk (dimensionality normalization for sampling)
    valid = F.grid_sample(msks.permute(0, 3, 1, 2), pix[:, None], align_corners=True)[:, 0, 0]  # whether this is background
    valid = (valid > 0.5).float().sum(dim=0)  # N, 1, 1, P -> N, P
    valid = (valid / vhull_camera_count > vhull_thresh) & (vhull_camera_count > count_thresh)  # P, ratio of cameras sees this

    # Find valid points on the voxels
    import ipdb; ipdb.set_trace()
    inds = valid.nonzero()  # MARK: SYNC
    vhull = world_xyz.reshape(-1, 3)[inds]
    vhull = remove_outlier(vhull[None], K=5, std_ratio=5.0)[0]
    import ipdb; ipdb.set_trace()

    bounds = torch.stack([vhull.min(dim=0)[0], vhull.max(dim=0)[0]])
    bounds = bounds + bounds.new_tensor([-padding, padding])[:, None]

    return vhull, bounds

def get_padding_masks(msks):
    hs, ws = [msk.shape[1] for msk in msks], [msk.shape[2] for msk in msks]
    hmax, wmax = np.max(hs), np.max(ws)
    msk_len = len(msks)
    new_msks = torch.zeros((msk_len, hmax, wmax), dtype=torch.float32, device=msks[0].device)
    for i in range(msk_len):
        new_msks[i, :hs[i], :ws[i]] = msks[i][0]
    return new_msks

def multi_indexing(indices: torch.Tensor, shape: torch.Size, dim=-2):
    # index will first be augmented to match the values' dimentionality at the back
    # then we will try to broatcast index's shape to values shape
    shape = list(shape)
    back_pad = len(shape) - indices.ndim
    for _ in range(back_pad): indices = indices.unsqueeze(-1)
    expand_shape = shape
    expand_shape[dim] = -1
    return indices.expand(*expand_shape)

def multi_gather(values: torch.Tensor, indices: torch.Tensor, dim=-2):
    # Gather the value at the -2th dim of values, augment index shape on the back
    # Example: values: B, P, 3, index: B, N, -> B, N, 3
    # index will first be augmented to match the values' dimentionality at the back
    # take care of batch dimension of, and acts like a linear indexing in the target dimention
    # we assume that the values's second to last dimension is the dimension to be indexed on
    return values.gather(dim, multi_indexing(indices, values.shape, dim))

def remove_outlier(pts: torch.Tensor, K: int = 20, std_ratio=2.0, ret_inds=False):
    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pts.view(-1, 3).detach().cpu().numpy())
    cl, ind = pcd.remove_statistical_outlier(nb_neighbors=K, std_ratio=std_ratio)
    if not ret_inds:
        return torch.as_tensor(np.array(pcd.points)[np.array(ind)]).to(pts.device, pts.dtype, non_blocking=True).view(pts.shape[0], -1, 3)
    else:
        return torch.as_tensor(np.array(pcd.points)[np.array(ind)]).to(pts.device, pts.dtype, non_blocking=True).view(pts.shape[0], -1, 3), ind
        

def get_grid_from_mask(batch, vhull_thresh=0.9, count_thresh=6):
    msks, exts, ixts = batch['msks'], batch['exts'], batch['ixts']
    hs, ws = torch.as_tensor([msk.shape[1] for msk in msks], device=exts.device), torch.as_tensor([msk.shape[2] for msk in msks], device=exts.device)
    msks = get_padding_masks(msks)[..., None]
    pts = batch['pts']
    world_xyz = pts[0, ..., 0, :] 
    xyz = world_xyz.reshape(-1, 3)
    xyz = torch.cat([xyz, torch.ones_like(xyz[..., :1])], dim=-1)
    xyz = xyz[None] @ exts[0].mT 
    pix = xyz[:, :, :3] @ ixts[0].mT
    pix = pix[..., :-1] / pix[..., -1:]  # N, P, 2, pixel coords

    # Use max size, instead of the one passed in before sampling
    pixel_uv_range = pix / torch.stack([ws, hs], dim=-1)[..., None, :] * 2 - 1  # N, P, 2 to sample the msk
    should_count_camera = ((pixel_uv_range > -1.0) & (pixel_uv_range < 1.0)).all(dim=-1)  # N, P
    vhull_camera_count = should_count_camera.sum(dim=0)  # P,
    
    H, W = msks.shape[-3:-1]
    pix = pix / msks.new_tensor([W, H]) * 2 - 1  # N, P, 2 to sample the msk (dimensionality normalization for sampling)
    valid = F.grid_sample(msks.permute(0, 3, 1, 2), pix[:, None], align_corners=True)[:, 0, 0]  # whether this is background
    valid = (valid > 0.5).float().sum(dim=0)  # N, 1, 1, P -> N, P
    valid = (valid / vhull_camera_count > vhull_thresh) & (vhull_camera_count > count_thresh)  # P, ratio of cameras sees this
    
    inds = valid.nonzero()  # MARK: SYNC
    vhull = multi_gather(world_xyz.reshape(-1, 3), inds)  # P, 3; V -> V, 3
    vhull, valid_inds = remove_outlier(vhull[None], K=5, std_ratio=5.0, ret_inds=True)
    inds = inds[torch.tensor(valid_inds, device=inds.device, dtype=torch.long)]
    
    binary = torch.zeros_like(world_xyz[..., 0])
    binary.view(-1, 1)[inds] = 1.
    # import ipdb; ipdb.set_trace()
    if cfg.padding_grid:
        structuring_element = torch.ones(1, 1, 3, 3, 3, device=binary.device)
        dilated_tensor = F.conv3d(binary[None, None].float(), structuring_element.float(), padding=1) > 0 
        binary = dilated_tensor[0, 0].bool()
    return {'binary': binary.detach().cpu().numpy(), 'bounds': batch['bounds'][0].detach().cpu().numpy()}
    
    

def save_ply(points, ply_path, num_clusters=10):
    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points.reshape(-1, 3))
    from sklearn.cluster import KMeans
    kmeans = KMeans(n_clusters=num_clusters).fit(np.asarray(pcd.points))
    labels = kmeans.labels_
    colors = np.random.rand(num_clusters, 3)
    point_colors = colors[labels]
    pcd.colors = o3d.utility.Vector3dVector(point_colors)

    if os.path.dirname(ply_path) != '':
        os.system('mkdir -p ' + os.path.dirname(ply_path))
    o3d.io.write_point_cloud(ply_path, pcd)

def read_and_denoise(ply_path, vox_size=0.02):
    import open3d as o3d
    pcd = o3d.io.read_point_cloud(ply_path)
    # import ipdb; ipdb.set_trace()
    # o3d.visualization.draw_geometries([pcd])
    voxel_down_pcd = pcd.voxel_down_sample(voxel_size=vox_size)
    # o3d.visualization.draw_geometries([voxel_down_pcd])

    def display_inlier_outlier(cloud, ind):
        inlier_cloud = cloud.select_by_index(ind)
        outlier_cloud = cloud.select_by_index(ind, invert=True)
        outlier_cloud.paint_uniform_color([1, 0, 0])
        inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
        o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])

    cl, ind = voxel_down_pcd.remove_statistical_outlier(nb_neighbors=20,
                                                    std_ratio=2.0)
    # display_inlier_outlier(voxel_down_pcd, ind)
    return voxel_down_pcd.select_by_index(ind)




def read_cam_file(filename):
    with open(filename) as f:
        lines = [line.rstrip() for line in f.readlines()]
    # extrinsics: line [1,5), 4x4 matrix
    extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ')
    extrinsics = extrinsics.reshape((4, 4))
    # intrinsics: line [7-10), 3x3 matrix
    intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ')
    intrinsics = intrinsics.reshape((3, 3))
    # depth_min & depth_interval: line 11
    depth_min = float(lines[11].split()[0])
    return intrinsics, extrinsics, depth_min

def read_pmn_cam_file(filename):
    with open(filename) as f:
        lines = [line.rstrip() for line in f.readlines()]
    # extrinsics: line [1,5), 4x4 matrix
    extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ')
    extrinsics = extrinsics.reshape((4, 4))
    # intrinsics: line [7-10), 3x3 matrix
    intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ')
    intrinsics = intrinsics.reshape((3, 3))
    # depth_min & depth_interval: line 11
    depth_min = float(lines[11].split()[0])
    depth_max = float(lines[11].split()[1])
    return intrinsics, extrinsics, depth_min, depth_max

def read_pfm(filename):
    file = open(filename, 'rb')
    color = None
    width = None
    height = None
    scale = None
    endian = None

    header = file.readline().decode('utf-8').rstrip()
    if header == 'PF':
        color = True
    elif header == 'Pf':
        color = False
    else:
        raise Exception('Not a PFM file.')

    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('utf-8'))
    if dim_match:
        width, height = map(int, dim_match.groups())
    else:
        raise Exception('Malformed PFM header.')

    scale = float(file.readline().rstrip())
    if scale < 0:  # little-endian
        endian = '<'
        scale = -scale
    else:
        endian = '>'  # big-endian

    data = np.fromfile(file, endian + 'f')
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    file.close()
    return data, scale

def to_tensor(rgb):
    rgb = rgb.astype(np.float32) / 255.
    rgb = (rgb-data_config.mean_rgb) / data_config.std_rgb
    return rgb.transpose(2, 0, 1)

def to_img(tensor):
    tensor = tensor.detach().cpu().clone().numpy()
    tensor = tensor.transpose(1, 2, 0)
    tensor = tensor * data_config.std_rgb + data_config.mean_rgb
    tensor *= 255.
    return tensor.astype(np.uint8)


def resize_images(imgs, masks, ixt, input_size):
    ori_h, ori_w = imgs[0].shape[0], imgs[0].shape[1]
    tar_h, tar_w = input_size
    imgs = [cv2.resize(img, input_size, interpolation=cv2.INTER_LINEAR) for img in imgs]
    masks = [cv2.resize(mask.astype(np.uint8), input_size, interpolation=cv2.INTER_NEAREST) for mask in masks]
    ixt[0][0] *= (tar_h/ori_h)
    ixt[0][2] *= (tar_h/ori_h)
    ixt[1][1] *= (tar_w/ori_w)
    ixt[1][2] *= (tar_w/ori_w)
    return imgs, masks, ixt

def resize_image(img, mask, ixt, input_size):
    ori_h, ori_w = img.shape[0], img.shape[1]
    tar_h, tar_w = input_size
    img = cv2.resize(img, input_size, interpolation=cv2.INTER_LINEAR)
    mask = cv2.resize(mask.astype(np.uint8), input_size, interpolation=cv2.INTER_NEAREST)
    ixt[0][0] *= (tar_h/ori_h)
    ixt[0][2] *= (tar_h/ori_h)
    ixt[1][1] *= (tar_w/ori_w)
    ixt[1][2] *= (tar_w/ori_w)
    return img, mask, ixt

def load_matrix(path):
    lines = [[float(w) for w in line.strip().split()] for line in open(path)]
    if len(lines[0]) == 2:
        lines = lines[1:]
    if len(lines[-1]) == 2:
        lines = lines[:-1]
    return np.array(lines).astype(np.float32)

def load_nsvf_intrinsics(filepath, resized_width=None, invert_y=False):
    try:
        intrinsics = load_matrix(filepath)
        if intrinsics.shape[0] == 3 and intrinsics.shape[1] == 3:
            _intrinsics = np.zeros((4, 4), np.float32)
            _intrinsics[:3, :3] = intrinsics
            _intrinsics[3, 3] = 1
            intrinsics = _intrinsics
        if intrinsics.shape[0] == 1 and intrinsics.shape[1] == 16:
            intrinsics = intrinsics.reshape(4, 4)
        return intrinsics
    except ValueError:
        pass

    # Get camera intrinsics
    with open(filepath, 'r') as file:

        f, cx, cy, _ = map(float, file.readline().split())
    fx = f
    if invert_y:
        fy = -f
    else:
        fy = f

    # Build the intrinsic matrices
    full_intrinsic = np.array([[fx, 0., cx, 0.],
                               [0., fy, cy, 0],
                               [0., 0, 1, 0],
                               [0, 0, 0, 1]])
    return full_intrinsic

def gaussian_radius(det_size, min_overlap=0.7):
    height, width = det_size

    a1 = 1
    b1 = (height + width)
    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
    r1 = (b1 + sq1) / 2

    a2 = 4
    b2 = 2 * (height + width)
    c2 = (1 - min_overlap) * width * height
    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
    r2 = (b2 + sq2) / 2

    a3 = 4 * min_overlap
    b3 = -2 * min_overlap * (height + width)
    c3 = (min_overlap - 1) * width * height
    if b3 ** 2 - 4 * a3 * c3 < 0:
        r3 = min(r1, r2)
    else:
        sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
        r3 = (b3 + sq3) / 2
    return min(r1, r2, r3)


def gaussian2D(shape, sigma=(1, 1), rho=0):
    if not isinstance(sigma, tuple):
        sigma = (sigma, sigma)
    sigma_x, sigma_y = sigma

    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m+1, -n:n+1]

    energy = (x * x) / (sigma_x * sigma_x) - 2 * rho * x * y / (sigma_x * sigma_y) + (y * y) / (sigma_y * sigma_y)
    h = np.exp(-energy / (2 * (1 - rho * rho)))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h


def draw_umich_gaussian(heatmap, center, radius, k=1):
    diameter = 2 * radius + 1
    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)

    x, y = int(center[0]), int(center[1])

    height, width = heatmap.shape[0:2]

    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
    return heatmap


def draw_distribution(heatmap, center, sigma_x, sigma_y, rho, radius, k=1):
    diameter = 2 * radius + 1
    gaussian = gaussian2D((diameter, diameter), (sigma_x/3, sigma_y/3), rho)

    x, y = int(center[0]), int(center[1])

    height, width = heatmap.shape[0:2]

    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
    return heatmap


def draw_heatmap_np(hm, point, box_size):
    """point: [x, y]"""
    # radius = gaussian_radius(box_size)
    radius = box_size[0]
    radius = max(0, int(radius))
    ct_int = np.array(point, dtype=np.int32)
    draw_umich_gaussian(hm, ct_int, radius)
    return hm


def get_edge(mask):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    return mask - cv2.erode(mask, kernel)


def compute_gaussian_1d(dmap, sigma=1):
    """dmap: each entry means a distance"""
    prob = np.exp(-dmap / (2 * sigma * sigma))
    prob[prob < np.finfo(prob.dtype).eps * prob.max()] = 0
    return prob


def get_3rd_point(a, b):
    direct = a - b
    return b + np.array([-direct[1], direct[0]], dtype=np.float32)


def get_dir(src_point, rot_rad):
    sn, cs = np.sin(rot_rad), np.cos(rot_rad)

    src_result = [0, 0]
    src_result[0] = src_point[0] * cs - src_point[1] * sn
    src_result[1] = src_point[0] * sn + src_point[1] * cs

    return src_result


def get_affine_transform(center,
                         scale,
                         rot,
                         output_size,
                         shift=np.array([0, 0], dtype=np.float32),
                         inv=0):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        scale = np.array([scale, scale], dtype=np.float32)

    scale_tmp = scale
    src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]

    rot_rad = np.pi * rot / 180
    src_dir = get_dir([0, src_w * -0.5], rot_rad)
    dst_dir = np.array([0, dst_w * -0.5], np.float32)

    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale_tmp * shift
    src[1, :] = center + src_dir + scale_tmp * shift
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir

    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])

    if inv:
        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

    return trans


def affine_transform(pt, t):
    """pt: [n, 2]"""
    new_pt = np.dot(np.array(pt), t[:, :2].T) + t[:, 2]
    return new_pt


def homography_transform(pt, H):
    """pt: [n, 2]"""
    pt = np.concatenate([pt, np.ones([len(pt), 1])], axis=1)
    pt = np.dot(pt, H.T)
    pt = pt[..., :2] / pt[..., 2:]
    return pt


def get_border(border, size):
    i = 1
    while np.any(size - border // i <= border // i):
        i *= 2
    return border // i


def grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)


def lighting_(data_rng, image, alphastd, eigval, eigvec):
    alpha = data_rng.normal(scale=alphastd, size=(3, ))
    image += np.dot(eigvec, eigval * alpha)


def blend_(alpha, image1, image2):
    image1 *= alpha
    image2 *= (1 - alpha)
    image1 += image2


def saturation_(data_rng, image, gs, gs_mean, var):
    alpha = 1. + data_rng.uniform(low=-var, high=var)
    blend_(alpha, image, gs[:, :, None])


def brightness_(data_rng, image, gs, gs_mean, var):
    alpha = 1. + data_rng.uniform(low=-var, high=var)
    image *= alpha


def contrast_(data_rng, image, gs, gs_mean, var):
    alpha = 1. + data_rng.uniform(low=-var, high=var)
    blend_(alpha, image, gs_mean)


def color_aug(data_rng, image, eig_val, eig_vec):
    functions = [brightness_, contrast_, saturation_]
    random.shuffle(functions)

    gs = grayscale(image)
    gs_mean = gs.mean()
    for f in functions:
        f(data_rng, image, gs, gs_mean, 0.4)
    lighting_(data_rng, image, 0.1, eig_val, eig_vec)


def blur_aug(inp):
    if np.random.random() < 0.1:
        if np.random.random() < 0.8:
            inp = iaa.blur_gaussian_(inp, abs(np.clip(np.random.normal(0, 1.5), -3, 3)))
        else:
            inp = iaa.MotionBlur((3, 15), (-45, 45))(images=[inp])[0]


def gaussian_blur(image, sigma):
    from scipy import ndimage
    if image.ndim == 2:
        image[:, :] = ndimage.gaussian_filter(image[:, :], sigma, mode="mirror")
    else:
        nb_channels = image.shape[2]
        for channel in range(nb_channels):
            image[:, :, channel] = ndimage.gaussian_filter(image[:, :, channel], sigma, mode="mirror")


def inter_from_mask(pred, gt):
    pred = pred.astype(np.bool)
    gt = gt.astype(np.bool)
    intersection = np.logical_and(gt, pred).sum()
    return intersection


def draw_poly(mask, poly):
    cv2.fillPoly(mask, [poly], 255)
    return mask


def inter_from_poly(poly, gt, width, height):
    mask_small = np.zeros((1, height, width), dtype=np.uint8)
    mask_small = draw_poly(mask_small, poly)
    mask_gt = gt[..., 0]

    return inter_from_mask(mask_small, mask_gt)


def inter_from_polys(poly, w, h, gt_mask):
    inter = inter_from_poly(poly, gt_mask, w, h)
    if inter > 0:
        return False
    return True


def select_point(shape, poly, gt_mask):
    for i in range(cfg.max_iter):
        y = np.random.randint(shape[0] - poly['bbox'][3])
        x = np.random.randint(shape[1] - poly['bbox'][2])
        delta = np.array([poly['bbox'][0] - x, poly['bbox'][1] - y])
        poly_move = np.array(poly['poly']) - delta
        inter = inter_from_polys(poly_move, shape[1], shape[0], gt_mask)
        if inter:
            return x, y
    x, y = -1, -1
    return x, y


def transform_small_gt(poly, box, x, y):
    delta = np.array([poly['bbox'][0] - x, poly['bbox'][1] - y])
    poly['poly'] -= delta
    box[:2] -= delta
    box[2:] -= delta
    return poly, box


def get_mask_img(img, poly):
    mask = np.zeros(img.shape[:2])[..., np.newaxis]
    cv2.fillPoly(mask, [np.round(poly['poly']).astype(int)], 1)
    poly_img = img * mask
    mask = mask[..., 0]
    return poly_img, mask


def add_small_obj(img, gt_mask, poly, box, polys_gt):
    poly_img, mask = get_mask_img(img, poly)
    x, y = select_point(img.shape, poly.copy(), gt_mask)
    if x == -1:
        box = []
        return img, poly, box
    poly, box = transform_small_gt(poly, box, x, y)
    _, mask_ori = get_mask_img(img, poly)
    gt_mask += mask_ori[..., np.newaxis]
    img[mask_ori == 1] = poly_img[mask == 1]
    return img, poly, box[np.newaxis, :], gt_mask


def get_gt_mask(img, poly):
    mask = np.zeros(img.shape[:2])[..., np.newaxis]
    for i in range(len(poly)):
        for j in range(len(poly[i])):
            cv2.fillPoly(mask, [np.round(poly[i][j]['poly']).astype(int)], 1)
    return mask


def small_aug(img, poly, box, label, num):
    N = len(poly)
    gt_mask = get_gt_mask(img, poly)
    for i in range(N):
        if len(poly[i]) > 1:
            continue
        if poly[i][0]['area'] < 32*32:
            for k in range(num):
                img, poly_s, box_s, gt_mask = add_small_obj(img, gt_mask, poly[i][0].copy(), box[i].copy(), poly)
                if len(box_s) == 0:
                    continue
                poly.append([poly_s])
                box = np.concatenate((box, box_s))
                label.append(label[i])
    return img, poly, box, label


def truncated_normal(mean, sigma, low, high, data_rng=None):
    if data_rng is None:
        data_rng = np.random.RandomState()
    value = data_rng.normal(mean, sigma)
    return np.clip(value, low, high)


def _nms(heat, kernel=3):
    """heat: [b, c, h, w]"""
    pad = (kernel - 1) // 2

    # find the local minimum of heat within the neighborhood kernel x kernel
    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=pad)
    keep = (hmax == heat).float()
    return heat * keep


def _gather_feat(feat, ind, mask=None):
    dim = feat.size(2)
    ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
    feat = feat.gather(1, ind)
    if mask is not None:
        mask = mask.unsqueeze(2).expand_as(feat)
        feat = feat[mask]
        feat = feat.view(-1, dim)
    return feat


def _topk(scores, K=40):
    batch, cat, height, width = scores.size()

    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

    topk_inds = topk_inds % (height * width)
    topk_ys = (topk_inds / width).int().float()
    topk_xs = (topk_inds % width).int().float()

    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
    topk_clses = (topk_ind / K).int()
    topk_inds = _gather_feat(
        topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


def clip_to_image(bbox, h, w):
    bbox[..., :2] = torch.clamp(bbox[..., :2], min=0)
    bbox[..., 2] = torch.clamp(bbox[..., 2], max=w-1)
    bbox[..., 3] = torch.clamp(bbox[..., 3], max=h-1)
    return bbox


def load_ply(path):
    ply = PlyData.read(path)
    data = ply.elements[0].data
    x, y, z = data['x'], data['y'], data['z']
    model = np.stack([x, y, z], axis=-1)
    return model

def to_cuda(batch, device=torch.device('cuda:0')):
    if isinstance(batch, tuple) or isinstance(batch, list):
        #batch[k] = [b.cuda() for b in batch[k]]
        #batch[k] = [b.to(self.device) for b in batch[k]]
        batch = [to_cuda(b, device) for b in batch]
    elif isinstance(batch, dict):
        #batch[k] = {key: self.to_cuda(batch[k][key]) for key in batch[k]}
        batch_ = {}
        for key in batch:
            if key == 'meta':
                batch_[key] = batch[key]
            else:
                batch_[key] = to_cuda(batch[key], device)
        batch = batch_
    else:
        # batch[k] = batch[k].cuda()
        batch = batch.to(device)
    return batch

def load_K_Rt_from_P(filename, P=None):
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K / K[2, 2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3, 3] = (t[:3] / t[3])[:, 0]

    return intrinsics, pose


def get_ray_near_far_torch(bounds, ray_o, ray_d):
    """calculate intersections with 3d bounding box in pytorch"""
    norm_d = torch.norm(ray_d, dim=-1, keepdim=True)
    viewdir = ray_d / norm_d
    viewdir[(viewdir < 1e-5) & (viewdir > -1e-10)] = 1e-5
    viewdir[(viewdir > -1e-5) & (viewdir < 1e-10)] = -1e-5
    tmin = (bounds[:1] - ray_o[:1]) / viewdir
    tmax = (bounds[1:2] - ray_o[:1]) / viewdir
    t1 = torch.min(tmin, tmax)
    t2 = torch.max(tmin, tmax)
    near = torch.max(t1, dim=-1).values
    far = torch.min(t2, dim=-1).values
    mask_at_box = near < far
    near = near[mask_at_box] / norm_d[mask_at_box, 0]
    far = far[mask_at_box] / norm_d[mask_at_box, 0]
    return near, far, mask_at_box

def get_ray_near_far(bounds, ray_o, ray_d):
    """calculate intersections with 3d bounding box"""
    norm_d = np.linalg.norm(ray_d, axis=-1, keepdims=True)
    viewdir = ray_d / norm_d
    viewdir[(viewdir < 1e-5) & (viewdir > -1e-10)] = 1e-5
    viewdir[(viewdir > -1e-5) & (viewdir < 1e-10)] = -1e-5
    tmin = (bounds[:1] - ray_o[:1]) / viewdir
    tmax = (bounds[1:2] - ray_o[:1]) / viewdir
    t1 = np.minimum(tmin, tmax)
    t2 = np.maximum(tmin, tmax)
    near = np.max(t1, axis=-1)
    far = np.min(t2, axis=-1)
    mask_at_box = near < far
    near = near[mask_at_box] / norm_d[mask_at_box, 0]
    far = far[mask_at_box] / norm_d[mask_at_box, 0]
    return near, far, mask_at_box

def read_camera(intri_name, extri_name, cam_names=[]):
    assert os.path.exists(intri_name), intri_name
    assert os.path.exists(extri_name), extri_name

    intri = FileStorage(intri_name)
    extri = FileStorage(extri_name)
    cams, P = {}, {}
    cam_names = intri.read('names', dt='list')
    for cam in cam_names:
        # 内参只读子码流的
        cams[cam] = {}
        cams[cam]['K'] = intri.read('K_{}'.format( cam))
        cams[cam]['invK'] = np.linalg.inv(cams[cam]['K'])
        Rvec = extri.read('R_{}'.format(cam))
        Tvec = extri.read('T_{}'.format(cam))
        R = cv2.Rodrigues(Rvec)[0]
        RT = np.hstack((R, Tvec))

        cams[cam]['RT'] = RT
        cams[cam]['R'] = R
        cams[cam]['T'] = Tvec
        P[cam] = cams[cam]['K'] @ cams[cam]['RT']
        cams[cam]['P'] = P[cam]

        cams[cam]['dist'] = intri.read('dist_{}'.format(cam))
    cams['basenames'] = cam_names
    return cams

class FileStorage(object):
    def __init__(self, filename, isWrite=False):
        version = cv2.__version__
        self.major_version = int(version.split('.')[0])
        self.second_version = int(version.split('.')[1])

        if isWrite:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_WRITE)
        else:
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_READ)

    def __del__(self):
        cv2.FileStorage.release(self.fs)

    def write(self, key, value, dt='mat'):
        if dt == 'mat':
            cv2.FileStorage.write(self.fs, key, value)
        elif dt == 'list':
            if self.major_version == 4: # 4.4
                self.fs.startWriteStruct(key, cv2.FileNode_SEQ)
                for elem in value:
                    self.fs.write('', elem)
                self.fs.endWriteStruct()
            else: # 3.4
                self.fs.write(key, '[')
                for elem in value:
                    self.fs.write('none', elem)
                self.fs.write('none', ']')

    def read(self, key, dt='mat'):
        if dt == 'mat':
            output = self.fs.getNode(key).mat()
        elif dt == 'list':
            results = []
            n = self.fs.getNode(key)
            for i in range(n.size()):
                val = n.at(i).string()
                if val == '':
                    val = str(int(n.at(i).real()))
                if val != 'none':
                    results.append(val)
            output = results
        else:
            raise NotImplementedError
        return output

    def close(self):
        self.__del__(self)


def get_mask(schp, mask):
    msk = schp.astype(np.int32)
    msk = (msk * [-1, 10, 100]).sum(axis=-1)
    palette = np.array([0, 128, 1280, 1408, 12800, 12928, 14080, 14208, 64,
                        192, 1344, 1472, 12864, 12992, 14144, 14272, 640, 768,
                        1920, 2048])
    leg_msk = (msk == palette[9]) | (msk == palette[16]) | (msk == palette[18])
    msk = (mask > 100).astype(np.uint8)
    msk[leg_msk] = 1
    return msk


def save_mesh_with_extracted_fields(ply_path, u, bounds, grid_resolution, threshold=0.5):
    import trimesh
    import mcubes
    vertices, triangles = mcubes.marching_cubes(u, threshold)
    vertices = vertices / (np.array(grid_resolution)) * (bounds[1] - bounds[0])[None, :] + bounds[0][None, :]
    mesh = trimesh.Trimesh(vertices, triangles)
    os.system('mkdir -p {}'.format(os.path.dirname(ply_path)))
    mesh.export(ply_path)
    print('Save mesh to: {}'.format(ply_path))
    
    
def export_path_inp_cameras(input_cam, render_cam, path='test.ply'):
    print('Save render cams and input cams to: {}'.format(path))
    save_ply(np.concatenate([input_cam, render_cam], axis=0), path)
    
def get_space_points_torch(bounds, grid_resolution):
    x = torch.linspace(bounds[0, 0], bounds[1, 0], grid_resolution[0], device=bounds.device)
    y = torch.linspace(bounds[0, 1], bounds[1, 1], grid_resolution[1], device=bounds.device)
    z = torch.linspace(bounds[0, 2], bounds[1, 2], grid_resolution[2], device=bounds.device)
    X, Y, Z = torch.meshgrid(x, y, z)
    space_points = torch.stack([X, Y, Z], dim=-1)
    return space_points


def get_ratios(length, start_ratio, end_ratio, inter='linear'):
    funcs = {
        'linear': lambda x:x,
        'smoothstep': lambda x: x * x * (3 - 2 * x),
        'smootherstep': lambda x: x * x * x * (x * (x * 6 - 15) + 10),
        'smoothererstep': lambda x: x * x * x * x * (x * (x * (x * (-20) + 70) - 84) + 35)
    }
    weights = funcs[inter](np.linspace(0, 1, length))
    ratios = start_ratio + (end_ratio - start_ratio) * weights
    return ratios


def opencv2tranform(c2w):
    c2w = c2w.copy()
    c2w[0:3, 1:3] *= -1
    c2w = c2w[np.array([1, 0, 2, 3]), :]
    c2w[2, :] *= -1
    return c2w

def transform2opencv(c2w):
    c2w = c2w.copy()
    c2w[2, :] *= -1
    c2w = c2w[np.array([1, 0, 2, 3]), :]
    c2w[0:3, 1:3] *= -1
    return c2w
    
def transform_points_opengl_to_opencv(points):
    new_points = points.copy()
    new_points[:, 2] *= -1
    new_points[:, :2] = new_points[:, [1, 0]]
    return new_points

def get_near_far(points, c2w, ixt, h, w):
    ext = np.linalg.inv(c2w)
    points = points @ ext[:3, :3].T + ext[:3, 3].T
    points = points @ ixt.T
    uv = points[:, :2] / points[:, 2:]
    uv_msk = (uv[:, 0] >= 0) & (uv[:, 0] < w) & (uv[:, 1] >= 0) & (uv[:, 1] < h)
    uv_msk = np.logical_and(uv_msk, points[:, 2] > 0)
    assert(uv_msk.sum() > 0)
    z_vals = points[..., 2][uv_msk]
    near_far = np.asarray([np.percentile(z_vals, 0.1), np.percentile(z_vals, 99.9)])
    return near_far

def cv2gl(c2w):
    c2w = c2w.copy()
    c2w[0:3, 1:3] *= -1
    c2w = c2w[np.array([1, 0, 2, 3]), :]
    c2w[2, :] *= -1
    return c2w

def gl2cv(c2w):
    c2w = c2w.copy()
    c2w[2, :] *= -1
    c2w = c2w[np.array([1, 0, 2, 3]), :]
    c2w[0:3, 1:3] *= -1
    return c2w

def load_path_pose(path, point_path):
    camera_path = json.load(open(path))
    h, w = camera_path['render_height'], camera_path['render_width']
    
    points = np.asarray(trimesh.load(point_path).vertices)
    points = transform_points_opengl_to_opencv(points)
    
    poses, ixts, near_fars = [], [], []
    for frame in camera_path['camera_path']:
        c2w = np.asarray(frame['camera_to_world']).reshape((4, 4))
        c2w = gl2cv(c2w)
        
        fov = frame['fov']
        # fx = 0.5 * w / np.tan(0.5 * fov / 180.0 * np.pi)
        fy = 0.5 * h / np.tan(0.5 * fov / 180.0 * np.pi)
        # aspect = frame['aspect']
        # 2 if aspect > 1.0:
            # fx = fy * aspect
        # fx = fx / aspect
        fx = fy
        # fy = fx
        # fx = fy
        ixt = np.asarray([fx, 0, w/2, 0, fy, h/2, 0, 0, 1]).reshape(3, 3)
        
        near_far = get_near_far(points, c2w, ixt, h, w)
        poses.append(np.linalg.inv(c2w).astype(np.float32))
        ixts.append(ixt.astype(np.float32))
        near_fars.append(near_far.astype(np.float32))
    return np.asarray(poses), np.asarray(ixts), np.asarray(near_fars), h, w


def read_bbox_renbody(data_root, frame_id, transform, bound_padding=0.1):
    pcd_path = join(data_root, 'pointcloud_denoise', '{:05d}.ply'.format(frame_id + 1))
    if os.path.exists(pcd_path): points = np.asarray(trimesh.load(pcd_path).vertices)
    else:
        ply_path = join(data_root, 'pointcloud', '{:05d}.ply'.format(frame_id + 1))
        points = np.asarray(read_and_denoise(ply_path).points)
        save_ply(points, ply_path.replace('pointcloud', 'pointcloud_denoise'))
    points = points @ transform[:3, :3].T + transform[:3, 3:].T
    bounds = np.array([points.min(0)-bound_padding, points.max(0)+bound_padding]).astype(np.float32)
    return bounds

def read_bbox_nhr(data_root, frame_id, transform, bound_padding=0.1):
    if 'NHR' in data_root: 
        if 'basketball' in data_root: pcd_path = join(data_root, 'vertices', '{:06d}.npy'.format(frame_id))
        else: pcd_path = join(data_root, 'vertices', '{}.npy'.format(frame_id))
    elif 'zju-mocap' in data_root: pcd_path = join(data_root, 'vertices', '{:06d}.npy'.format(frame_id))
    else: pass
    assert(os.path.exists(pcd_path))
    points = np.load(pcd_path)
    points = points @ transform[:3, :3].T + transform[:3, 3:].T
    bounds = np.array([points.min(0)-bound_padding, points.max(0)+bound_padding]).astype(np.float32)
    return bounds

def export_depth_to_pcd(depth, ixt, path='debug_depth.ply', crop_ratio=0.1):
    h, w = depth.shape
    crop_h, crop_w = int(crop_ratio * h), int(crop_ratio * w)
    ixt[0, 2] -= crop_w
    ixt[1, 2] -= crop_h
    tar_h, tar_w = int(h - crop_h*2), int(w - crop_w*2)
    depth = depth[crop_h:crop_h+tar_h, crop_w:crop_w+tar_w]
    h, w = depth.shape
    x = np.linspace(0, w-1, w)
    y = np.linspace(0, h-1, h)
    X, Y = np.meshgrid(x, y)
    X = X.reshape(-1)
    Y = Y.reshape(-1)
    Z = depth.reshape(-1)
    points = np.stack([X, Y, Z], axis=-1)
    points[..., :2] *= points[..., 2:]
    points = points @ np.linalg.inv(ixt).T
    save_ply(points, path)
lib/utils/data_utils.py
zju3dv/im4d
from __future__ import annotations
from copy import copy
from typing import Mapping, TypeVar, Union, Iterable, Callable, Dict
# these are generic type vars to tell mapping to accept any type vars when creating a type
KT = TypeVar("KT")  # key type
VT = TypeVar("VT")  # value type

# TODO: move this to engine implementation
# TODO: this is a special type just like Config
# ? However, dotdict is a general purpose data passing object, instead of just designed for config
# The only reason we defined those special variables are for type annotations
# If removed, all will still work flawlessly, just no editor annotation for output, type and meta
import torch

def return_dotdict(func: Callable):
    def inner(*args, **kwargs):
        return dotdict(func(*args, **kwargs))
    return inner


class dotdict(dict, Dict[KT, VT]):
    """
    This is the default data passing object used throughout the codebase
    Main function: dot access for dict values & dict like merging and updates

    a dictionary that supports dot notation 
    as well as dictionary access notation 
    usage: d = make_dotdict() or d = make_dotdict{'val1':'first'})
    set attributes: d.val2 = 'second' or d['val2'] = 'second'
    get attributes: d.val2 or d['val2']
    """

    def update(self, dct: Dict = None, **kwargs):
        dct = copy(dct)  # avoid modifying the original dict, use super's copy to avoid recursion

        # Handle different arguments
        if dct is None:
            dct = kwargs
        elif isinstance(dct, Mapping):
            dct.update(kwargs)
        else:
            super().update(dct, **kwargs)
            return

        # Recursive updates
        for k, v in dct.items():
            if k in self:

                # Handle type conversions
                target_type = type(self[k])
                if not isinstance(v, target_type):
                    # NOTE: bool('False') will be True
                    if target_type == bool and isinstance(v, str):
                        dct[k] = v == 'True'
                    else:
                        dct[k] = target_type(v)

                if isinstance(v, dict):
                    self[k].update(v)  # recursion from here
                else:
                    self[k] = v
            else:
                if isinstance(v, dict):
                    self[k] = dotdict(v)  # recursion?
                else:
                    self[k] = v
        return self

    def __init__(self, *args, **kwargs):
        self.update(*args, **kwargs)

    copy = return_dotdict(dict.copy)
    fromkeys = return_dotdict(dict.fromkeys)

    # def __hash__(self):
    #     # return hash(''.join([str(self.values().__hash__())]))
    #     return super(dotdict, self).__hash__()

    # def __init__(self, *args, **kwargs):
    #     super(dotdict, self).__init__(*args, **kwargs)

    """
    Uncomment following lines and 
    comment out __getattr__ = dict.__getitem__ to get feature:
    
    returns empty numpy array for undefined keys, so that you can easily copy things around
    TODO: potential caveat, harder to trace where this is set to np.array([], dtype=np.float32)
    """

    def __getitem__(self, key):
        try:
            return dict.__getitem__(self, key)
        except KeyError as e:
            raise AttributeError(e)
    # MARK: Might encounter exception in newer version of pytorch
    # Traceback (most recent call last):
    #   File "/home/xuzhen/miniconda3/envs/torch/lib/python3.9/multiprocessing/queues.py", line 245, in _feed
    #     obj = _ForkingPickler.dumps(obj)
    #   File "/home/xuzhen/miniconda3/envs/torch/lib/python3.9/multiprocessing/reduction.py", line 51, in dumps
    #     cls(buf, protocol).dump(obj)
    # KeyError: '__getstate__'
    # MARK: Because you allow your __getattr__() implementation to raise the wrong kind of exception.
    # FIXME: not working typing hinting code
    __getattr__: Callable[..., 'torch.Tensor'] = __getitem__  # type: ignore # overidden dict.__getitem__
    __getattribute__: Callable[..., 'torch.Tensor']  # type: ignore
    # __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

    # TODO: better ways to programmically define these special variables?

    @property
    def meta(self) -> dotdict:
        # Special variable used for storing cpu tensor in batch
        if 'meta' not in self:
            self.meta = dotdict()
        return self.__getitem__('meta')

    @meta.setter
    def meta(self, meta):
        self.__setitem__('meta', meta)

    @property
    def output(self) -> dotdict:  # late annotation needed for this
        # Special entry for storing output tensor in batch
        if 'output' not in self:
            self.output = dotdict()
        return self.__getitem__('output')

    @output.setter
    def output(self, output):
        self.__setitem__('output', output)

    @property
    def type(self) -> str:  # late annotation needed for this
        # Special entry for type based construction system
        return self.__getitem__('type')

    @type.setter
    def type(self, type):
        self.__setitem__('type', type)


class default_dotdict(dotdict):
    def __init__(self, default_type=object, *arg, **kwargs):
        super().__init__(*arg, **kwargs)
        dict.__setattr__(self, 'default_type', default_type)

    def __getitem__(self, key):
        try:
            return super().__getitem__(key)
        except (AttributeError, KeyError) as e:
            super().__setitem__(key, dict.__getattribute__(self, 'default_type')())
            return super().__getitem__(key)


context = dotdict()  # a global context object. Forgot why I did this. TODO: remove this


class Camera:
    # Helper class to manage camera parameters
    def __init__(self,
                 H: int = 512, 
                 W: int = 512,
                 K: torch.Tensor = torch.tensor([[512.0, 0.0, 256], [0.0, 512.0, 256.0], [0.0, 0.0, 1.0]]),  # intrinsics
                 R: torch.Tensor = torch.tensor([[-1.0, 0.0, 0.0,], [0.0, 0.0, -1.0,], [0.0, -1.0, 0.0,]]),  # extrinsics
                 T: torch.Tensor = torch.tensor([[0.0], [0.0], [-3.0],]),  # extrinsics
                 n: float = 0.002,  # bounds limit
                 f: float = 100,  # bounds limit
                 t: float = 0.0,  # temporal dimension (implemented as a float instead of int)
                 v: float = 0.0,  # view dimension (implemented as a float instead of int)
                 bounds: torch.Tensor = torch.tensor([[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]),  # bounding box

                 # camera update hyperparameters
                 origin: torch.Tensor = torch.tensor([0.0, 0.0, 0.0]),
                 world_up: torch.Tensor = torch.tensor([0.0, 0.0, 1.0]),
                 movement_speed: float = 1.0,  # gui movement speed

                 batch: dotdict = None,  # will ignore all other inputs
                 **kwargs,
                 ) -> None:

        # Batch (network input parameters)
        if batch is None:
            batch = dotdict()
            batch.H, batch.W, batch.K, batch.R, batch.T, batch.n, batch.f, batch.t, batch.v, batch.bounds = H, W, K, R, T, n, f, t, v, bounds
        self.from_batch(batch)

        # Other configurables
        self.origin = vec3(*origin)
        self.world_up = vec3(*world_up)
        self.movement_speed = movement_speed
        # self.front = self.front  # will trigger an update

        # Internal states to facilitate camera position change
        self.is_dragging = False  # rotation
        self.about_origin = False  # about origin rotation
        self.is_panning = False  # translation
        self.lock_fx_fy = True

    @property
    def w2p(self):
        ixt = mat4(self.ixt)
        ixt[3, 3] = 0
        ixt[2, 3] = 1
        return ixt @ self.ext  # w2c -> c2p = w2p

    @property
    def V(self): return self.c2w

    @property
    def ixt(self): return self.K

    @property
    def gl_ext(self):
        gl_c2w = self.c2w
        gl_c2w[0] *= 1  # flip x
        gl_c2w[1] *= -1  # flip y
        gl_c2w[2] *= -1  # flip z
        gl_ext = glm.affineInverse(gl_c2w)
        return gl_ext  # use original opencv ext since we've taken care of the intrinsics in gl_ixt

    @property
    def gl_ixt(self):
        # Construct opengl camera matrix with projection & clipping
        # https://fruty.io/2019/08/29/augmented-reality-with-opencv-and-opengl-the-tricky-projection-matrix/
        # https://gist.github.com/davegreenwood/3a32d779f81f08dce32f3bb423672191
        # fmt: off
        gl_ixt = mat4(
                      2 * self.fx / self.W,                          0,                                       0,  0,
                       2 * self.s / self.W,       2 * self.fy / self.H,                                       0,  0,
                1 - 2 * (self.cx / self.W), 2 * (self.cy / self.H) - 1,   (self.f + self.n) / (self.n - self.f), -1,
                                         0,                          0, 2 * self.f * self.n / (self.n - self.f),  0,
        )
        # fmt: on

        return gl_ixt

    @property
    def ext(self): return self.w2c

    @property
    def w2c(self):
        w2c = mat4(self.R)
        w2c[3] = vec4(*self.T, 1.0)
        return w2c

    @property
    def c2w(self):
        return glm.affineInverse(self.w2c)

    @property
    def right(self) -> vec3: return vec3(self.R[0, 0], self.R[1, 0], self.R[2, 0])  # c2w R, 0 -> 3,

    @property
    def down(self) -> vec3: return vec3(self.R[0, 1], self.R[1, 1], self.R[2, 1])  # c2w R, 1 -> 3,

    @property
    def front(self) -> vec3: return vec3(self.R[0, 2], self.R[1, 2], self.R[2, 2])  # c2w R, 2 -> 3,

    @front.setter
    def front(self, v: vec3):
        front = v  # the last row of R
        self.R[0, 2], self.R[1, 2], self.R[2, 2] = front.x, front.y, front.z
        right = glm.normalize(glm.cross(self.front, self.world_up))  # right
        self.R[0, 0], self.R[1, 0], self.R[2, 0] = right.x, right.y, right.z
        down = glm.cross(self.front, self.right)  # down
        self.R[0, 1], self.R[1, 1], self.R[2, 1] = down.x, down.y, down.z

    @property
    def center(self): return -glm.transpose(self.R) @ self.T  # 3,

    @center.setter
    def center(self, v: vec3):
        self.T = -self.R @ v  # 3, 1

    @property
    def s(self): return self.K[1, 0]

    @s.setter
    def s(self, s): self.K[1, 0] = s

    @property
    def fx(self): return self.K[0, 0]

    @fx.setter
    def fx(self, v: float):
        v = min(v, 1e5)
        v = max(v, 1e-3)
        if self.lock_fx_fy:
            self.K[1, 1] = v / self.K[0, 0] * self.K[1, 1]
        self.K[0, 0] = v

    @property
    def fy(self): return self.K[1, 1]

    @fy.setter
    def fy(self, v: float):
        if self.lock_fx_fy:
            self.K[0, 0] = v / self.K[1, 1] * self.K[0, 0]
        self.K[1, 1] = v

    @property
    def cx(self): return self.K[2, 0]

    @cx.setter
    def cx(self, v: float):
        self.K[2, 0] = v

    @property
    def cy(self): return self.K[2, 1]

    @cy.setter
    def cy(self, v: float):
        self.K[2, 1] = v

    def begin_dragging(self,
                       x: float, y: float,
                       is_panning: bool,
                       about_origin: bool,
                       ):
        self.is_dragging = True
        self.is_panning = is_panning
        self.about_origin = about_origin
        self.drag_start = vec2([x, y])

        # Record internal states # ? Will this make a copy?
        self.drag_start_front = self.front  # a recording
        self.drag_start_down = self.down
        self.drag_start_right = self.right
        self.drag_start_center = self.center
        self.drag_start_origin = self.origin
        self.drag_start_world_up = self.world_up

        # Need to find the max or min delta y to align with world_up
        dot = glm.dot(self.world_up, self.drag_start_front)
        self.drag_ymin = -np.arccos(-dot) + 0.01  # drag up, look down
        self.drag_ymax = np.pi + self.drag_ymin - 0.02  # remove the 0.01 of drag_ymin

    def end_dragging(self):
        self.is_dragging = False

    def update_dragging(self, x: float, y: float):
        if not self.is_dragging:
            return

        current = vec2(x, y)
        delta = current - self.drag_start
        delta *= self.movement_speed / max(self.H, self.W)
        delta *= -1

        if self.is_panning:
            center_delta = delta[0] * self.drag_start_right + delta[1] * self.drag_start_down
            self.center = self.drag_start_center + center_delta
            if self.about_origin:
                self.origin = self.drag_start_origin + center_delta
        else:
            m = mat4(1.0)
            m = glm.rotate(m, delta.x % 2 * np.pi, self.world_up)
            m = glm.rotate(m, np.clip(delta.y, self.drag_ymin, self.drag_ymax), self.drag_start_right)
            self.front = m @ self.drag_start_front  # might overshoot

            if self.about_origin:
                self.center = -m @ (self.origin - self.drag_start_center) + self.origin

    def move(self, x_offset: float, y_offset: float):
        speed_factor = 1e-1
        movement = y_offset * speed_factor
        movement = movement * self.front * self.movement_speed
        self.center += movement

        if self.is_dragging:
            self.drag_start_center += movement

    def to_batch(self):
        meta = dotdict()
        meta.H = torch.as_tensor(self.H)
        meta.W = torch.as_tensor(self.W)
        meta.K = torch.as_tensor(self.K.to_list()).mT
        meta.R = torch.as_tensor(self.R.to_list()).mT
        meta.T = torch.as_tensor(self.T.to_list())[..., None]
        meta.n = torch.as_tensor(self.n)
        meta.f = torch.as_tensor(self.f)
        meta.t = torch.as_tensor(self.t)
        meta.v = torch.as_tensor(self.v)
        meta.bounds = torch.as_tensor(self.bounds.to_list())  # no transpose for bounds

        batch = dotdict()
        batch.update(meta)
        batch.meta.update(meta)
        return batch

    def to_easymocap(self):
        batch = self.to_batch()
        camera = to_numpy(batch)
        return camera

    def from_batch(self, batch: dotdict):
        H, W, K, R, T, n, f, t, v, bounds = batch.H, batch.W, batch.K, batch.R, batch.T, batch.n, batch.f, batch.t, batch.v, batch.bounds

        # Batch (network input parameters)
        self.H = int(H)
        self.W = int(W)
        self.K = mat3(*K.mT.ravel())
        self.R = mat3(*R.mT.ravel())
        self.T = vec3(*T.ravel())  # 3,
        self.n = float(n)
        self.f = float(f)
        self.t = float(t)
        self.v = float(v)
        self.bounds = mat2x3(*bounds.ravel())  # 2, 3
        return self

    def custom_pose(self, R: torch.Tensor, T: torch.Tensor, K: torch.Tensor):
        # self.K = mat3(*K.mT.ravel())
        self.R = mat3(*R.mT.ravel())
        self.T = vec3(*T.ravel())

    def from_easymocap(self, camera: dict):
        batch = to_tensor(camera)
        self.from_batch(batch)
        return self

lib/utils/easyvv/base_utils.py
zju3dv/im4d
# easymocap utility functions
import os
import cv2
import numpy as np
from lib.utils.easyvv.base_utils import dotdict


class FileStorage(object):
    def __init__(self, filename, isWrite=False):
        version = cv2.__version__
        self.major_version = int(version.split('.')[0])
        self.second_version = int(version.split('.')[1])

        if isWrite:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            self.fs = open(filename, 'w')
            self.fs.write('%YAML:1.0\r\n')
            self.fs.write('---\r\n')
        else:
            assert os.path.exists(filename), filename
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_READ)
        self.isWrite = isWrite

    def __del__(self):
        if self.isWrite:
            self.fs.close()
        else:
            cv2.FileStorage.release(self.fs)

    def _write(self, out):
        self.fs.write(out + '\r\n')

    def write(self, key, value, dt='mat'):
        if dt == 'mat':
            self._write('{}: !!opencv-matrix'.format(key))
            self._write('  rows: {}'.format(value.shape[0]))
            self._write('  cols: {}'.format(value.shape[1]))
            self._write('  dt: d')
            self._write('  data: [{}]'.format(', '.join(['{:.10f}'.format(i) for i in value.reshape(-1)])))
        elif dt == 'list':
            self._write('{}:'.format(key))
            for elem in value:
                self._write('  - "{}"'.format(elem))
        elif dt == 'real':
            self._write('{}: {:.10f}'.format(key, value))  # as accurate as possible
        else:
            raise NotImplementedError

    def read(self, key, dt='mat'):
        if dt == 'mat':
            output = self.fs.getNode(key).mat()
        elif dt == 'list':
            results = []
            n = self.fs.getNode(key)
            for i in range(n.size()):
                val = n.at(i).string()
                if val == '':
                    val = str(int(n.at(i).real()))
                if val != 'none':
                    results.append(val)
            output = results
        elif dt == 'real':
            output = self.fs.getNode(key).real()
        else:
            raise NotImplementedError
        return output

    def close(self):
        self.__del__(self)


def read_camera(intri_path: str, extri_path: str, cam_names=[], read_rvec=True):
    assert os.path.exists(intri_path), intri_path
    assert os.path.exists(extri_path), extri_path

    intri = FileStorage(intri_path)
    extri = FileStorage(extri_path)
    cams = dotdict()
    cam_names = intri.read('names', dt='list')
    for cam in cam_names:
        # Intrinsics
        cams[cam] = dotdict()
        cams[cam].K = intri.read('K_{}'.format(cam))
        cams[cam].H = intri.read('H_{}'.format(cam), dt='real') or -1
        cams[cam].W = intri.read('W_{}'.format(cam), dt='real') or -1
        cams[cam].invK = np.linalg.inv(cams[cam]['K'])

        # Extrinsics
        Tvec = extri.read('T_{}'.format(cam))
        Rvec = extri.read('R_{}'.format(cam))
        if Rvec is not None and read_rvec: R = cv2.Rodrigues(Rvec)[0]
        else:
            R = extri.read('Rot_{}'.format(cam))
            Rvec = cv2.Rodrigues(R)[0]
        RT = np.hstack((R, Tvec))

        cams[cam].R = R
        cams[cam].T = Tvec
        cams[cam].C = - Rvec.T @ Tvec
        cams[cam].RT = RT
        cams[cam].Rvec = Rvec
        cams[cam].P = cams[cam].K @ cams[cam].RT

        # Distortion
        D = intri.read('D_{}'.format(cam))
        if D is None: D = intri.read('dist_{}'.format(cam))
        cams[cam].D = D

        # Time input
        cams[cam].t = extri.read('t_{}'.format(cam), dt='real') or 0  # temporal index, might all be 0

        # Bounds, could be overwritten
        cams[cam].n = extri.read('n_{}'.format(cam), dt='real') or 0.0001  # temporal index, might all be 0
        cams[cam].f = extri.read('f_{}'.format(cam), dt='real') or 1e6  # temporal index, might all be 0
        cams[cam].bounds = extri.read('bounds_{}'.format(cam))
        cams[cam].bounds = np.array([[-1e6, -1e6, -1e6], [1e6, 1e6, 1e6]]) if cams[cam].bounds is None else cams[cam].bounds

        # CCM
        cams[cam].ccm = intri.read('ccm_{}'.format(cam))
        cams[cam].ccm = np.eye(3) if cams[cam].ccm is None else cams[cam].ccm

    return dotdict(cams)


def write_camera(cameras: dict, path: str, intri_path: str = '', extri_path: str = ''):
    from os.path import join
    os.makedirs(path, exist_ok=True)
    if not intri_path or not extri_path:
        intri_name = join(path, 'intri.yml')  # TODO: make them arguments
        extri_name = join(path, 'extri.yml')
    intri = FileStorage(intri_name, True)
    extri = FileStorage(extri_name, True)
    cam_names = [key_.split('.')[0] for key_ in cameras.keys()]
    intri.write('names', cam_names, 'list')
    extri.write('names', cam_names, 'list')

    cameras = dotdict(cameras)
    for key_, val in cameras.items():
        if key_ == 'basenames': continue
        key = key_.split('.')[0]
        # Intrinsics
        intri.write('K_{}'.format(key), val.K)
        if 'H' in val: intri.write('H_{}'.format(key), val.H, 'real')
        if 'W' in val: intri.write('W_{}'.format(key), val.W, 'real')

        # Distortion
        if 'D' not in val:
            if 'dist' in val: val.D = val.dist
            else: val.D = np.zeros((5, 1))
        intri.write('D_{}'.format(key), val.D.reshape(5, 1))

        # Extrinsics
        if 'R' not in val: val.R = cv2.Rodrigues(val.Rvec)[0]
        if 'Rvec' not in val: val.Rvec = cv2.Rodrigues(val.R)[0]
        extri.write('R_{}'.format(key), val.Rvec)
        extri.write('Rot_{}'.format(key), val.R)
        extri.write('T_{}'.format(key), val.T.reshape(3, 1))

        # Temporal
        if 't' in val: extri.write('t_{}'.format(key), val.t, 'real')

        # Bounds
        if 'n' in val: extri.write('n_{}'.format(key), val.n, 'real')
        if 'f' in val: extri.write('f_{}'.format(key), val.f, 'real')
        if 'bounds' in val: extri.write('bounds_{}'.format(key), val.bounds)

        # Color correction matrix
        if 'ccm' in val: intri.write('ccm_{}'.format(key), val.ccm)

lib/utils/easyvv/easy_utils.py
zju3dv/im4d
import torch

class FastRendUtils:
    @staticmethod
    def prepare_wpts(wbounds, grid_res, subgrid_res):
        grid_res = torch.tensor(grid_res, device=wbounds.device, dtype=torch.float32)
        globol_domain_min, global_domain_max = wbounds[0], wbounds[1]
        global_domain_size = global_domain_max - globol_domain_min
        voxel_size = global_domain_size / grid_res
        voxel_offset_min = globol_domain_min
        voxel_offset_max = voxel_offset_min + voxel_size
        voxel_samples = []
        for dim in range(3): # TODO: for subgrid resolution
            voxel_samples.append(torch.linspace(voxel_offset_min[dim], voxel_offset_max[dim], subgrid_res[dim], device=wbounds.device))
        voxel_samples = torch.stack(torch.meshgrid(*voxel_samples, indexing='ij'), axis=-1).reshape(-1, 3)
        voxel_samples = voxel_samples[None, None, None] # [1, 1, 1, r, 3]
        voxel_ranges = []
        for dim in range(3):
            voxel_ranges.append(torch.arange(0, int(grid_res[dim].item()), device=wbounds.device))    
        voxel_grid = torch.stack(torch.meshgrid(*voxel_ranges, indexing='ij'), axis=-1)
        voxel_grid = (voxel_grid * voxel_size)[..., None, :] # [r1, r2, r3, -1, 3]
        pts = voxel_grid + voxel_samples # [r1, r2, r3, r, 3]
        return pts

lib/utils/im4d/fastrend_utils.py
zju3dv/im4d
import torch
import numpy as np
import cv2
    
import trimesh
class Im4DUtils:
    @staticmethod
    def render_median_depth(weights, z_vals):
        weights = weights[..., None]
        z_vals = z_vals[..., None]
        cumulative_weights = torch.cumsum(weights[..., 0], dim=-1)
        split = torch.ones((*weights.shape[:-2], 1), device=weights.device) * 0.5  # [..., 1]
        median_index = torch.searchsorted(cumulative_weights, split, side="left")  # [..., 1]
        median_index = torch.clamp(median_index, 0, z_vals.shape[-2] - 1)  # [..., 1]
        median_depth = torch.gather(z_vals[..., 0], dim=-1, index=median_index)  # [..., 1]
        return median_depth[..., 0]
    
    @staticmethod
    def extract_mesh(pts, msks, hws, rgbs, mask_at_boxs, export_path, erode_pixel=15):
        points, colors = [], []
        for pt, msk, hw, rgb, mask_at_box in zip(pts, msks, hws, rgbs, mask_at_boxs):
            h, w = hw
            if erode_pixel >= 2:
                errode_msk = torch.zeros((h, w), device=msk.device).float()
                errode_msk[mask_at_box.reshape(h, w)!=0] = msk
                errode_msk = (errode_msk > 0.5).detach().cpu().numpy().astype(np.uint8)
                errode_msk = errode_msk.reshape(h, w)
                kernel = np.ones((erode_pixel,erode_pixel),np.uint8)
                errode_msk = cv2.erode(errode_msk,kernel,iterations = 1)
                msk = errode_msk[mask_at_box.cpu().numpy().reshape(h, w)!=0]
            points.append(pt.detach().cpu().numpy()[msk.nonzero()])
            colors.append(rgb.detach().cpu().numpy()[msk.nonzero()])
        points = np.concatenate(points, axis=0)
        colors = (np.concatenate(colors, axis=0) * 255.).astype(np.uint8)
        mesh = trimesh.PointCloud(points, colors=colors)
        mesh.export(export_path)
        
    @staticmethod
    def ndc2world(points, f, H, W):
        points = points.clone()
        points[..., 2] = torch.clamp_max_(points[..., 2], 0.9999)
        points[..., 2] = 2 / (points[..., 2] -1)
        points[..., 0] = -points[..., 0] * points[..., 2] * W / (2 * f)
        points[..., 1] = -points[..., 1] * points[..., 2] * H / (2 * f)
        return points
        
lib/utils/im4d/im4d_utils.py
zju3dv/im4d
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})")

lib/utils/kplanes/intrinsics.py
zju3dv/im4d
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)

lib/utils/kplanes/lr_scheduling.py
zju3dv/im4d
from typing import Tuple, Optional

import numpy as np
import torch

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) -> 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
    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):
    """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
    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

lib/utils/kplanes/ray_utils.py
zju3dv/im4d
import numpy as np
import os, imageio
from lib.utils.kplanes import ray_utils

########## Slightly modified version of LLFF data loading code
##########  see https://github.com/Fyusion/LLFF for original

def _load_data(basedir, factor=None, width=None, height=None, load_imgs=False):
    poses_arr = np.load(os.path.join(basedir, 'poses_bounds.npy'))
    poses = poses_arr[:, :-2].reshape([-1, 3, 5]).transpose([1,2,0])
    bds = poses_arr[:, -2:].transpose([1,0])
    # sh = imageio.imread(imgfiles[0]).shape
    # poses[:2, 4, :] = np.array(sh[:2]).reshape([2, 1])
    poses[2, 4, :] = poses[2, 4, :] * 1./factor
    if not load_imgs:
        return poses, bds

def normalize(x):
    return x / np.linalg.norm(x)

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

def ptstocam(pts, c2w):
    tt = np.matmul(c2w[:3,:3].T, (pts-c2w[:3,3])[...,np.newaxis])[...,0]
    return tt

def poses_avg(poses):

    hwf = poses[0, :3, -1:]

    center = poses[:, :3, 3].mean(0)
    vec2 = normalize(poses[:, :3, 2].sum(0))
    up = poses[:, :3, 1].sum(0)
    c2w = np.concatenate([viewmatrix(vec2, up, center), hwf], 1)

    return c2w



def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, rots, N):
    render_poses = []
    rads = np.array(list(rads) + [1.])
    hwf = c2w[:,4:5]

    for theta in np.linspace(0., 2. * np.pi * rots, N+1)[:-1]:
        c = np.dot(c2w[:3,:4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta*zrate), 1.]) * rads)
        z = normalize(c - np.dot(c2w[:3,:4], np.array([0,0,-focal, 1.])))
        render_poses.append(np.concatenate([viewmatrix(z, up, c), hwf], 1))
    return render_poses



def recenter_poses(poses):

    poses_ = poses+0
    bottom = np.reshape([0,0,0,1.], [1,4])
    c2w = poses_avg(poses)
    c2w = np.concatenate([c2w[:3,:4], bottom], -2)
    bottom = np.tile(np.reshape(bottom, [1,1,4]), [poses.shape[0],1,1])
    poses = np.concatenate([poses[:,:3,:4], bottom], -2)

    poses = np.linalg.inv(c2w) @ poses
    poses_[:,:3,:4] = poses[:,:3,:4]
    poses = poses_
    return poses


#####################


def spherify_poses(poses, bds):

    p34_to_44 = lambda p : np.concatenate([p, np.tile(np.reshape(np.eye(4)[-1,:], [1,1,4]), [p.shape[0], 1,1])], 1)

    rays_d = poses[:,:3,2:3]
    rays_o = poses[:,:3,3:4]

    def min_line_dist(rays_o, rays_d):
        A_i = np.eye(3) - rays_d * np.transpose(rays_d, [0,2,1])
        b_i = -A_i @ rays_o
        pt_mindist = np.squeeze(-np.linalg.inv((np.transpose(A_i, [0,2,1]) @ A_i).mean(0)) @ (b_i).mean(0))
        return pt_mindist

    pt_mindist = min_line_dist(rays_o, rays_d)

    center = pt_mindist
    up = (poses[:,:3,3] - center).mean(0)

    vec0 = normalize(up)
    vec1 = normalize(np.cross([.1,.2,.3], vec0))
    vec2 = normalize(np.cross(vec0, vec1))
    pos = center
    c2w = np.stack([vec1, vec2, vec0, pos], 1)

    poses_reset = np.linalg.inv(p34_to_44(c2w[None])) @ p34_to_44(poses[:,:3,:4])

    rad = np.sqrt(np.mean(np.sum(np.square(poses_reset[:,:3,3]), -1)))

    sc = 1./rad
    poses_reset[:,:3,3] *= sc
    bds *= sc
    rad *= sc

    centroid = np.mean(poses_reset[:,:3,3], 0)
    zh = centroid[2]
    radcircle = np.sqrt(rad**2-zh**2)
    new_poses = []

    for th in np.linspace(0.,2.*np.pi, 120):

        camorigin = np.array([radcircle * np.cos(th), radcircle * np.sin(th), zh])
        up = np.array([0,0,-1.])

        vec2 = normalize(camorigin)
        vec0 = normalize(np.cross(vec2, up))
        vec1 = normalize(np.cross(vec2, vec0))
        pos = camorigin
        p = np.stack([vec0, vec1, vec2, pos], 1)

        new_poses.append(p)

    new_poses = np.stack(new_poses, 0)

    new_poses = np.concatenate([new_poses, np.broadcast_to(poses[0,:3,-1:], new_poses[:,:3,-1:].shape)], -1)
    poses_reset = np.concatenate([poses_reset[:,:3,:4], np.broadcast_to(poses[0,:3,-1:], poses_reset[:,:3,-1:].shape)], -1)

    return poses_reset, new_poses, bds

def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, rots, N):
    render_poses = []
    rads = np.array(list(rads) + [1.])
    # hwf = c2w[:,4:5]

    for theta in np.linspace(0., 2. * np.pi * rots, N+1)[:-1]:
        c = np.dot(c2w[:3,:4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta*zrate), 1.]) * rads)
        z = normalize(c - np.dot(c2w[:3,:4], np.array([0,0,-focal, 1.])))
        render_poses.append(viewmatrix(z, up, c))
    return render_poses

def generate_path(poses, bds, c2w, path_zflat=False):
    print('recentered', c2w.shape)
    print(c2w[:3,:4])
        ## Get spiral
        # Get average pose
    up = normalize(poses[:, :3, 1].sum(0))

        # Find a reasonable "focus depth" for this dataset
    close_depth, inf_depth = bds.min()*.9, bds.max()*5.
    dt = .75
    mean_dz = 1./(((1.-dt)/close_depth + dt/inf_depth))
    focal = mean_dz

    # Get radii for spiral path
    shrink_factor = .8
    zdelta = close_depth * .2
    tt = poses[:,:3,3] # ptstocam(poses[:3,3,:].T, c2w).T
    rads = np.percentile(np.abs(tt), 90, 0)
    c2w_path = c2w
    N_views = 120
    N_rots = 2
    if path_zflat:
        zloc = np.percentile(tt, 10, 0)[2]
        zloc = -close_depth * .1
        c2w_path[:3,3] = c2w_path[:3,3] + zloc * c2w_path[:3,2]
        rads[2] = 0.
        N_rots = 1
        N_views/=2

    # Generate poses for spiral path
    render_poses = render_path_spiral(c2w_path, up, rads, focal, zdelta, zrate=.5, rots=N_rots, N=N_views)
    return render_poses


def load_llff_data(basedir, recenter=True, bd_factor=.95, spherify=False, path_zflat=False):
    poses, bds = _load_data(basedir, factor=1) # factor=8 downsamples original imgs by 8x
    print('Loaded', basedir, bds.min(), bds.max())
    # Correct rotation matrix ordering and move variable dim to axis 0
    poses = np.concatenate([poses[:, 1:2, :], -poses[:, 0:1, :], poses[:, 2:, :]], 1)
    poses = np.moveaxis(poses, -1, 0).astype(np.float32)
    bds = np.moveaxis(bds, -1, 0).astype(np.float32)
    hwf = poses[:, :3, 4:]
    # Rescale if bd_factor is provided
    # sc = 1. if bd_factor is None else 1./(bds.min() * bd_factor)
    # poses[:,:3,3] *= sc
    # bds *= sc
    assert(recenter)
    if recenter:
        poses, pose_avg = ray_utils.center_poses(poses[:, :3, :4])

    near_original = np.min(bds)
    scale_factor = near_original * bd_factor
    bds /= scale_factor
    poses[..., 3] /= scale_factor

    render_poses = ray_utils.generate_spiral_path(poses, bds)
    c2w = pose_avg
    print('Data:')
    print(poses.shape, bds.shape)
    dists = np.sum(np.square(c2w[:3,3] - poses[:,:3,3]), -1)
    i_test = np.argmin(dists)
    print('HOLDOUT view is', i_test)
    poses = poses.astype(np.float32)
    poses = np.concatenate([poses, hwf], axis=-1)
    return poses, bds, render_poses, i_test
lib/utils/llff/load_llff.py
zju3dv/im4d
import os
import requests
import json

def send_msg(msg, cfg, config_file='configs/local/dingtalk.txt'):
    if os.path.exists(config_file):
        with open(config_file, 'r') as f:
            webhook = f.readline().strip()
            proxy = f.readline().strip()
    else:
        return False
    if cfg.log_level not in ['INFO', 'WARNINING'] or cfg.local_rank != 0: 
        return False
    header = {
        'Content-Type': 'application/json',
        'Charset': 'UTF-8'
    }
    message ={
        'msgtype': 'text',
        'text': {
            'content': msg
        },
    }
    message_json = json.dumps(message)
    print('send msg begin')
    os.environ['http_proxy'] = proxy
    os.environ['https_proxy'] = proxy
    try:
        info = requests.post(url=webhook, data=message_json,headers=header, timeout=5)
        print('send msg end')
        os.environ['http_proxy'] = ''
        os.environ['https_proxy'] = ''
        return True
    except:
        print('send msg error')
        os.environ['http_proxy'] = ''
        os.environ['https_proxy'] = ''
        return False
    
lib/utils/msg_utils.py
zju3dv/im4d
import torch
import os
from torch import nn
import numpy as np
import torch.nn.functional
from collections import OrderedDict
from termcolor import colored
import sys
import yaml
from lib.config import cfg


def gen_rays_bbox(rays, bounds):
    rays_o, rays_d = rays[..., :3], rays[..., 3:6]
    norm_d = torch.norm(rays_d, dim=-1, keepdim=True)
    viewdir = rays_d / norm_d
    viewdir[(viewdir < 1e-5) & (viewdir > -1e-10)] = 1e-5
    viewdir[(viewdir > -1e-5) & (viewdir < 1e-10)] = -1e-5

    tmin = (bounds[0:1] - rays_o[:1]) / viewdir
    tmax = (bounds[1:2] - rays_o[:1]) / viewdir
    t1 = torch.min(tmin, tmax)
    t2 = torch.max(tmin, tmax)

    near = torch.max(t1, dim=-1)[0]
    far = torch.min(t2, dim=-1)[0]
    mask_at_box = near < far
    return mask_at_box

import time
class perf_timer:
    def __init__(self, msg="Elapsed time: {}s", logf=lambda x: print(colored(x, 'yellow')), sync_cuda=True, use_ms=False, disabled=False):
        self.logf = logf
        self.msg = msg
        self.sync_cuda = sync_cuda
        self.use_ms = use_ms
        self.disabled = disabled

        self.loggedtime = None

    def __enter__(self,):
        if self.sync_cuda:
            torch.cuda.synchronize()
        self.start = time.perf_counter()

    def __exit__(self, exc_type, exc_value, traceback):
        if self.sync_cuda:
            torch.cuda.synchronize()
        self.logtime(self.msg)

    def logtime(self, msg=None, logf=None):
        if self.disabled:
            return
        # SAME CLASS, DIFFERENT FUNCTIONALITY, is this good?
        # call the logger for timing code sections
        if self.sync_cuda:
            torch.cuda.synchronize()

        # always remember current time
        prev = self.loggedtime
        self.loggedtime = time.perf_counter()

        # print it if we've remembered previous time
        if prev is not None and msg:
            logf = logf or self.logf
            diff = self.loggedtime-prev
            diff *= 1000 if self.use_ms else 1
            logf(msg.format(diff))

        return self.loggedtime

def sigmoid(x):
    y = torch.clamp(x.sigmoid(), min=1e-4, max=1 - 1e-4)
    return y


def _neg_loss(pred, gt):
    ''' Modified focal loss. Exactly the same as CornerNet.
        Runs faster and costs a little bit more memory
        Arguments:
            pred (batch x c x h x w)
            gt_regr (batch x c x h x w)
    '''
    pos_inds = gt.eq(1).float()
    neg_inds = gt.lt(1).float()

    neg_weights = torch.pow(1 - gt, 4)

    loss = 0

    pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
    neg_loss = torch.log(1 - pred) * torch.pow(pred,
                                               2) * neg_weights * neg_inds

    num_pos = pos_inds.float().sum()
    pos_loss = pos_loss.sum()
    neg_loss = neg_loss.sum()

    if num_pos == 0:
        loss = loss - neg_loss
    else:
        loss = loss - (pos_loss + neg_loss) / num_pos
    return loss


class FocalLoss(nn.Module):
    '''nn.Module warpper for focal loss'''
    def __init__(self):
        super(FocalLoss, self).__init__()
        self.neg_loss = _neg_loss

    def forward(self, out, target):
        return self.neg_loss(out, target)


def smooth_l1_loss(vertex_pred,
                   vertex_targets,
                   vertex_weights,
                   sigma=1.0,
                   normalize=True,
                   reduce=True):
    """
    :param vertex_pred:     [b, vn*2, h, w]
    :param vertex_targets:  [b, vn*2, h, w]
    :param vertex_weights:  [b, 1, h, w]
    :param sigma:
    :param normalize:
    :param reduce:
    :return:
    """
    b, ver_dim, _, _ = vertex_pred.shape
    sigma_2 = sigma**2
    vertex_diff = vertex_pred - vertex_targets
    diff = vertex_weights * vertex_diff
    abs_diff = torch.abs(diff)
    smoothL1_sign = (abs_diff < 1. / sigma_2).detach().float()
    in_loss = torch.pow(diff, 2) * (sigma_2 / 2.) * smoothL1_sign \
              + (abs_diff - (0.5 / sigma_2)) * (1. - smoothL1_sign)

    if normalize:
        in_loss = torch.sum(in_loss.view(b, -1), 1) / (
            ver_dim * torch.sum(vertex_weights.view(b, -1), 1) + 1e-3)

    if reduce:
        in_loss = torch.mean(in_loss)

    return in_loss


class SmoothL1Loss(nn.Module):
    def __init__(self):
        super(SmoothL1Loss, self).__init__()
        self.smooth_l1_loss = smooth_l1_loss

    def forward(self,
                preds,
                targets,
                weights,
                sigma=1.0,
                normalize=True,
                reduce=True):
        return self.smooth_l1_loss(preds, targets, weights, sigma, normalize,
                                   reduce)


class AELoss(nn.Module):
    def __init__(self):
        super(AELoss, self).__init__()

    def forward(self, ae, ind, ind_mask):
        """
        ae: [b, 1, h, w]
        ind: [b, max_objs, max_parts]
        ind_mask: [b, max_objs, max_parts]
        obj_mask: [b, max_objs]
        """
        # first index
        b, _, h, w = ae.shape
        b, max_objs, max_parts = ind.shape
        obj_mask = torch.sum(ind_mask, dim=2) != 0

        ae = ae.view(b, h * w, 1)
        seed_ind = ind.view(b, max_objs * max_parts, 1)
        tag = ae.gather(1, seed_ind).view(b, max_objs, max_parts)

        # compute the mean
        tag_mean = tag * ind_mask
        tag_mean = tag_mean.sum(2) / (ind_mask.sum(2) + 1e-4)

        # pull ae of the same object to their mean
        pull_dist = (tag - tag_mean.unsqueeze(2)).pow(2) * ind_mask
        obj_num = obj_mask.sum(dim=1).float()
        pull = (pull_dist.sum(dim=(1, 2)) / (obj_num + 1e-4)).sum()
        pull /= b

        # push away the mean of different objects
        push_dist = torch.abs(tag_mean.unsqueeze(1) - tag_mean.unsqueeze(2))
        push_dist = 1 - push_dist
        push_dist = nn.functional.relu(push_dist, inplace=True)
        obj_mask = (obj_mask.unsqueeze(1) + obj_mask.unsqueeze(2)) == 2
        push_dist = push_dist * obj_mask.float()
        push = ((push_dist.sum(dim=(1, 2)) - obj_num) /
                (obj_num * (obj_num - 1) + 1e-4)).sum()
        push /= b
        return pull, push


class PolyMatchingLoss(nn.Module):
    def __init__(self, pnum):
        super(PolyMatchingLoss, self).__init__()

        self.pnum = pnum
        batch_size = 1
        pidxall = np.zeros(shape=(batch_size, pnum, pnum), dtype=np.int32)
        for b in range(batch_size):
            for i in range(pnum):
                pidx = (np.arange(pnum) + i) % pnum
                pidxall[b, i] = pidx

        device = torch.device('cuda')
        pidxall = torch.from_numpy(
            np.reshape(pidxall, newshape=(batch_size, -1))).to(device)

        self.feature_id = pidxall.unsqueeze_(2).long().expand(
            pidxall.size(0), pidxall.size(1), 2).detach()

    def forward(self, pred, gt, loss_type="L2"):
        pnum = self.pnum
        batch_size = pred.size()[0]
        feature_id = self.feature_id.expand(batch_size,
                                            self.feature_id.size(1), 2)
        device = torch.device('cuda')

        gt_expand = torch.gather(gt, 1,
                                 feature_id).view(batch_size, pnum, pnum, 2)

        pred_expand = pred.unsqueeze(1)

        dis = pred_expand - gt_expand

        if loss_type == "L2":
            dis = (dis**2).sum(3).sqrt().sum(2)
        elif loss_type == "L1":
            dis = torch.abs(dis).sum(3).sum(2)

        min_dis, min_id = torch.min(dis, dim=1, keepdim=True)
        # print(min_id)

        # min_id = torch.from_numpy(min_id.data.cpu().numpy()).to(device)
        # min_gt_id_to_gather = min_id.unsqueeze_(2).unsqueeze_(3).long().\
        #                         expand(min_id.size(0), min_id.size(1), gt_expand.size(2), gt_expand.size(3))
        # gt_right_order = torch.gather(gt_expand, 1, min_gt_id_to_gather).view(batch_size, pnum, 2)

        return torch.mean(min_dis)


class AttentionLoss(nn.Module):
    def __init__(self, beta=4, gamma=0.5):
        super(AttentionLoss, self).__init__()

        self.beta = beta
        self.gamma = gamma

    def forward(self, pred, gt):
        num_pos = torch.sum(gt)
        num_neg = torch.sum(1 - gt)
        alpha = num_neg / (num_pos + num_neg)
        edge_beta = torch.pow(self.beta, torch.pow(1 - pred, self.gamma))
        bg_beta = torch.pow(self.beta, torch.pow(pred, self.gamma))

        loss = 0
        loss = loss - alpha * edge_beta * torch.log(pred) * gt
        loss = loss - (1 - alpha) * bg_beta * torch.log(1 - pred) * (1 - gt)
        return torch.mean(loss)


def _gather_feat(feat, ind, mask=None):
    dim = feat.size(2)
    ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
    feat = feat.gather(1, ind)
    if mask is not None:
        mask = mask.unsqueeze(2).expand_as(feat)
        feat = feat[mask]
        feat = feat.view(-1, dim)
    return feat


def _tranpose_and_gather_feat(feat, ind):
    feat = feat.permute(0, 2, 3, 1).contiguous()
    feat = feat.view(feat.size(0), -1, feat.size(3))
    feat = _gather_feat(feat, ind)
    return feat


class Ind2dRegL1Loss(nn.Module):
    def __init__(self, type='l1'):
        super(Ind2dRegL1Loss, self).__init__()
        if type == 'l1':
            self.loss = torch.nn.functional.l1_loss
        elif type == 'smooth_l1':
            self.loss = torch.nn.functional.smooth_l1_loss

    def forward(self, output, target, ind, ind_mask):
        """ind: [b, max_objs, max_parts]"""
        b, max_objs, max_parts = ind.shape
        ind = ind.view(b, max_objs * max_parts)
        pred = _tranpose_and_gather_feat(output,
                                         ind).view(b, max_objs, max_parts,
                                                   output.size(1))
        mask = ind_mask.unsqueeze(3).expand_as(pred)
        loss = self.loss(pred * mask, target * mask, reduction='sum')
        loss = loss / (mask.sum() + 1e-4)
        return loss


class IndL1Loss1d(nn.Module):
    def __init__(self, type='l1'):
        super(IndL1Loss1d, self).__init__()
        if type == 'l1':
            self.loss = torch.nn.functional.l1_loss
        elif type == 'smooth_l1':
            self.loss = torch.nn.functional.smooth_l1_loss

    def forward(self, output, target, ind, weight):
        """ind: [b, n]"""
        output = _tranpose_and_gather_feat(output, ind)
        weight = weight.unsqueeze(2)
        loss = self.loss(output * weight, target * weight, reduction='sum')
        loss = loss / (weight.sum() * output.size(2) + 1e-4)
        return loss


class GeoCrossEntropyLoss(nn.Module):
    def __init__(self):
        super(GeoCrossEntropyLoss, self).__init__()

    def forward(self, output, target, poly):
        output = torch.nn.functional.softmax(output, dim=1)
        output = torch.log(torch.clamp(output, min=1e-4))
        poly = poly.view(poly.size(0), 4, poly.size(1) // 4, 2)
        target = target[..., None, None].expand(poly.size(0), poly.size(1), 1,
                                                poly.size(3))
        target_poly = torch.gather(poly, 2, target)
        sigma = (poly[:, :, 0] - poly[:, :, 1]).pow(2).sum(2, keepdim=True)
        kernel = torch.exp(-(poly - target_poly).pow(2).sum(3) / (sigma / 3))
        loss = -(output * kernel.transpose(2, 1)).sum(1).mean()
        return loss


def load_model(net,
               optim,
               scheduler,
               recorder,
               model_dir,
               resume=True,
               epoch=-1,
               strict=True):
    if not resume:
        print(colored('remove contents of model directory %s' % model_dir, 'red'))
        assert(model_dir != '/')
        os.system('rm -rf {}'.format(os.path.join(model_dir, '*')))
        return 0

    if not os.path.exists(model_dir):
        return 0

    pths = [
        int(pth.split('.')[0]) for pth in os.listdir(model_dir)
        if pth != 'latest.pth'
    ]
    if len(pths) == 0 and 'latest.pth' not in os.listdir(model_dir):
        return 0
    if epoch == -1:
        if 'latest.pth' in os.listdir(model_dir):
            pth = 'latest'
        else:
            pth = max(pths)
    else:
        pth = epoch
    print('load model: {}'.format(os.path.join(model_dir,
                                               '{}.pth'.format(pth))))
    pretrained_model = torch.load(
        os.path.join(model_dir, '{}.pth'.format(pth)), 'cpu')
    net.load_state_dict(pretrained_model['net'], strict=strict)
    if 'optim' in pretrained_model:
        try:
            optim.load_state_dict(pretrained_model['optim'])
            scheduler.load_state_dict(pretrained_model['scheduler'])
            recorder.load_state_dict(pretrained_model['recorder'])
            return pretrained_model['epoch'] + 1
        except:
            if cfg.bkgd_render:
                print('Warning: optimizer is not loaded, but okay.')
                return pretrained_model['epoch'] + 1
            else:
                print('Warning: optimizer is not loaded, but okay.')
                return pretrained_model['epoch'] + 1
                print('Error: optimizer is not loaded, not okay.')
                sys.exit(1)
    else:
        return 0


def save_model(net, optim, scheduler, recorder, model_dir, epoch, last=False):
    os.system('mkdir -p {}'.format(model_dir))
    model = {
        'net': net.state_dict(),
        'epoch': epoch
    }
    model_last = {
        'net': net.state_dict(),
        'optim': optim.state_dict(),
        'scheduler': scheduler.state_dict(),
        'recorder': recorder.state_dict(),
        'epoch': epoch
    }
    if last:
        torch.save(model_last, os.path.join(model_dir, 'latest.pth'))
    else:
        torch.save(model, os.path.join(model_dir, '{}.pth'.format(epoch)))

    # remove previous pretrained model if the number of models is too big
    pths = [
        int(pth.split('.')[0]) for pth in os.listdir(model_dir)
        if pth != 'latest.pth'
    ]
    if len(pths) <= 10:
        return
    os.system('rm {}'.format(
        os.path.join(model_dir, '{}.pth'.format(min(pths)))))

def load_network(net, model_dir, resume=True, epoch=-1, strict=True):
    if not resume:
        return 0

    if not os.path.exists(model_dir):
        print(colored('pretrained model does not exist: {}'.format(model_dir), 'red'))
        return 0

    if os.path.isdir(model_dir):
        pths = [
            int(pth.split('.')[0]) for pth in os.listdir(model_dir)
            if pth != 'latest.pth'
        ]
        if len(pths) == 0 and 'latest.pth' not in os.listdir(model_dir):
            return 0
        if epoch == -1:
            if 'latest.pth' in os.listdir(model_dir):
                pth = 'latest'
            else:
                pth = max(pths)
        else:
            pth = epoch
        model_path = os.path.join(model_dir, '{}.pth'.format(pth))
    else:
        model_path = model_dir

    print('load model: {}'.format(model_path))
    pretrained_model = torch.load(model_path)
    pretrain_net = pretrained_model['net']
    # pretrain_net = {k:pretrain_net[k] for k in pretrain_net if 'rendering_head' not in k}
    net.load_state_dict(pretrain_net, strict=False)
    if 'epoch' in pretrained_model:
        return pretrained_model['epoch'] + 1
    else:
        return 0


def remove_net_prefix(net, prefix):
    net_ = OrderedDict()
    for k in net.keys():
        if k.startswith(prefix):
            net_[k[len(prefix):]] = net[k]
        else:
            net_[k] = net[k]
    return net_


def add_net_prefix(net, prefix):
    net_ = OrderedDict()
    for k in net.keys():
        net_[prefix + k] = net[k]
    return net_


def replace_net_prefix(net, orig_prefix, prefix):
    net_ = OrderedDict()
    for k in net.keys():
        if k.startswith(orig_prefix):
            net_[prefix + k[len(orig_prefix):]] = net[k]
        else:
            net_[k] = net[k]
    return net_


def remove_net_layer(net, layers):
    keys = list(net.keys())
    for k in keys:
        for layer in layers:
            if k.startswith(layer):
                del net[k]
    return net

def save_trained_config(cfg, split='train'):
    train_config_path = os.path.join(cfg.result_dir, '{}_config_{}.yaml'.format(split, os.popen('git describe --tags --always').readline().strip()))
    open(train_config_path, 'w').write(cfg.dump())

def load_pretrain_head(net, model):
    model_path = os.path.join(cfg.workspace, 'pretrain', model)
    print('Load pretrain model: {}'.format(model_path))
    model = torch.load(model_path, 'cpu')
    if 'finetune' in model_path:
        feature_net = {'.'.join(k.split('.')[1:]):model['net'][k] for k in model['net'].keys() if 'feature_net' in k}
        net.feature_net.load_state_dict(feature_net)
        rendering_head = {'.'.join(k.split('.')[2:]):model['net'][k] for k in model['net'].keys() if 'rendering_head' in k}
        net.implicit_decoder.rendering_head.load_state_dict(rendering_head)
        return 0
    elif 'cnn_ibr' in model_path:
        feature_net = {'.'.join(k.split('.')[1:]):model['net'][k] for k in model['net'].keys() if 'feature_net' in k}
        net.feature_net.load_state_dict(feature_net)
        rendering_head = {'.'.join(k.split('.')[1:]):model['net'][k] for k in model['net'].keys() if 'nerf.agg' in k or 'nerf.color' in k}
        net.implicit_decoder.rendering_head.load_state_dict(rendering_head, strict=True)
        return 0
    feature_net = {'.'.join(k.split('.')[1:]):model['net'][k] for k in model['net'].keys() if 'feature_net' in k}
    net.feature_net.load_state_dict(feature_net)

    rendering_head = {'.'.join(k.split('.')[1:]):model['net'][k] for k in model['net'].keys() if 'nerf_1.agg' in k or 'nerf_1.color' in k}
    rendering_head.update({ '.'.join(k.split('.')[1:]).replace('lr1', 'lr0'):model['net'][k] for k in model['net'].keys() if 'nerf_1.lr1' in k})
    net.implicit_decoder.rendering_head.load_state_dict(rendering_head, strict=True)
    return 0


def load_pretrain(net, model_dir):
    model_dir = os.path.join(cfg.workspace, 'trained_model', cfg.task, cfg.scene, model_dir)
    if not os.path.exists(model_dir):
        return 1

    pths = [int(pth.split('.')[0]) for pth in os.listdir(model_dir) if pth != 'latest.pth']
    if len(pths) == 0 and 'latest.pth' not in os.listdir(model_dir):
        return 1

    if 'latest.pth' in os.listdir(model_dir):
        pth = 'latest'
    else:
        pth = max(pths)
  
    print('Load pretrain model: {}'.format(os.path.join(model_dir, '{}.pth'.format(pth))))
    pretrained_model = torch.load(os.path.join(model_dir, '{}.pth'.format(pth)), 'cpu')
    net.load_state_dict(pretrained_model['net'])
    return 0

def save_pretrain(net, task, model_dir):
    model_dir = os.path.join('data/trained_model', task, model_dir)
    os.system('mkdir -p ' +  model_dir)
    model = {'net': net.state_dict()}
    torch.save(model, os.path.join(model_dir, 'latest.pth'))



lib/utils/net_utils.py
zju3dv/im4d
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
    elif activation_type == 'none':
        return None
    else:
        raise ValueError(activation_type)

lib/utils/ops/activations.py
zju3dv/im4d
from bisect import bisect_right
from collections import Counter

import torch


class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):
    def __init__(
        self,
        optimizer,
        milestones,
        gamma=0.1,
        warmup_factor=1.0 / 3,
        warmup_iters=5,
        warmup_method="linear",
        last_epoch=-1,
    ):
        if not list(milestones) == sorted(milestones):
            raise ValueError(
                "Milestones should be a list of" " increasing integers. Got {}",
                milestones,
            )

        if warmup_method not in ("constant", "linear"):
            raise ValueError(
                "Only 'constant' or 'linear' warmup_method accepted"
                "got {}".format(warmup_method)
            )
        self.milestones = milestones
        self.gamma = gamma
        self.warmup_factor = warmup_factor
        self.warmup_iters = warmup_iters
        self.warmup_method = warmup_method
        super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        warmup_factor = 1
        if self.last_epoch < self.warmup_iters:
            if self.warmup_method == "constant":
                warmup_factor = self.warmup_factor
            elif self.warmup_method == "linear":
                alpha = float(self.last_epoch) / self.warmup_iters
                warmup_factor = self.warmup_factor * (1 - alpha) + alpha
        return [
            base_lr
            * warmup_factor
            * self.gamma ** bisect_right(self.milestones, self.last_epoch)
            for base_lr in self.base_lrs
        ]


class MultiStepLR(torch.optim.lr_scheduler._LRScheduler):

    def __init__(self, optimizer, milestones, gamma=0.1, last_epoch=-1):
        self.milestones = Counter(milestones)
        self.gamma = gamma
        super(MultiStepLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch not in self.milestones:
            return [group['lr'] for group in self.optimizer.param_groups]
        return [group['lr'] * self.gamma ** self.milestones[self.last_epoch]
                for group in self.optimizer.param_groups]


class ExponentialLR(torch.optim.lr_scheduler._LRScheduler):

    def __init__(self, optimizer, decay_epochs, gamma=0.1, last_epoch=-1):
        self.decay_epochs = decay_epochs
        self.gamma = gamma
        super(ExponentialLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        return [base_lr * self.gamma ** (self.last_epoch / self.decay_epochs)
                for base_lr in self.base_lrs]

lib/utils/optimizer/lr_scheduler.py
zju3dv/im4d
# copy from EasyVolumetricVideo
# Author: Zhen Xu https://github.com/dendenxu

from typing import Callable, List, Dict
from multiprocessing.pool import ThreadPool
from tqdm import tqdm
# from lib.utils.console_utils import *
from threading import Thread
import torch
from torch import Tensor
from typing import List




def cat_dict(dict_list, dim=0):
    return {k: torch.cat([item[k] for item in dict_list], dim) for k in dict_list[0] if isinstance(dict_list[0][k], Tensor)}

def cat_list(list_list: List[List[Tensor]], dim: int = 0):
    return [torch.cat([item[i] for item in list_list], dim=dim)  for i in range(len(list_list[0]))]

def cat_tensor(tensor_list: List[Tensor], dim=0):
    return torch.cat(tensor_list, dim=dim)

slice_func = lambda chunk_index, chunk_dim, chunk_size: [slice(None)] * chunk_dim + [slice(chunk_index, chunk_index+chunk_size)]
def chunkify(func, cat_func, chunk_tensors: List[Tensor], chunk_dim: int, chunk_size: int, **kwargs):
    '''
    func: function to be chunkified
    cat: function to concatenate the results
    chunk_tensors: list of tensors to be chunkified
    chunk_dim: dimension to be chunkified
    chunk_size: size of each chunk
    '''
    total_chunk_size = chunk_tensors[0].shape[chunk_dim]
    assert all([total_chunk_size == chunk_tensors[i].shape[chunk_dim] for i in range(1, len(chunk_tensors))])
    return cat_func([func(*[chunk_tensor[slice_func(i, chunk_dim, chunk_size)] for chunk_tensor in chunk_tensors], **kwargs) for i in range(0, total_chunk_size, chunk_size)], chunk_dim)
    
    all_ret = {}
    for i in range(0, chunk_size, chunk_num):
        ret = func(*[chunk_tensor[:, i:i + chunk_num] for chunk_tensor in chunk_tensors], **kwargs)
        if isinstance(ret, dict):
            for k in ret:
                if ret[k] is None:
                    continue
                if k not in all_ret:
                    all_ret[k] = []
                all_ret[k].append(ret[k])
        elif isinstance(ret, list) or isinstance(ret, tuple):
            for k in range(len(ret)):
                if ret[k] is None:
                    continue
                if k not in all_ret:
                    all_ret[k] = []
                all_ret[k].append(ret[k])
        elif isinstance(ret, torch.Tensor):
            if 0 not in all_ret:
                all_ret[0] = []
            all_ret[0].append(ret)
    if isinstance(ret, dict):
        return {k: torch.cat(all_ret[k], dim=chunk_dim) for k in all_ret}
    elif isinstance(ret, list) or isinstance(ret, tuple):
        return [torch.cat(all_ret[k], dim=chunk_dim) for k in all_ret]
    elif isinstance(ret, torch.Tensor):
        return torch.cat(all_ret[0], dim=chunk_dim)

def async_call(fn):
    def wrapper(*args, **kwargs):
        Thread(target=fn, args=args, kwargs=kwargs).start()
    return wrapper

def parallel_execution(*args, action: Callable, num_processes=32, print_progress=False, sequential=False, async_return=False, desc=None, **kwargs):
    # NOTE: we expect first arg / or kwargs to be distributed
    # NOTE: print_progress arg is reserved

    def get_length(args: List, kwargs: Dict):
        for a in args:
            if isinstance(a, list):
                return len(a)
        for v in kwargs.values():
            if isinstance(v, list):
                return len(v)
        raise NotImplementedError

    def get_action_args(length: int, args: List, kwargs: Dict, i: int):
        action_args = [(arg[i] if isinstance(arg, list) and len(arg) == length else arg) for arg in args]
        # TODO: Support all types of iterable
        action_kwargs = {key: (kwargs[key][i] if isinstance(kwargs[key], list) and len(kwargs[key]) == length else kwargs[key]) for key in kwargs}
        return action_args, action_kwargs

    if not sequential:
        # Create ThreadPool
        pool = ThreadPool(processes=num_processes)

        # Spawn threads
        results = []
        asyncs = []
        length = get_length(args, kwargs)
        for i in range(length):
            action_args, action_kwargs = get_action_args(length, args, kwargs, i)
            async_result = pool.apply_async(action, action_args, action_kwargs)
            asyncs.append(async_result)

        # Join threads and get return values
        if not async_return:
            for async_result in tqdm(asyncs, desc=desc, disable=not print_progress):
                results.append(async_result.get())  # will sync the corresponding thread
            pool.close()
            pool.join()
            return results
        else:
            return pool
    else:
        results = []
        length = get_length(args, kwargs)
        for i in tqdm(range(length), desc=desc, disable=not print_progress):
            action_args, action_kwargs = get_action_args(length, args, kwargs, i)
            async_result = action(*action_args, **action_kwargs)
            results.append(async_result)
        return results

lib/utils/parallel_utils.py
zju3dv/im4d
import numpy as np

class PixelSampler:
    @staticmethod
    def sample_patches(num_patch, patch_size, patch_ratio, H, W, msk_sample=None):
        assert(patch_size % 2 == 0)
        true_patch_size = patch_size
        patch_size = int(patch_size * patch_ratio)
        half_true_patch_size = true_patch_size // 2
        half_patch_size = patch_size // 2
        
        if msk_sample is not None and msk_sample.sum() > 0:
            num_fg_patch = num_patch
            non_zero = msk_sample.nonzero()
            permutation = np.random.permutation(msk_sample.sum())[:num_fg_patch].astype(np.int32)
            X_, Y_ = non_zero[1][permutation], non_zero[0][permutation]
            X_ = np.clip(X_, half_patch_size, W-half_patch_size)
            Y_ = np.clip(Y_, half_patch_size, H-half_patch_size)
        else:
            num_fg_patch = 0
            
        num_patch = num_patch - num_fg_patch
        X = np.random.randint(low=half_patch_size, high=W-half_patch_size, size=num_patch)
        Y = np.random.randint(low=half_patch_size, high=H-half_patch_size, size=num_patch)
        if num_fg_patch > 0:
            X = np.concatenate([X, X_]).astype(np.int32)
            Y = np.concatenate([Y, Y_]).astype(np.int32)
        grid = np.meshgrid((np.arange(true_patch_size)-half_true_patch_size)*patch_ratio, (np.arange(true_patch_size)-half_true_patch_size)*patch_ratio)
        return np.concatenate([grid[0].reshape(-1) + x for x in X]), np.concatenate([grid[1].reshape(-1) + y for y in Y])
lib/utils/pixel_sampling/pixel_sampler.py
zju3dv/im4d
import numpy as np
import cv2

def normalize(x):
    return x / np.linalg.norm(x)

def ptstocam(pts, c2w):
    tt = np.matmul(c2w[:3, :3].T, (pts-c2w[:3, 3])[..., np.newaxis])[..., 0]
    return tt

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

def gen_path(RT, center=None, num_views=100):
    lower_row = np.array([[0., 0., 0., 1.]])

    # transfer RT to camera_to_world matrix
    RT = np.array(RT)
    RT[:] = np.linalg.inv(RT[:])

    RT = np.concatenate([RT[:, :, 1:2], RT[:, :, 0:1],
                         -RT[:, :, 2:3], RT[:, :, 3:4]], 2)

    up = normalize(RT[:, :3, 0].sum(0))  # average up vector
    z = normalize(RT[0, :3, 2])
    vec1 = normalize(np.cross(z, up))
    vec2 = normalize(np.cross(up, vec1))
    z_off = 0

    if center is None:
        center = RT[:, :3, 3].mean(0)
        z_off = 1.3

    c2w = np.stack([up, vec1, vec2, center], 1)

    # get radii for spiral path
    tt = ptstocam(RT[:, :3, 3], c2w).T
    rads = np.percentile(np.abs(tt), 80, -1)
    rads = rads * 1.3
    rads = np.array(list(rads) + [1.])

    render_w2c = []
    for theta in np.linspace(0., 2 * np.pi, num_views + 1)[:-1]:
        # camera position
        cam_pos = np.array([0, np.sin(theta), np.cos(theta), 1] * rads)
        cam_pos_world = np.dot(c2w[:3, :4], cam_pos)
        # z axis
        z = normalize(cam_pos_world -
                      np.dot(c2w[:3, :4], np.array([z_off, 0, 0, 1.])))
        # vector -> 3x4 matrix (camera_to_world)
        mat = viewmatrix(z, up, cam_pos_world)

        mat = np.concatenate([mat[:, 1:2], mat[:, 0:1],
                              -mat[:, 2:3], mat[:, 3:4]], 1)
        mat = np.concatenate([mat, lower_row], 0)
        mat = np.linalg.inv(mat)
        render_w2c.append(mat)

    return render_w2c

def create_center_radius(center, radius=5., up='y', ranges=[0, 360, 36], angle_x=0, **kwargs):
    center = np.array(center).reshape(1, 3)
    thetas = np.deg2rad(np.linspace(*ranges))
    st = np.sin(thetas)
    ct = np.cos(thetas)
    zero = np.zeros_like(st)
    Rotx = cv2.Rodrigues(np.deg2rad(angle_x) * np.array([1., 0., 0.]))[0]
    if up == 'z':
        center = np.stack([radius*ct, radius*st, zero], axis=1) + center
        R = np.stack([-st, ct, zero, zero, zero, zero-1, -ct, -st, zero], axis=-1)
    elif up == 'y':
        center = np.stack([radius*ct, zero, radius*st, ], axis=1) + center
        R = np.stack([
            +st,  zero,  -ct,
            zero, zero-1, zero,
            -ct,  zero, -st], axis=-1)
    R = R.reshape(-1, 3, 3)
    if up == 'y':
        Rotz = cv2.Rodrigues(np.deg2rad(180) * np.array([0., 0., 1.]))[0]
        Rotx = Rotx @ Rotz
    R = np.einsum('ab,fbc->fac', Rotx, R)
    center = center.reshape(-1, 3, 1)
    T = - R @ center
    RT = np.dstack([R, T])
    return RT


def gen_path(RT, center=None, radius=3.2, up='z', ranges=[0, 360, 90], angle_x=27, **kwargs):
    ranges = [0, 360, kwargs['num_views']]
    c2ws = np.linalg.inv(RT[:])
    if center is None:
        center = c2ws[:, :3, 3].mean(0)
    center[0], center[1] = 0., 0.

    RTs = []
    center = np.array(center).reshape(1, 3)
    thetas = np.deg2rad(np.linspace(*ranges))
    st = np.sin(thetas)
    ct = np.cos(thetas)
    zero = np.zeros_like(st)
    Rotx = cv2.Rodrigues(np.deg2rad(angle_x) * np.array([1., 0., 0.]))[0]
    if up == 'z':
        center = np.stack([radius*ct, radius*st, zero], axis=1) + center
        R = np.stack([-st, ct, zero, zero, zero, zero-1, -ct, -st, zero], axis=-1)
    elif up == 'y':
        center = np.stack([radius*ct, zero, radius*st, ], axis=1) + center
        R = np.stack([
            +st,  zero,  -ct,
            zero, zero-1, zero,
            -ct,  zero, -st], axis=-1)
    R = R.reshape(-1, 3, 3)
    R = np.einsum('ab,fbc->fac', Rotx, R)
    center = center.reshape(-1, 3, 1)
    T = - R @ center
    RT = np.dstack([R, T])
    RT_bottom = np.zeros_like(RT[:, :1])
    RT_bottom[:, :, 3] = 1
    # __import__('ipdb').set_trace()
    ext = np.concatenate([RT, RT_bottom], axis=1)
    c2w = np.linalg.inv(ext)
    # __import__('ipdb').set_trace()
    # import matplotlib.pyplot as plt
    # plt.plot(c2ws[:, 0, 3], c2ws[:, 1, 3], '.')
    # plt.plot(c2w[:, 0, 3], c2w[:, 1, 3], '.')
    # plt.show()
    return ext

def gen_nerf_path(c2ws, depth_ranges, rads_scale=.5,  N_views=60):
    c2w = poses_avg(c2ws)
    up = normalize(c2ws[:, :3, 1].sum(0))

    close_depth, inf_depth = depth_ranges
    dt = .75
    mean_dz = 1./(( (1.-dt)/close_depth + dt/inf_depth ))
    focal = mean_dz

    shrink_factor = .8
    zdelta = close_depth * .2
    tt = c2ws[:, :3, 3] = c2w[:3, 3][None]
    rads = np.percentile(np.abs(tt), 70, 0)*rads_scale

    render_poses = render_path_spiral(c2w, up, rads, focal, zdelta, zrate=.5, N=N_views)
    return render_poses

def poses_avg(poses):
    center = poses[:, :3, 3].mean(0)
    vec2 = normalize(poses[:, :3, 2].sum(0))
    up = poses[:, :3, 1].sum(0)
    c2w = viewmatrix(vec2, up, center)
    return c2w

def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, N_rots=2, N=120):
    render_poses = []
    rads = np.array(list(rads) + [1.])

    for theta in np.linspace(0., 2. * np.pi * N_rots, N+1)[:-1]:
        c = np.dot(c2w[:3,:4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta*zrate), 1.]) * rads)
        z = normalize(c - np.dot(c2w[:3,:4], np.array([0,0,-focal, 1.])))
        render_poses.append(viewmatrix(z, up, c))
    return render_poses

lib/utils/rend_utils.py
zju3dv/im4d
import numpy as np
import cv2
from sklearn.decomposition import PCA

def vis_point_cloud(points, colors=None):
    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    if colors is not None: pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.visualization.draw_geometries([pcd])
    
def pcd_array2o3d(points, colors=None):
    import open3d as o3d
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    if colors is not None: pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd

def compute_pca_features(
    feature2D_map, feature_rescal_factor=8, save_path=None
):
    """
    feature2D_map: H*W*c
    feature3D: N*C
    """
    # PCA decrease dim of input feature:
    pca = PCA(n_components=3, svd_solver="arpack")

    H, W, C = feature2D_map.shape

    # feat_concated = np.concatenate(
        # [ feature2D_map.reshape(-1, C)], axis=0
    # )  # 2N * C
    pca_feature = pca.fit_transform(feature2D_map.reshape(-1, C))

    # Convert decreased feature to color map:
    pca_feature = cv2.normalize(
        pca_feature,
        None,
        alpha=0,
        beta=255,
        norm_type=cv2.NORM_MINMAX,
        dtype=cv2.CV_8UC3,
    )
    feature2D_map_pca = pca_feature.reshape(H, W, 3)
    return feature2D_map_pca

def get_bound_corners(bounds):
    min_x, min_y, min_z = bounds[0]
    max_x, max_y, max_z = bounds[1]
    corners_3d = np.array([
        [min_x, min_y, min_z],
        [min_x, min_y, max_z],
        [min_x, max_y, min_z],
        [min_x, max_y, max_z],
        [max_x, min_y, min_z],
        [max_x, min_y, max_z],
        [max_x, max_y, min_z],
        [max_x, max_y, max_z],
    ])
    return corners_3d

def plt_show_bbox(img, bounds, ext, ixt, path):
    import matplotlib.pyplot as plt
    plt.imshow(img)
    bbox_3d = get_bound_corners(bounds)
    bbox_2d = project(bbox_3d, ixt, ext)
    lines = np.array([[0, 0, 0, 7, 7, 7, 1, 1, 2, 2, 4, 4], [1, 2, 4, 5, 6, 3, 3, 5, 3, 6, 5, 6]]).T
    for id, line in enumerate(lines):
        pt1 = bbox_2d[line[0]]
        pt2 = bbox_2d[line[1]]
        plt.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]], 'm')
    plt.savefig(path)

def bbox_img(image, bounds, ext, ixt, linewidth=3, color=(0, 255, 0)):
    bbox_3d = get_bound_corners(bounds)
    bbox_2d = project(bbox_3d, ixt, ext)
    lines = np.array([[0, 0, 0, 7, 7, 7, 1, 1, 2, 2, 4, 4], [1, 2, 4, 5, 6, 3, 3, 5, 3, 6, 5, 6]]).T
    for id, line in enumerate(lines):
        pt1 = bbox_2d[line[0]].astype(int)
        pt2 = bbox_2d[line[1]].astype(int)
        image = cv2.line(image, tuple(pt1), tuple(pt2), color, linewidth)
    return image
            
def project(xyz, K, RT):
    """
    xyz: [N, 3]
    K: [3, 3]
    RT: [3, 4]
    """
    xyz = np.dot(xyz, RT[:, :3].T) + RT[:, 3:].T

    xyz = np.dot(xyz, K.T)
    xy = xyz[:, :2] / xyz[:, 2:]
    return xy

row_col_ = {
    2: (2, 1),
    7: (2, 4),
    8: (2, 4),
    9: (3, 3),
    26: (4, 7)
}

row_col_square = {
    2: (2, 1),
    7: (3, 3),
    8: (3, 3),
    9: (3, 3),
    26: (5, 5)
}

def get_row_col(l, square):
    if square and l in row_col_square.keys():
        return row_col_square[l]
    if l in row_col_.keys():
        return row_col_[l]
    else:
        from math import sqrt
        row = int(sqrt(l) + 0.5)
        col = int(l/ row + 0.5)
        if row*col<l:
            col = col + 1
        if row > col:
            row, col = col, row
        return row, col

def merge(images, row=-1, col=-1, resize=False, ret_range=False, square=False, **kwargs):
    if row == -1 and col == -1:
        row, col = get_row_col(len(images), square)
    height = images[0].shape[0]
    width = images[0].shape[1]
    # special case
    if height > width:
        if len(images) == 3:
            row, col = 1, 3
    if len(images[0].shape) > 2:
        ret_img = np.zeros((height * row, width * col, images[0].shape[2]), dtype=np.uint8) + 255
    else:
        ret_img = np.zeros((height * row, width * col), dtype=np.uint8) + 255
    ranges = []
    for i in range(row):
        for j in range(col):
            if i*col + j >= len(images):
                break
            img = images[i * col + j]
            # resize the image size
            img = cv2.resize(img, (width, height))
            ret_img[height * i: height * (i+1), width * j: width * (j+1)] = img
            ranges.append((width*j, height*i, width*(j+1), height*(i+1)))
    if resize:
        min_height = 1000
        if ret_img.shape[0] > min_height:
            scale = min_height/ret_img.shape[0]
            ret_img = cv2.resize(ret_img, None, fx=scale, fy=scale)
    if ret_range:
        return ret_img, ranges
    return ret_img

lib/utils/vis_utils.py
zju3dv/im4d
from lib.config import cfg, args
import numpy as np
import os
import cv2
from os.path import join
import torch
torch.backends.cudnn.benchmark = False

def run_dataset():
    from lib.datasets import make_data_loader
    import tqdm
    
    from lib.utils.data_utils import save_img

    # cfg.train.num_workers = 8
    data_loader = make_data_loader(cfg, is_train=True)
    for batch in tqdm.tqdm(data_loader):
        print(batch['near_far'])

def run_network():
    from lib.networks import make_network
    from lib.datasets import make_data_loader
    from lib.utils.net_utils import load_network
    from lib.utils.data_utils import to_cuda
    import tqdm
    import torch
    import time

    network = make_network(cfg).cuda()
    load_network(network, cfg.trained_model_dir, epoch=cfg.test.epoch)
    network.eval()

    data_loader = make_data_loader(cfg, is_train=False)
    net_time = []
    for batch in tqdm.tqdm(data_loader):
        batch = to_cuda(batch)
        with torch.no_grad():
            torch.cuda.synchronize()
            start = time.time()
            with torch.cuda.amp.autocast(enabled=cfg.eval_fp16):
                network(batch)
            torch.cuda.synchronize()
            net_time.append(time.time() - start)
    if len(net_time) > 1:
        print('net_time: ', np.mean(net_time[1:]))
        print('fps: ', 1./np.mean(net_time[1:]))
    else:
        print('net_time: ', np.mean(net_time))
        print('fps: ', 1./np.mean(net_time))

def run_evaluate():
    from lib.datasets import make_data_loader
    from lib.evaluators import make_evaluator
    import tqdm
    import torch
    from lib.networks import make_network
    from lib.utils import net_utils
    import time
    from lib.utils.net_utils import save_trained_config as save_config

    network = make_network(cfg).cuda()
    net_utils.load_network(network,
                           cfg.trained_model_dir,
                           resume=cfg.resume,
                           epoch=cfg.test.epoch)
    network.eval()
    save_config(cfg, 'test')

    data_loader = make_data_loader(cfg, is_train=False)
    evaluator = make_evaluator(cfg)
    net_time = []
    for batch in tqdm.tqdm(data_loader):
        for k in batch:
            if k != 'meta':
                batch[k] = batch[k].cuda()
        with torch.no_grad():
            torch.cuda.synchronize()
            start_time = time.time()
            with torch.cuda.amp.autocast(enabled=cfg.eval_fp16):
                output = network(batch)
            torch.cuda.synchronize()
            end_time = time.time()
        net_time.append(end_time - start_time)
        evaluator.evaluate(output, batch)
    evaluator.summarize()
    if len(net_time) > 1:
        print('net_time: ', np.mean(net_time[1:]))
        print('fps: ', 1./np.mean(net_time[1:]))
    else:
        print('net_time: ', np.mean(net_time))
        print('fps: ', 1./np.mean(net_time))

def run_cache_grid():
    from lib.datasets import make_data_loader
    from lib.evaluators import make_evaluator
    import tqdm
    import torch
    from lib.networks import make_network
    from lib.utils import net_utils
    import time
    from lib.utils.net_utils import save_trained_config as save_config

    network = make_network(cfg).cuda()
    net_utils.load_network(network,
                           cfg.trained_model_dir,
                           resume=cfg.resume,
                           epoch=cfg.test.epoch)
    network.eval()
    save_config(cfg, 'test')

    data_loader = make_data_loader(cfg, is_train=False)
    binarys = []
    bounds = []
    net_time = []
    for batch in tqdm.tqdm(data_loader):
        for k in batch:
            if k != 'meta':
                batch[k] = batch[k].cuda()
        with torch.no_grad():
            torch.cuda.synchronize()
            start_time = time.time()
            with torch.cuda.amp.autocast(enabled=cfg.eval_fp16):
                binary, bound = network.cache_grid(batch, cfg.save_mesh)
                binarys.append(binary)
                bounds.append(bound)
            torch.cuda.synchronize()
            end_time = time.time()
        net_time.append(end_time - start_time)
    np.savez_compressed(join(cfg.grid_dir, 'binarys.npz'), np.array(binarys))
    np.savez_compressed(join(cfg.grid_dir, 'bounds.npz'), np.array(bounds))
    if len(net_time) > 1:
        print('net_time: ', np.mean(net_time[1:]))
        print('fps: ', 1./np.mean(net_time[1:]))
    else:
        print('net_time: ', np.mean(net_time))
        print('fps: ', 1./np.mean(net_time))

def run_export_pcd():
    from lib.datasets import make_data_loader
    from lib.evaluators import make_evaluator
    import tqdm
    import torch
    from lib.networks import make_network
    from lib.utils import net_utils
    import time
    from lib.utils.net_utils import save_trained_config as save_config
    from lib.utils.im4d.im4d_utils import Im4DUtils

    network = make_network(cfg).cuda()
    net_utils.load_network(network,
                           cfg.trained_model_dir,
                           resume=cfg.resume,
                           epoch=cfg.test.epoch)
    network.eval()
    save_config(cfg, 'test')

    data_loader = make_data_loader(cfg, is_train=False)
    evaluator = make_evaluator(cfg)
    net_time = []
    frame_id = None
    pts, msks, hws, rgbs, mask_at_boxs = [], [], [], [], []
    for batch in tqdm.tqdm(data_loader):
        if frame_id is not None and batch['meta']['frame_id'].item() != frame_id:
            export_path = join(cfg.result_dir, 'meshes', '{:06d}.ply'.format(frame_id))
            os.makedirs(os.path.dirname(export_path), exist_ok=True)
            Im4DUtils.extract_mesh(pts, msks, hws, rgbs, mask_at_boxs, export_path)
            pts, msks, hws, rgbs, mask_at_boxs = [], [], [], [], []
        frame_id = batch['meta']['frame_id'].item()
        for k in batch:
            if k != 'meta':
                batch[k] = batch[k].cuda()
        with torch.no_grad():
            torch.cuda.synchronize()
            start_time = time.time()
            with torch.cuda.amp.autocast(enabled=cfg.eval_fp16):
                output = network(batch)
            pts.append(output['pts_0'][0])
            msks.append(batch['msk'][0])
            rgbs.append(batch['rgb'][0])
            hws.append((batch['meta']['H'].item(), batch['meta']['W'].item()))
            mask_at_boxs.append(batch['mask_at_box'][0])
            torch.cuda.synchronize()
            end_time = time.time()
        net_time.append(end_time - start_time)
        evaluator.evaluate(output, batch)
    export_path = join(cfg.result_dir, 'meshes', '{:06d}.ply'.format(frame_id))
    Im4DUtils.extract_mesh(pts, msks, hws, rgbs, mask_at_boxs, export_path)
    evaluator.summarize()
    if len(net_time) > 1:
        print('net_time: ', np.mean(net_time[1:]))
        print('fps: ', 1./np.mean(net_time[1:]))
    else:
        print('net_time: ', np.mean(net_time))
        print('fps: ', 1./np.mean(net_time))

if __name__ == '__main__':
    globals()['run_' + args.type]()

run.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

import azure.cognitiveservices.speech as speechsdk
speech_config = speechsdk.SpeechConfig(subscription=os.environ.get('SPEECH_KEY'), region=os.environ.get('SPEECH_REGION'))
speech_config.set_speech_synthesis_output_format(speechsdk.SpeechSynthesisOutputFormat.Audio16Khz32KBitRateMonoMp3) 
audio_config = speechsdk.audio.AudioOutputConfig(use_default_speaker=True)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, help='input text file')
    parser.add_argument('--output', type=str, default=None, help='output audio file')
    parser.add_argument('--voice_type', type=str, default='en-US-AriaNeural', help='output audio file')
    args = parser.parse_args()
    return args

def print_output(speech_synthesis_result, text='Hello, world!'):
    if speech_synthesis_result.reason == speechsdk.ResultReason.SynthesizingAudioCompleted:
        print("Speech synthesized for text [{}]".format(text))
    elif speech_synthesis_result.reason == speechsdk.ResultReason.Canceled:
        cancellation_details = speech_synthesis_result.cancellation_details
        print("Speech synthesis canceled: {}".format(cancellation_details.reason))
        if cancellation_details.reason == speechsdk.CancellationReason.Error:
            if cancellation_details.error_details:
                print("Error details: {}".format(cancellation_details.error_details))
                print("Did you set the speech resource key and region values?")
                
def parse_lines(input_path):
    with open(input_path, 'r') as f:
        lines = f.readlines()
    texts = []
    text = ''
    for line in lines:
        if line == '\n':
            texts.append(text)
            text = ''
        else:
            text += line
    if text != '':
        texts.append(text)
    return texts

def gen_audio(text, filename):
    audio_config = speechsdk.audio.AudioOutputConfig(use_default_speaker=True, filename=filename)
    speech_synthesizer = speechsdk.SpeechSynthesizer(speech_config=speech_config, audio_config=audio_config)
    speech_synthesis_result = speech_synthesizer.speak_text_async(text).get()
    print_output(speech_synthesis_result)

def main(args):
    speech_config.speech_synthesis_voice_name = args.voice_type
    if args.input[-4:] != '.txt': # Debug only
        gen_audio(args.input, args.output)
        return
    texts = parse_lines(args.input)
    os.makedirs(args.output, exist_ok=True)
    for idx, text in enumerate(texts):
        filename = join(args.output, '{:03d}_{}.mp3'.format(idx, text[:50]))
        gen_audio(text, filename)
if __name__ == '__main__':
    args = parse_args()
    main(args)
### usage ###
# 1. python scripts/azure/tts.py --input "Hello, world!" --output data/outputaudio.mp3
# This will generate an audio file named outputaudio.mp3 in the data folder.
# 2. python scripts/azure/tts.py --input data/neural_scene_chronology.txt --output data/neural_scene_chronology
# This will generate a folder named neural_scene_chronology in the data folder, and each audio file is named as 000_The scene opens on a dark night.mp3, 001_The scene opens on a dark night.mp3, etc.
# One empty line in the input text file means split.

scripts/azure/tts.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from trdparty.colmap.read_write_model import read_model, read_points3d_binary
import trimesh

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--thresh', type=int, default=3)
    parser.add_argument('--use_roimask', type=bool, default=True)
    args = parser.parse_args()
    return args

def main(args):
    input_dir = args.input
    output_path = args.output
    # cameras, images, points3D = read_model(input_dir)
    points3D = read_points3d_binary(join(input_dir, 'points3D.bin'))
    points = np.array([points3D[k].xyz for k in points3D if len(points3D[k].image_ids) > args.thresh ])
    rgbs = np.array([points3D[k].rgb for k in points3D if len(points3D[k].image_ids) > args.thresh])
    if args.use_roimask:
        roi = np.stack([np.percentile(points, 0.1, axis=0), np.percentile(points, 99.9, axis=0)])
        mask = np.logical_and(points > roi[:1], points < roi[1:])
        mask = mask.sum(axis=-1) == 3
        points = points[mask]
        rgbs = rgbs[mask]
    print('Number of points: ', len(points))
    pcd = trimesh.PointCloud(vertices=points, colors=rgbs)
    pcd.export(output_path)

if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/colmap/colmap2ply.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from trdparty.colmap.read_write_model import read_model, qvec2rotmat

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    args = parser.parse_args()
    return args

def main(args,
         num_thresh = 3, # 每个点最少应该被3+1张图片看到
         ):
    input_path = args.input
    cameras, images, points3D = read_model(join(input_path, 'dense', 'sparse'))
    cam_dict = {}
    
    for im_k in tqdm(images):
        
        ext = np.eye(4)
        ext[:3, :3] = qvec2rotmat(images[im_k].qvec)
        ext[:3, 3:] = images[im_k].tvec.reshape(3, 1)
        
        cam_params = cameras[images[im_k].camera_id].params
        assert(len(cam_params) == 4) # only support PINHOLE camera
        ixt = np.eye(3)
        ixt[0, 0], ixt[1, 1] = cam_params[0], cam_params[1]
        ixt[0, 2], ixt[1, 2] = cam_params[2], cam_params[3]
        
        h, w = cameras[images[im_k].camera_id].height, cameras[images[im_k].camera_id].width
        
        points = [points3D[k].xyz for k in images[im_k].point3D_ids if k != -1 and points3D[k] and len(points3D[k].image_ids) > num_thresh]
        points = np.array(points)
        
        if len(points) <= 50:
            continue
        points = points @ ext[:3, :3].T + ext[:3, 3:].T
        near_far = np.array([points[:, 2].min(), points[:, 2].max()])
        
        print(im_k, near_far)
    
        cam = {
            'img_path': images[im_k].name,
            'ext': ext,
            'ixt': ixt, 
            'near_far': near_far,
            'h': h,
            'w': w,
        }
        cam_dict[im_k] = cam
    np.save(join(input_path, 'annots/cam_dict.npy'), cam_dict)

if __name__ == '__main__':
    args = parse_args()
    main(args)
    


scripts/colmap/colmap2sc.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from trdparty.colmap.read_write_model import read_model, read_points3d_binary
import trimesh
import torch

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_sparse', type=str)
    parser.add_argument('--input_ply', type=str)
    parser.add_argument('--output_path', type=str)
    parser.add_argument('--thresh', type=float, default=3)
    args = parser.parse_args()
    return args

def main(args):
    '''
    generate the list of interested images from a COLMAP sparse model and ply.
    specifically, select images that rely on the point cloud of interest to complete the registration.
    '''
    cameras, images, points3D = read_model(args.input_sparse)
    roi_ply = trimesh.load_mesh(args.input_ply)
    roi_points = np.array(roi_ply.vertices)
    points = np.array([points3D[k].xyz for k in points3D])
    points_id = np.array([points3D[k].id for k in points3D])
    
    roi_points_id = []
    chunk = 128
    roi_points_cuda = torch.tensor(roi_points).cuda().float()
    points_cuda = torch.tensor(points).cuda().float()
    for i in tqdm(range(0, roi_points.shape[0], chunk)):
        chunk_points = roi_points_cuda[i:i+chunk]
        distance = (chunk_points[:, None] - points_cuda[None]).norm(dim=-1)
        roi_points_id += points_id[distance.argmin(dim=-1).cpu().numpy()].tolist()
    roi_points_id = np.array(roi_points_id)
    
    save_images = []
    save_conds = []
    for im_id, image in tqdm(images.items()):
        im_points_ids = image.point3D_ids[image.point3D_ids != -1]
        save_conds.append(np.isin(im_points_ids, roi_points_id).mean())
        save_images.append(image.name)
    threshs = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] 
    for thresh in threshs:
        save_list = np.array(save_images)[(np.array(save_conds) > thresh)].tolist()
        print('thresh: {}, number of images: {}'.format(thresh, len(save_list)))
        open(join(args.output_path, '{}_{}.txt'.format(os.path.basename(args.input_ply).split('.')[0], thresh)), 'w').write('\n'.join(save_list))

if __name__ == '__main__':
    args = parse_args()
    main(args)





# def main(input_sparse, input_ply, output_path, thresh):
#     '''
#     Generate the list of interested images from a COLMAP sparse model and ply.
#     Specifically, select images that rely on the point cloud of interest to complete the registration.
#     '''

#     cameras, images, points3D = read_model(input_sparse)
#     roi_ply = trimesh.load_mesh(input_ply)
#     roi_points = np.array(roi_ply.vertices)
#     points = np.array([points3D[k].xyz for k in points3D])
#     points_id = np.array([points3D[k].id for k in points3D])

#     # Generate points_id for the interested point cloud
#     roi_points_id = []
#     chunk = 128 # Reduce it if CUDA out of memory occurs
#     roi_points_cuda = torch.tensor(roi_points).cuda().float()
#     points_cuda = torch.tensor(points).cuda().float()
#     for i in tqdm.tqdm(range(0, roi_points.shape[0], chunk)):
#         chunk_points = roi_points_cuda[i:i+chunk]
#         distance = (chunk_points[:, None] - points_cuda[None]).norm(dim=-1)
#         roi_points_id += points_id[distance.argmin(dim=-1).cpu().numpy()].tolist()
#     roi_points_id = np.array(roi_points_id)

#     # Select interested images
#     save_images = []
#     save_conds = []
#     for im_id, image in tqdm.tqdm(images.items()):
#         im_points_ids = image.point3D_ids[image.point3D_ids!=-1]
#         save_conds.append(np.isin(im_points_ids, roi_points_id).mean())
#         save_images.append(image.name)
#     save_list = np.array(save_images)[(np.array(save_conds) > thresh)].tolist()
#     print('Interested: {}'.format(len(save_list)))
#     open(output_path, 'w').writelines([item + '\n' for item in save_list])
scripts/colmap/genlist_by_ply.py
zju3dv/im4d
'''
On the NHR dataset, mlp_maps found that the mask provided was relatively poor, so the edge part of the mask was excluded when testing PSNR.
'''

import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
import imageio.v2 as imageio
import matplotlib.pyplot as plt 
from skimage.metrics import peak_signal_noise_ratio as psnr_func
sys.path.append('.')

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--gt_path', type=str)
    parser.add_argument('--output_path', type=str)
    args = parser.parse_args()
    return args


def eval_img(pred_path, gt_path, msk_path, edge_path):
    pred = (imageio.imread(pred_path)[..., :3] / 255.).astype(np.float32)
    gt = (imageio.imread(gt_path) / 255.).astype(np.float32)
    msk = (imageio.imread(msk_path) / 255.)
    edge = (imageio.imread(edge_path) / 255.)
    
    eval_msk = np.logical_and(msk > 0.5, edge < 0.5)
    return psnr_func(pred[eval_msk], gt[eval_msk])
    # plt.imshow(eval_msk)
    # plt.savefig('test.jpg')
    
def main(args):
    psnrs = []
    for img_name in os.listdir(args.output_path):
        frame_id, view_id = int(img_name[5:9]), int(img_name[14:18])
        pred_path = join(args.output_path, img_name)
        gt_path = join(args.gt_path, '{:04d}_{:02d}_gt.png'.format(frame_id, view_id))
        msk_path = join(args.gt_path, '{:04d}_{:02d}_mask.png'.format(frame_id, view_id))
        edge_path = join(args.gt_path, '{:04d}_{:02d}_edge.png'.format(frame_id, view_id))
        psnrs.append(eval_img(pred_path, gt_path, msk_path, edge_path))
    print('Mean psnr: ', np.mean(psnrs))

if __name__ == '__main__':
    args = parse_args()
    main(args)
scripts/evaluate/im4d/eval_nhr.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
import open3d as o3d
import trimesh
sys.path.append('.')
from lib.utils import data_utils
import torch
import torch.nn.functional as F

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_grid_dir', type=str, default='/home/linhaotong/workspace/im4d/data/htcode/result/0013_01/im4d/grid/sample')
    parser.add_argument('--input_mesh_dir', type=str, default='/home/linhaotong/workspace/im4d/data/htcode/result/0013_01/im4d/0013_01_demo/pcd_1003/meshes')
    parser.add_argument('--output_grid_dir', type=str, default='/home/linhaotong/workspace/im4d/data/htcode/result/0013_01/im4d/grid/pcd')
    parser.add_argument('--padding_grid', type=bool, default=False)
    parser.add_argument('--conv_size', type=int, default=3)
    args = parser.parse_args()
    return args

def main(args):
    grid_size = (128, 256, 128)
    padding_size = 0.1
    bounds = np.load(join(args.input_grid_dir, 'bounds.npz'))['arr_0']
    binarys = np.load(join(args.input_grid_dir, 'binarys.npz'))['arr_0']
    
    new_bounds = []
    new_binarys = []
    for frame_id in tqdm(range(len(bounds))):
        bound = bounds[frame_id]
        if not os.path.exists(join(args.input_mesh_dir, '{:06d}.ply'.format(frame_id))) and frame_id not in [0, 20, 40, 60, 80, 100, 120, 140]:
            new_bounds.append(bounds[frame_id])
            new_binarys.append(binarys[frame_id])
            continue
        pcd_denoise = data_utils.read_and_denoise(join(args.input_mesh_dir, '{:06d}.ply'.format(frame_id)), vox_size=0.001)
        points = np.asarray(pcd_denoise.points)
        
        # cropping points out of bound
        points = points[(points[:, 0] > bound[0, 0]) & (points[:, 0] < bound[1, 0])]
        points = points[(points[:, 1] > bound[0, 1]) & (points[:, 1] < bound[1, 1])]
        points = points[(points[:, 2] > bound[0, 2]) & (points[:, 2] < bound[1, 2])]
        
        mesh = trimesh.PointCloud(points)
        mesh_path = join(args.output_grid_dir, 'meshes', '{:06d}.ply'.format(frame_id))
        os.makedirs(os.path.dirname(mesh_path), exist_ok=True)
        print(mesh_path)
        mesh.export(mesh_path)
        
        # min_bound = points.min(axis=0) - padding_size
        # max_bound = points.max(axis=0) + padding_size
        min_bound = bound[0]
        max_bound = bound[1]
        voxel_size = (max_bound - min_bound) / np.array(grid_size)
        indices = ((points - min_bound) // voxel_size).astype(int)
        voxel_grid = np.zeros(grid_size, dtype=np.uint8)
        voxel_grid[indices[:, 0], indices[:, 1], indices[:, 2]] = 1
        
        if args.padding_grid:
            conv_size = args.conv_size
            binary = torch.tensor(voxel_grid).cuda()
            structuring_element = torch.ones(1, 1, conv_size, conv_size, conv_size, device=binary.device)
            dilated_tensor = F.conv3d(binary[None, None].float(), structuring_element.float(), padding=conv_size//2) > 0
            voxel_grid = (dilated_tensor.cpu().numpy()).astype(np.uint8)[0, 0]
       
        print(frame_id, voxel_size, binarys[frame_id].sum(), voxel_grid.sum())
        
        new_bounds.append(np.stack([min_bound, max_bound]).astype(np.float32))
        new_binarys.append(voxel_grid.astype(np.bool))
    os.makedirs(args.output_grid_dir, exist_ok=True)
    np.savez_compressed(join(args.output_grid_dir, 'bounds.npz'), np.array(new_bounds))
    np.savez_compressed(join(args.output_grid_dir, 'binarys.npz'), np.array(new_binarys))
        

if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/im4d/voxelize_pcd.py
zju3dv/im4d
# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#
#     * 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.
#
#     * Neither the name of ETH Zurich and UNC Chapel Hill 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 HOLDERS 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.
#
# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)

import os
import sys
import collections
import numpy as np
import struct
import cv2
import sys
sys.path.append('.')

from trdparty.EasyMocapPublic.easymocap.mytools.camera_utils import write_camera

CameraModel = collections.namedtuple(
    "CameraModel", ["model_id", "model_name", "num_params"])
Camera = collections.namedtuple(
    "Camera", ["id", "model", "width", "height", "params"])
BaseImage = collections.namedtuple(
    "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
Point3D = collections.namedtuple(
    "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])

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


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])
CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model)
                           for camera_model in CAMERA_MODELS])

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
            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 read_points3D_text(path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadPoints3DText(const std::string& path)
        void Reconstruction::WritePoints3DText(const std::string& path)
    """
    points3D = {}
    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()
                point3D_id = int(elems[0])
                xyz = np.array(tuple(map(float, elems[1:4])))
                rgb = np.array(tuple(map(int, elems[4:7])))
                error = float(elems[7])
                image_ids = np.array(tuple(map(int, elems[8::2])))
                point2D_idxs = np.array(tuple(map(int, elems[9::2])))
                points3D[point3D_id] = Point3D(id=point3D_id, xyz=xyz, rgb=rgb,
                                               error=error, image_ids=image_ids,
                                               point2D_idxs=point2D_idxs)
    return points3D


def read_points3d_binary(path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadPoints3DBinary(const std::string& path)
        void Reconstruction::WritePoints3DBinary(const std::string& path)
    """
    points3D = {}
    with open(path_to_model_file, "rb") as fid:
        num_points = read_next_bytes(fid, 8, "Q")[0]
        for point_line_index in range(num_points):
            binary_point_line_properties = read_next_bytes(
                fid, num_bytes=43, format_char_sequence="QdddBBBd")
            point3D_id = binary_point_line_properties[0]
            xyz = np.array(binary_point_line_properties[1:4])
            rgb = np.array(binary_point_line_properties[4:7])
            error = np.array(binary_point_line_properties[7])
            track_length = read_next_bytes(
                fid, num_bytes=8, format_char_sequence="Q")[0]
            track_elems = read_next_bytes(
                fid, num_bytes=8*track_length,
                format_char_sequence="ii"*track_length)
            image_ids = np.array(tuple(map(int, track_elems[0::2])))
            point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
            points3D[point3D_id] = Point3D(
                id=point3D_id, xyz=xyz, rgb=rgb,
                error=error, image_ids=image_ids,
                point2D_idxs=point2D_idxs)
    return points3D


def read_model(path, ext):
    if ext == ".txt":
        cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
        images = read_images_text(os.path.join(path, "images" + ext))
        points3D = read_points3D_text(os.path.join(path, "points3D") + ext)
    else:
        cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
        images = read_images_binary(os.path.join(path, "images" + ext))
        points3D = read_points3d_binary(os.path.join(path, "points3D") + ext)
    return cameras, images, points3D


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 qvec


def write_cameras_text(cameras, path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasText(const std::string& path)
        void Reconstruction::ReadCamerasText(const std::string& path)
    """
    HEADER = '# Camera list with one line of data per camera:\n'
    '#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n'
    '# Number of cameras: {}\n'.format(len(cameras))
    with open(path, "w") as fid:
        fid.write(HEADER)
        for _, cam in cameras.items():
            to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
            line = " ".join([str(elem) for elem in to_write])
            fid.write(line + "\n")

def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
    """pack and write to a binary file.
    :param fid:
    :param data: data to send, if multiple elements are sent at the same time,
    they should be encapsuled either in a list or a tuple
    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
    should be the same length as the data list or tuple
    :param endian_character: Any of {@, =, <, >, !}
    """
    if isinstance(data, (list, tuple)):
        bytes = struct.pack(endian_character + format_char_sequence, *data)
    else:
        bytes = struct.pack(endian_character + format_char_sequence, data)
    fid.write(bytes)

def write_cameras_binary(cameras, path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::WriteCamerasBinary(const std::string& path)
        void Reconstruction::ReadCamerasBinary(const std::string& path)
    """
    with open(path_to_model_file, "wb") as fid:
        write_next_bytes(fid, len(cameras), "Q")
        for _, cam in cameras.items():
            model_id = CAMERA_MODEL_NAMES[cam.model].model_id
            camera_properties = [cam.id,
                                 model_id,
                                 cam.width,
                                 cam.height]
            write_next_bytes(fid, camera_properties, "iiQQ")
            for p in cam.params:
                write_next_bytes(fid, float(p), "d")
    return cameras


def write_images_binary(images, path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesBinary(const std::string& path)
        void Reconstruction::WriteImagesBinary(const std::string& path)
    """
    with open(path_to_model_file, "wb") as fid:
        write_next_bytes(fid, len(images), "Q")
        for _, img in images.items():
            write_next_bytes(fid, img.id, "i")
            write_next_bytes(fid, img.qvec.tolist(), "dddd")
            write_next_bytes(fid, img.tvec.tolist(), "ddd")
            write_next_bytes(fid, img.camera_id, "i")
            for char in img.name:
                write_next_bytes(fid, char.encode("utf-8"), "c")
            write_next_bytes(fid, b"\x00", "c")
            write_next_bytes(fid, len(img.point3D_ids), "Q")
            for xy, p3d_id in zip(img.xys, img.point3D_ids):
                write_next_bytes(fid, [*xy, p3d_id], "ddq")

def write_images_text(images, path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesText(const std::string& path)
        void Reconstruction::WriteImagesText(const std::string& path)
    """
    if len(images) == 0:
        mean_observations = 0
    else:
        mean_observations = sum((len(img.point3D_ids) for _, img in images.items()))/len(images)
    HEADER = '# Image list with two lines of data per image:\n'
    '#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n'
    '#   POINTS2D[] as (X, Y, POINT3D_ID)\n'
    '# Number of images: {}, mean observations per image: {}\n'.format(len(images), mean_observations)

    with open(path, "w") as fid:
        fid.write(HEADER)
        for _, img in images.items():
            image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
            first_line = " ".join(map(str, image_header))
            fid.write(first_line + "\n")

            points_strings = []
            for xy, point3D_id in zip(img.xys, img.point3D_ids):
                points_strings.append(" ".join(map(str, [*xy, point3D_id])))
            fid.write(" ".join(points_strings) + "\n")

def write_points3D_text(points3D, path):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadPoints3DText(const std::string& path)
        void Reconstruction::WritePoints3DText(const std::string& path)
    """
    if len(points3D) == 0:
        mean_track_length = 0
    else:
        mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items()))/len(points3D)
    HEADER = '# 3D point list with one line of data per point:\n'
    '#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n'
    '# Number of points: {}, mean track length: {}\n'.format(len(points3D), mean_track_length)

    with open(path, "w") as fid:
        fid.write(HEADER)
        for _, pt in points3D.items():
            point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
            fid.write(" ".join(map(str, point_header)) + " ")
            track_strings = []
            for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
                track_strings.append(" ".join(map(str, [image_id, point2D])))
            fid.write(" ".join(track_strings) + "\n")


def write_points3d_binary(points3D, path_to_model_file):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadPoints3DBinary(const std::string& path)
        void Reconstruction::WritePoints3DBinary(const std::string& path)
    """
    with open(path_to_model_file, "wb") as fid:
        write_next_bytes(fid, len(points3D), "Q")
        for _, pt in points3D.items():
            write_next_bytes(fid, pt.id, "Q")
            write_next_bytes(fid, pt.xyz.tolist(), "ddd")
            write_next_bytes(fid, pt.rgb.tolist(), "BBB")
            write_next_bytes(fid, pt.error, "d")
            track_length = pt.image_ids.shape[0]
            write_next_bytes(fid, track_length, "Q")
            for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
                write_next_bytes(fid, [image_id, point2D_id], "ii")


class FileStorage(object):
    def __init__(self, filename, isWrite=False):
        version = cv2.__version__
        self.version = int(version.split('.')[0])
        if isWrite:
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_WRITE)
        else:
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_READ)

    def __del__(self):
        cv2.FileStorage.release(self.fs)

    def write(self, key, value, dt='mat'):
        if dt == 'mat':
            cv2.FileStorage.write(self.fs, key, value)
        elif dt == 'list':
            if self.version == 4: # 4.4
                # self.fs.write(key, '[')
                # for elem in value:
                #     self.fs.write('none', elem)
                # self.fs.write('none', ']')
                # import ipdb; ipdb.set_trace()
                self.fs.startWriteStruct(key, cv2.FileNode_SEQ)
                for elem in value:
                    self.fs.write('', elem)
                self.fs.endWriteStruct()
            else: # 3.4
                self.fs.write(key, '[')
                for elem in value:
                    self.fs.write('none', elem)
                self.fs.write('none', ']')

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str, default='/home/')
    parser.add_argument('--ext', type=str, default='.bin')
    parser.add_argument('--min_views', type=int, default=4)
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()

    cameras, images, points3D = read_model(path=args.path, ext=args.ext)
    points3D_new = {}
    for key, val in points3D.items():
        if val.image_ids.shape[0] < args.min_views:
            continue
        points3D_new[key] = val
    print('[Read Colmap] filter {}/{} points3D, min view = {}'.format(len(points3D_new), len(points3D), args.min_views))
    points3D = points3D_new
    import cv2
    cameras_out = {}
    for key in cameras.keys():
        p = cameras[key].params
        if cameras[key].model == 'SIMPLE_RADIAL':
            f, cx, cy, k = p
            K = np.array([f, 0, cx, 0, f, cy, 0, 0, 1]).reshape(3, 3)
            dist = np.array([[k, 0, 0, 0, 0]])
        elif cameras[key].model == 'PINHOLE':
            fx, fy, cx, cy = p
            K = np.array([fx, 0, cx, 0, fy, cy, 0, 0, 1]).reshape(3, 3)
            dist = np.array([[0, 0, 0, 0, 0]])
        else:
            K = np.array([[p[0], 0, p[2], 0, p[1], p[3], 0, 0, 1]]).reshape(3, 3)
            dist = np.array([[p[4], p[5], p[6], p[7], 0.]])
        cameras_out[key] = {'K': K, 'dist': dist, 'H': cameras[key].height, 'W': cameras[key].width}
    mapkey = {}
    cameras_new = {}
    for key, val in images.items():
        cam = cameras_out[val.camera_id].copy()
        t = val.tvec.reshape(3, 1)
        R = qvec2rotmat(val.qvec)
        cam['R'] = R
        cam['Rvec'] = cv2.Rodrigues(R)[0]
        cam['T'] = t
        # mapkey[val.name.split('.')[0]] = val.camera_id
        
        cameras_new[val.name.split('.')[0]] = cam
        # cameras_new[val.name.split('.')[0].split('/')[0]] = cam
    keys = sorted(list(cameras_new.keys()))
    cameras_new = {key:cameras_new[key] for key in keys}
    print("num_cameras: {}/{}".format(len(cameras), len(cameras_new)))
    print("num_images:", len(images))
    print("num_points3D:", len(points3D))
    if len(points3D) > 0:
        keys = list(points3D.keys())
        xyz = np.stack([points3D[k].xyz for k in keys])
        rgb = np.stack([points3D[k].rgb for k in keys])
        import open3d as o3d
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(xyz)
        pcd.colors = o3d.utility.Vector3dVector(rgb/255.)
        from os.path import join
        pcdname = join(sys.argv[1], 'sparse.ply')
        o3d.io.write_point_cloud(pcdname, pcd)
        if False:
            o3d.visualization.draw_geometries([pcd])
    write_camera(cameras_new, sys.argv[1])

if __name__ == "__main__":
    main()
scripts/nerfstudio/colmap2easyvv.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

# from lib.utils.data_utils import read_camera
from trdparty.colmap.read_write_model import read_model, qvec2rotmat
import json

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--input_list', type=str, default=None)
    parser.add_argument('--output', type=str)
    args = parser.parse_args()
    return args

def main(args):
    cameras, images, point3Ds = read_model(args.input)
    if args.input_list is not None:
        interested = open(args.input_list).read().splitlines()
    else:
        interested = []
    
    frames = []
    
    for k, v in tqdm(images.items()):
        if len(interested) != 0 and v.name not in interested:
            continue
        file_path = join('images', v.name)
        colmap_im_id = k
        
        w2c = np.eye(4)
        w2c[:3, :3] = v.qvec2rotmat()
        w2c[:3, 3] = v.tvec
        
        c2w = np.linalg.inv(w2c)
        c2w[0:3, 1:3] *= -1
        c2w = c2w[np.array([1, 0, 2, 3]), :]
        c2w[2, :] *= -1
        transform_matrix = c2w
        
        cam_id = v.camera_id
        
        frame = {
            'file_path': file_path,
            'transform_matrix': transform_matrix.tolist(),
            'colmap_im_id': colmap_im_id,
            'camera_model': 'OPENCV',
            'fl_x': cameras[v.camera_id].params[0],
            'fl_y': cameras[v.camera_id].params[1],
            'cx': cameras[v.camera_id].params[2],
            'cy': cameras[v.camera_id].params[3],
            'w': cameras[cam_id].width,
            'h': cameras[cam_id].height,
            'k1': cameras[cam_id].params[4],
            'k2': cameras[cam_id].params[5],
            'p1': cameras[cam_id].params[6],
            'p2': cameras[cam_id].params[7],
            'k3': 0,
        }
        frames.append(frame)
    out = {}
    out['frames'] = frames
    
    applied_transform = np.eye(4)[:3, :]
    applied_transform = applied_transform[np.array([1, 0, 2]), :]
    applied_transform[2, :] *= -1
    out["applied_transform"] = applied_transform.tolist()
    with open(join(f'{args.output}', 'transforms.json'), 'w', encoding='utf-8') as f:
        json.dump(out, f, indent=4)
        
if __name__ == '__main__':
    args = parse_args()
    main(args)

scripts/nerfstudio/colmap2nerfstudio.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from lib.utils.data_utils import read_camera
from lib.utils.easyvv.easy_utils import read_camera
import json
import imageio

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--frame_id', type=int)
    parser.add_argument('--use_msk', type=bool, default=False)
    parser.add_argument('--write_alpha', type=bool, default=False)
    parser.add_argument('--white_bkgd', type=bool, default=False)
    args = parser.parse_args()
    return args

def main(args):
    cams = os.listdir(join(args.input, 'images'))
    os.system(f'mkdir -p {args.output}/images')
    for cam in tqdm(cams):
        frame_path = join(args.input, 'images', cam, f'{args.frame_id:06d}.jpg')
        if not args.use_msk:
            tar_path = join(args.output, 'images', f'{cam}.jpg')
            os.system(f'cp {frame_path} {tar_path}')
        else:
            tar_path = join(args.output, 'images', f'{cam}.png') if args.write_alpha else join(args.output, 'images', f'{cam}.jpg')
            img = imageio.imread(frame_path)
            msk = np.array(imageio.imread(join(args.input, 'maskes', cam, f'{args.frame_id:06d}.jpg')))
            if args.write_alpha:
                out_img = np.concatenate([img, msk[:, :, None]], -1)
                imageio.imwrite(tar_path, out_img)
            elif not args.write_alpha and args.white_bkgd:
                msk = (msk / 255).astype(np.float32)
                img = (img / 255).astype(np.float32)
                out_img = img * msk[..., None] + (1 - msk[..., None])
                out_img = (out_img * 255).astype(np.uint8)
                imageio.imwrite(tar_path, out_img)
            elif not args.write_alpha and not args.white_bkgd:
                msk = (msk / 255).astype(np.float32)
                img = (img / 255).astype(np.float32)
                out_img = img * msk[..., None]
                out_img = (out_img * 255).astype(np.uint8)
                imageio.imwrite(tar_path, out_img)
    
    extri_path = join(args.input, 'extri.yml')
    intri_path = join(args.input, 'intri.yml')
    cams_ = cams
    cams = read_camera(intri_path, extri_path)
    
    # transform.json
    frames = []
    for idx, cam in enumerate(cams.keys()):
        if cam == 'basenames':
            continue
        file_path = join('images', f'{cam}.jpg' if not args.write_alpha else f'{cam}.png')
        colmap_im_id = idx
        
        w2c = cams[cam]['RT']
        
        w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0)
        c2w = np.linalg.inv(w2c)
        c2w[0:3, 1:3] *= -1
        c2w = c2w[np.array([1, 0, 2, 3]), :]
        c2w[2, :] *= -1
        transform_matrix = c2w
        
        if cam not in cams_:
            continue
        
        
        frame = {
            'file_path': file_path,
            'transform_matrix': transform_matrix.tolist(),
            'colmap_im_id': colmap_im_id,
            'camera_model': 'OPENCV',
            'fl_x': cams[cam]['K'][0, 0],
            'fl_y': cams[cam]['K'][1, 1],
            'cx': cams[cam]['K'][0, 2],
            'cy': cams[cam]['K'][1, 2],
            'w': cams[cam]['W'],
            'h': cams[cam]['H'],
            'k1': cams[cam]['D'][0, 0],
            'k2': cams[cam]['D'][1, 0],
            'p1': cams[cam]['D'][2, 0],
            'p2': cams[cam]['D'][3, 0],
            'k3': cams[cam]['D'][4, 0],
        }
        frames.append(frame)
    out = {}
    out['frames'] = frames
    
    applied_transform = np.eye(4)[:3, :]
    applied_transform = applied_transform[np.array([1, 0, 2]), :]
    applied_transform[2, :] *= -1
    out["applied_transform"] = applied_transform.tolist()
    with open(join(f'{args.output}', 'transforms.json'), 'w', encoding='utf-8') as f:
        json.dump(out, f, indent=4)
        
if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/nerfstudio/easyvv2nerfstudio.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from lib.utils.data_utils import read_camera
import json

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    args = parser.parse_args()
    return args

def main(args):
    extri_path = join(args.input, 'extri.yml')
    intri_path = join(args.input, 'intri.yml')
    cams = read_camera(intri_path, extri_path)
    
    # transform.json
    frames = []
    for idx, cam in enumerate(cams.keys()):
        if cam == 'basenames':
            continue
        file_path = join('images', f'{cam}.jpg')
        colmap_im_id = idx
        
        w2c = cams[cam]['RT']
        
        w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0)
        c2w = np.linalg.inv(w2c)
        c2w[0:3, 1:3] *= -1
        c2w = c2w[np.array([1, 0, 2, 3]), :]
        c2w[2, :] *= -1
        transform_matrix = c2w
        
        frame = {
            'file_path': file_path,
            'transform_matrix': transform_matrix.tolist(),
            'colmap_im_id': colmap_im_id,
            'camera_model': 'OPENCV',
            'fl_x': cams[cam]['K'][0, 0],
            'fl_y': cams[cam]['K'][1, 1],
            'cx': cams[cam]['K'][0, 2],
            'cy': cams[cam]['K'][1, 2],
            'w': 1920,
            'h': 1080,
            'k1': cams[cam]['dist'][0, 0],
            'k2': cams[cam]['dist'][0, 1],
            'p1': cams[cam]['dist'][0, 2],
            'p2': cams[cam]['dist'][0, 3],
            'k3': cams[cam]['dist'][0, 4],
        }
        frames.append(frame)
    out = {}
    out['frames'] = frames
    
    applied_transform = np.eye(4)[:3, :]
    applied_transform = applied_transform[np.array([1, 0, 2]), :]
    applied_transform[2, :] *= -1
    out["applied_transform"] = applied_transform.tolist()
    with open(join(f'{args.output}', 'transforms.json'), 'w', encoding='utf-8') as f:
        json.dump(out, f, indent=4)
        
if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/nerfstudio/easyvv2transform.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from lib.utils.data_utils import read_camera
import json
from copy import deepcopy

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    args = parser.parse_args()
    return args

def main(args):
    # extri_path = join(args.input, 'extri.yml')
    # intri_path = join(args.input, 'intri.yml')
    # cams = read_camera(intri_path, extri_path)
    transforms = json.load(open(join(args.input, 'transforms.json'), 'r'))
    transforms_train = json.load(open(join(args.input, 'transforms_train.json'), 'r'))
    
    refined_transforms = deepcopy(transforms)
    
    file_paths = [item['file_path'] for item in transforms_train]
    
    def get_index(inference_file_path):
        for file_path in file_paths:
            if inference_file_path in file_path:
                return file_paths.index(file_path)
    
    for frame in refined_transforms['frames']:
        file_path = frame['file_path']
        idx = get_index(file_path)
        transform_matrix = np.eye(4)
        transform_matrix[:3, :4] = np.array(transforms_train[idx]['transform'])
        frame['transform_matrix'] = transform_matrix.tolist()
        
    with open(join(f'{args.input}', 'refined_transforms.json'), 'w', encoding='utf-8') as f:
        json.dump(refined_transforms, f, indent=4)
        
if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/nerfstudio/refinedtransform2transform.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

from lib.utils.data_utils import read_camera
from lib.utils.camera_utils import write_camera
from lib.utils.data_utils import transform2opencv, opencv2tranform
import json
import cv2

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--output_basename', type=str)
    parser.add_argument('--json_file', type=str, default='transform.json')
    args = parser.parse_args()
    return args

def main(args):
    transforms = json.load(open(join(args.input, args.json_file)))
    cameras_out = {}
    cameras_new = {}
    for frame in transforms['frames']:
        key = os.path.basename(frame['file_path']).split('.')[0]
        if frame['camera_model'] == 'OPENCV':
            K = np.array([[frame['fl_x'], 0, frame['cx'], 0, frame['fl_y'], frame['cy'], 0, 0, 1]]).reshape(3, 3)
            dist = np.array([[frame['k1'], frame['k2'], frame['p1'], frame['p2'], frame['k3']]])
        else:
            import ipdb; ipdb.set_trace()
        cameras_out[key] = {'K': K, 'dist': dist, 'H': frame['h'], 'W': frame['w']}
        
        c2w = np.array(frame['transform_matrix']).reshape(4, 4)
        c2w = transform2opencv(c2w)
        ext = np.linalg.inv(c2w)
        
        cam = cameras_out[key].copy()
        t = ext[:3, 3:]
        R = ext[:3, :3]
        cam['R'] = R
        # cam['Rvec'] = cv2.Rodrigues(R)[0]
        cam['T'] = t
        # mapkey[val.name.split('.')[0]] = val.camera_id
        
        cameras_new[key] = cam
        # cameras_new[val.name.split('.')[0].split('/')[0]] = cam
    keys = sorted(list(cameras_new.keys()))
    cameras_new = {key:cameras_new[key] for key in keys}
    write_camera(cameras_new, args.output, args.output_basename)
    
if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/nerfstudio/transform2easyvv.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

import json
from lib.utils.data_utils import transform2opencv
from lib.utils.vis_utils import vis_point_cloud, pcd_array2o3d
import open3d as o3d

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--transform_name', type=str)
    args = parser.parse_args()
    return args

def transform_points_opengl_to_opencv(points):
    # 创建一个新的数组，因为我们不想修改原始数据
    new_points = points.copy()

    # 在OpenGL中，Z轴是向摄像机的反方向，所以我们需要翻转Z轴
    new_points[:, 2] *= -1
    new_points[:, :2] = new_points[:, [1, 0]]

    return new_points

def main(args):
    # datapaser_tranforms = json.load(open(join(args.input, 'dataparser_transforms.json')))
    # scale = datapaser_tranforms['scale']
    # global_trans = np.eye(4)
    # global_trans[:3, :4] = np.asarray(datapaser_tranforms['transform'])
    
    # transform2opencv = lambda x: x
    
    input_transforms = json.load(open(join(args.input, 'transforms.json')))
    input_poses = [transform2opencv(np.asarray(transform['transform_matrix'])) for transform in input_transforms['frames']]
    input_xyz = np.asarray([pose[:3, 3] for pose in input_poses])
    input_pcd = pcd_array2o3d(input_xyz)
    input_scene_pcd = o3d.io.read_point_cloud(join(args.input, 'exports/pcd/point_cloud.ply'))
    points = np.asarray(input_scene_pcd.points)
    points = transform_points_opengl_to_opencv(points)
    input_scene_pcd.points = o3d.utility.Vector3dVector(points)
    # o3d.visualization.draw_geometries([input_pcd, input_scene_pcd])
    # vis_point_cloud(input_xyz)
    
    tranforms = json.load(open(join(args.input, 'camera_paths', args.transform_name)))
    vis_poses = [transform2opencv(np.asarray(transform['camera_to_world']).reshape(4, 4)) for transform in tranforms['camera_path']]
    # vis_poses = [np.linalg.inv(np.linalg.inv(pose) @ np.linalg.inv(global_trans * scale)) for pose in vis_poses]
    # vis_scene_pcd = o3d.io.read_point_cloud(join(args.input, 'point_cloud.ply'))
    vis_xyz = np.asarray([pose[:3, 3] for pose in vis_poses])
    vis_pcd = pcd_array2o3d(vis_xyz)
    
    
    
    o3d.visualization.draw_geometries([vis_pcd, input_pcd, input_scene_pcd])

if __name__ == '__main__':
    args = parse_args()
    main(args)
scripts/nerfstudio/vis_transform_path.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

import objaverse
import multiprocessing
import random
import trimesh


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--download_all', type=bool, default=False)
    args = parser.parse_args()
    return args

def main(args):
    uids = objaverse.load_uids()
    annotations = objaverse.load_annotations(uids[:10])
    processes = multiprocessing.cpu_count()
    random.seed(42)
    uids = objaverse.load_uids()
    random_object_uids = random.sample(uids, 10)
    objects = objaverse.load_objects(
        uids=random_object_uids,
        download_processes=processes
    )
    # objects = objaverse.load_objects(uids=random_object_uids)
    # trimesh.load(list(objects.values())[0]).show() 

if __name__ == '__main__':
    args = parse_args()
    main(args)
    # scripts usage:
    # python scripts/objaverse/download_objects.py

scripts/objaverse/download_objects.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

import imageio
import cv2
import matplotlib.pyplot as plt 
import json
from matplotlib import cm

def is_overlapping(box1, box2, patch_size):
    x1, y1 = box1
    x2, y2 = box2

    overlap_x = (x1 < (x2 + patch_size)) and (x2 < (x1 + patch_size))
    overlap_y = (y1 < (y2 + patch_size)) and (y2 < (y1 + patch_size))
    
    return overlap_x and overlap_y

def find_most_dissimilar_patches(img1, img2, patch_size=128, num_patches=4):
    if img1.shape != img2.shape:
        raise ValueError("The two images must have the same shape.")

    diff = np.abs(img1 - img2)

    height, width = img1.shape[:2]
    best_patches = []
    best_scores = []
    candidate_patches = []

    for i in tqdm(range(0, height - patch_size), desc='traverse the image grid'):
        for j in range(0, width - patch_size):
            patch = diff[i:i+patch_size, j:j+patch_size]
            score = np.mean(patch)
            candidate_patches.append(((i, j), score))

    candidate_patches.sort(key=lambda x: x[1], reverse=True)
    
    for (i, j), score in tqdm(candidate_patches, desc='traverse patches'):
        if len(best_patches) >= num_patches:
            break
        if all([not is_overlapping((i, j), best_i_j, patch_size) for best_i_j, _ in best_patches]):
            best_patches.append(((i, j), score))

    return best_patches

# # 读取图像
# img1 = cv2.imread('image1.jpg', cv2.IMREAD_GRAYSCALE)
# img2 = cv2.imread('image2.jpg', cv2.IMREAD_GRAYSCALE)

# best_patches, best_scores = find_most_dissimilar_patches(img1, img2, patch_size=128, num_patches=4)

# print("Most dissimilar patches at:")
# for (i, j), score in zip(best_patches, best_scores):
#     print(f"Top-left corner: ({i}, {j}), Score: {score}")
#     cv2.rectangle(img1, (j, i), (j+128, i+128), (255, 0, 0), 2)
#     cv2.rectangle(img2, (j, i), (j+128, i+128), (255, 0, 0), 2)

# cv2.imshow('Image 1', img1)
# cv2.imshow('Image 2', img2)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

def vis_patches(img1, img2, patches, patch_size=128, save_path=None):
    prop_cycle = plt.rcParams['axes.prop_cycle']
    colors = prop_cycle.by_key()['color']
    colors_float = [(r, g, b) for r, g, b, _ in [cm.colors.to_rgba(c) for c in colors]]
    for idx, ((i, j), score) in enumerate(patches):
        color = colors_float[idx % len(colors_float)]
        cv2.rectangle(img1, (j, i), (j+patch_size, i+patch_size), color, 2)
        cv2.rectangle(img2, (j, i), (j+patch_size, i+patch_size), color, 2)
        
    imageio.imwrite(save_path, np.concatenate([img1, img2], axis=1))
    
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default='/mnt/data2/home/linhaotong/Datasets/dynerf/cut_beef')
    parser.add_argument('--output', type=str, default='/mnt/data2/home/linhaotong/Datasets/dynerf/cut_beef/im4d')
    parser.add_argument('--test_view', type=str, default='00')
    parser.add_argument('--frame_len', type=int, default=60)
    parser.add_argument('--eval_interval', type=int, default=30)
    parser.add_argument('--eval_ratio', type=float, default=0.5)
    parser.add_argument('--patch_size', type=int, default=128)
    parser.add_argument('--num_patches', type=int, default=4)
    args = parser.parse_args()
    return args

def main(args):
    imgs = []
    for i in tqdm(range(args.frame_len)):
        img_path = join(args.input, 'images', args.test_view, '{:04d}.jpg'.format(i))
        img = (imageio.imread(img_path) / 255.).astype(np.float32)
        imgs.append(img)
        
    os.makedirs(join(args.output, 'test', 'patches'), exist_ok=True)
    os.makedirs(join(args.output, 'test', 'vis'), exist_ok=True)
    
    for i in tqdm(range(0, args.frame_len, args.eval_interval)):
        img = cv2.resize(imgs[i], None, fx=args.eval_ratio, fy=args.eval_ratio, interpolation=cv2.INTER_AREA)
        avg_img = cv2.resize(np.mean(imgs[i:i+args.eval_interval], axis=0), None, fx=args.eval_ratio, fy=args.eval_ratio, interpolation=cv2.INTER_AREA)
        patches = find_most_dissimilar_patches(img, avg_img, patch_size=args.patch_size, num_patches=args.num_patches)
        vis_patches(img, avg_img, patches, patch_size=args.patch_size, save_path=join(args.output, 'test', 'vis', '{:03d}.png'.format(i)))
        save_dict = {'coords': [patch[0] for patch in patches], 'patch_size': args.patch_size}
        save_path = join(args.output, 'test', 'patches', '{:03d}.json'.format(i))
        with open(save_path, 'w', encoding='utf-8') as json_file:
            json.dump(save_dict, json_file, ensure_ascii=False, indent=4)

if __name__ == '__main__':
    args = parse_args()
    main(args)
scripts/preprocess/dynerf/gen_diff_patch.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

import torch
import torch.nn.functional as F
import cv2
import imageio
import glob
from lib.utils.parallel_utils import parallel_execution

np.random.seed(0)

def get_rays_d(h, w, ixt, R):
    x, y = np.meshgrid(np.arange(w), np.arange(h))
    z = np.ones_like(x)
    d = np.stack([x, y, z], axis=2)
    d = np.reshape(d, (-1, 3))
    d = d @ np.linalg.inv(ixt.T)
    d = d @ np.linalg.inv(R.T)
    d = d / np.linalg.norm(d, axis=1, keepdims=True)
    return d


def sample_envmap_image(image: torch.Tensor, ray_d: torch.Tensor):
    sh = ray_d.shape
    if image.ndim == 4:
        image = image[0]
    ray_d = ray_d.view(-1, 3)
    # envmap: H, W, C
    # viewdirs: N, 3

    # https://github.com/zju3dv/InvRender/blob/45e6cdc5e3c9f092b5d10e2904bbf3302152bb2f/code/model/sg_render.py
    image = image.permute(2, 0, 1).unsqueeze(0)

    theta = torch.arccos(ray_d[:, 2]).reshape(-1) - 1e-6
    phi = torch.atan2(ray_d[:, 1], ray_d[:, 0]).reshape(-1)  # 0 - pi

    # normalize to [-1, 1]
    query_y = (theta / torch.pi) * 2 - 1
    query_x = - phi / torch.pi
    grid = torch.stack((query_x, query_y)).permute(1, 0).unsqueeze(0).unsqueeze(0)

    rgb = F.grid_sample(image, grid, align_corners=False, padding_mode='border')
    rgb = rgb.squeeze().permute(1, 0)
    return rgb.view(sh)

def generate_rotation_matrices(num_samples):
    """
    生成旋转矩阵
    :param num_samples: 生成的视角数量
    :return: 旋转矩阵列表
    """
    # 生成球面均匀采样点
    theta = np.arccos(1 - 2 * np.random.rand(num_samples))  # 极角
    phi = 2 * np.pi * np.random.rand(num_samples)  # 方位角

    # 构建旋转矩阵
    rotation_matrices = []
    for i in range(num_samples):
        x = np.sin(theta[i]) * np.cos(phi[i])
        y = np.sin(theta[i]) * np.sin(phi[i])
        z = np.cos(theta[i])
        R = np.array([[np.cos(phi[i]), -np.sin(phi[i]), 0],
                      [np.sin(phi[i]), np.cos(phi[i]), 0],
                      [0, 0, 1]])
        V = np.array([[np.cos(theta[i]), 0, np.sin(theta[i])],
                      [0, 1, 0],
                      [-np.sin(theta[i]), 0, np.cos(theta[i])]])
        rotation_matrix = np.dot(R, V)
        rotation_matrices.append(rotation_matrix)

    return rotation_matrices

def generate_spherical_coordinates(num_samples):
    u = np.arange(0, num_samples+1) / num_samples
    v = np.arange(0, num_samples+1) / num_samples
    u, v = np.meshgrid(u, v)
    u, v = u.flatten(), v.flatten()
    # 将随机数映射到 [-1, 1] 范围内
    azimuth = 2 * np.pi * u
    elevation = np.arccos(2 * v - 1)

    # 将弧度转换为度
    a = np.degrees(azimuth)
    e = np.degrees(elevation)
    return a, e

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, help='Street view image dir')
    parser.add_argument('--output', type=str, help='Output dir')
    parser.add_argument('--size', type=int, default=640, help='image size')
    parser.add_argument('--focal', type=float, default=256, help='focal length')
    args = parser.parse_args()
    return args


def generate_a_e():
    gap_a = 30
    gap_e = 45
    a = np.arange(60, 121, gap_a).astype(np.float32)
    e = np.arange(0, 360, gap_e).astype(np.float32)
    a, e = np.meshgrid(a, e)
    for i in range(a.shape[0]):
        a[i, :] += (np.random.random(len(a[i, :])) * 2 - 1) * gap_a / 2
    for i in range(e.shape[1]):
        e[:, i] += np.random.random(len(e[:, i])) * gap_e
    a, e = a.flatten(), e.flatten()
    return a, e

def process_single_image(img_path, output, h, w, ixt):
    output_dir = join(output, os.path.basename(img_path).split('.')[0])
    os.system(f'mkdir -p {output_dir}')
    a, e = generate_a_e()
    img = (np.array(imageio.imread(img_path)) / 255.).astype(np.float32)
    env_map = torch.tensor(img).float().cuda() if torch.cuda.is_available() else torch.tensor(img).float()
    for i,(a,e) in enumerate(zip(a, e)):
        a_rad = np.radians(a)
        e_rad = np.radians(e)
        y = a_rad
        pitch = np.pi / 2 - e_rad
        roll = 0.
        rotation_matrix_x = cv2.Rodrigues(np.array([a_rad, 0, 0]))[0]
        rotation_matrix_y = cv2.Rodrigues(np.array([0, pitch, 0]))[0]
        rotation_matrix_z = cv2.Rodrigues(np.array([0, 0, roll]))[0]
        rotation_matrix = rotation_matrix_z @ rotation_matrix_y @ rotation_matrix_x
        rays_d = torch.tensor(get_rays_d(h, w, ixt, rotation_matrix)).float().cuda() if torch.cuda.is_available() else torch.tensor(get_rays_d(h, w, ixt, rotation_matrix)).float()
        img = sample_envmap_image(env_map, rays_d)
        imageio.imwrite(f'{output_dir}/{i:03d}.jpg', (img.reshape(h,w,3).cpu().numpy() * 255.).astype(np.uint8))
        
def main(args):
    img_dir = args.input
    h, w = args.size, args.size
    focal = args.focal
    ixt = np.array([[focal, 0., w/2], [0, focal, h/2], [0, 0, 1 ]])
    
    img_paths = glob.glob(f'{img_dir}/*.jpg')
    img_paths = sorted(img_paths)
    # process_single_image(img_paths[0], args.output, h, w, ixt)
    # return
    parallel_execution(
        img_paths, 
        args.output,
        h, w, ixt,
        action=process_single_image,
        print_progress=True,
        num_processes=4,
    )
    
if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/sphere_images/sample_from_sphere.py
zju3dv/im4d
import os
from os.path import join
import argparse
import sys
import numpy as np
from tqdm import tqdm
sys.path.append('.')

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    args = parser.parse_args()
    return args

def main(args):
    pass

if __name__ == '__main__':
    args = parse_args()
    main(args)


scripts/template.py
zju3dv/im4d
import os
import torch
import torch.distributed as dist
import torch.multiprocessing

from lib.config import cfg, logger
from lib.networks import make_network
from lib.train import make_trainer, make_optimizer, make_lr_scheduler, make_recorder, set_lr_scheduler
from lib.datasets import make_data_loader
from lib.utils.net_utils import load_model, save_model, save_trained_config, load_pretrain
from lib.evaluators import make_evaluator
from lib.utils.msg_utils import send_msg
torch.backends.cudnn.benchmark = False
# torch.autograd.set_detect_anomaly(True)

if cfg.fix_random:
    torch.manual_seed(0)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def train(cfg, network):
    train_loader = make_data_loader(cfg,
                                    is_train=True,
                                    is_distributed=cfg.distributed,
                                    max_iter=cfg.ep_iter)
    if cfg.skip_eval:
        val_loader = None
    else:
        val_loader = make_data_loader(cfg, is_train=False)
    trainer = make_trainer(cfg, network, train_loader)
    optimizer = make_optimizer(cfg, network)
    scheduler = make_lr_scheduler(cfg, optimizer)
    recorder = make_recorder(cfg)
    evaluator = make_evaluator(cfg)

    begin_epoch = load_model(network,
                             optimizer,
                             scheduler,
                             recorder,
                             cfg.trained_model_dir,
                             resume=cfg.resume)

    if begin_epoch == 0 and cfg.pretrain != '':
        load_pretrain(network, cfg.pretrain)
    set_lr_scheduler(cfg, scheduler)
    save_trained_config(cfg)
    
    for epoch in range(begin_epoch, cfg.train.epoch):
        recorder.epoch = epoch
        if cfg.distributed:
            train_loader.batch_sampler.sampler.set_epoch(epoch)

        train_loader.dataset.epoch = epoch

        trainer.train(epoch, train_loader, optimizer, recorder, scheduler)
        # scheduler.step()

        if (epoch + 1) % cfg.save_ep == 0 and cfg.local_rank == 0:
            save_model(network, optimizer, scheduler, recorder,
                       cfg.trained_model_dir, epoch)

        if (epoch + 1) % cfg.save_latest_ep == 0 and cfg.local_rank == 0:
            save_model(network,
                       optimizer,
                       scheduler,
                       recorder,
                       cfg.trained_model_dir,
                       epoch,
                       last=True)

        if not cfg.skip_eval and (epoch + 1) % cfg.eval_ep == 0 and cfg.local_rank == 0:
            trainer.val(epoch, val_loader, evaluator, recorder)

    return network

def synchronize():
    """
    Helper function to synchronize (barrier) among all processes when
    using distributed training
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()

def main():
    if cfg.distributed:
        cfg.local_rank = int(os.environ['RANK']) % torch.cuda.device_count()
        torch.cuda.set_device(cfg.local_rank)
        torch.distributed.init_process_group(backend="nccl",
                                             init_method="env://")
        synchronize()
        
    network = make_network(cfg)
    send_msg('Training start: ' + cfg.exp_name + ' Machine: ' + cfg.machine, cfg)
    train(cfg, network)
    send_msg('Training success: ' + cfg.exp_name + ' Machine: ' + cfg.machine, cfg)
    if cfg.local_rank == 0:
        logger.info('Success!')
        logger.info('='*80)
    os.system('kill -9 {}'.format(os.getpid()))

if __name__ == "__main__":
    main()
train_net.py
zju3dv/im4d