diff --git a/.circleci/config.yml b/.circleci/config.yml index ce939130db..fd4b3e2823 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -45,7 +45,8 @@ jobs: - run: name: run clang-format command: | - find . -not -path "*/\.*" -not -path "*/deps/*" -not -path "*/obsolete/*" -not -path "*/build/*" | grep -E ".*\.cpp$|.*\.h$|.*\.cu$|.*\.hpp$" | xargs -I {} bash -c "diff -u <(cat {}) <(clang-format -style=file {})" + find . -not -path "*/\.*" -not -path "*/deps/*" -not -path "*/obsolete/*" -not -path "*/build/*" | grep -E ".*\.(cpp|h|cu|hpp|cuh)$" | xargs -I {} bash -c "diff -u <(cat {}) <(clang-format -style=file {})" + js_lint_and_test: docker: - image: circleci/node:11.9 @@ -241,13 +242,13 @@ jobs: name: Run sim benchmark command: | while [ ! -f ./cuda_installed ]; do sleep 2; done # wait for CUDA - export PATH=$HOME/miniconda/bin:$PATH + export PATH=$HOME/miniconda/bin:/usr/local/cuda/bin:$PATH . activate habitat; cd habitat-sim python examples/example.py --scene data/scene_datasets/habitat-test-scenes/van-gogh-room.glb --silent --test_fps_regression $FPS_THRESHOLD - run: name: Run sim tests command: | - export PATH=$HOME/miniconda/bin:$PATH + export PATH=$HOME/miniconda/bin:/usr/local/cuda/bin:$PATH . activate habitat; cd habitat-sim export PYTHONPATH=$PYTHONPATH:$(pwd) diff --git a/.gitignore b/.gitignore index eae06e8406..7d89f72d83 100644 --- a/.gitignore +++ b/.gitignore @@ -164,7 +164,7 @@ run_example *.swp tests/17DRP5sb8fy -data/ +/data/ *.zip habitat_sim/_ext build_js/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1e4e7643a5..4521beca50 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,6 +16,7 @@ repos: rev: v1.9.2 hooks: - id: seed-isort-config + language_version: python3.6 - repo: https://github.com/pre-commit/mirrors-isort rev: v4.3.20 @@ -35,7 +36,7 @@ repos: name: Run clang-format entry: clang-format -i -style=file types: [text] - files: '.*\.cpp|.*\.h' + files: '.*\.(cpp|h|hpp|cu|cuh)$' language: system - repo: https://github.com/pre-commit/mirrors-eslint diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000000..6c9881b2c4 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +graft habitat_sim/sensors/noise_models/data diff --git a/README.md b/README.md index f06a0dd5a7..f200482d3e 100644 --- a/README.md +++ b/README.md @@ -359,6 +359,11 @@ The Habitat project would not have been possible without the support and contrib Specifically, the noise model used for the noisy control functions named `pyrobot_*` and defined in `habitat_sim/agent/controls/pyrobot_noisy_controls.py` +* If you use the Redwood Depth Noise Model, please cite their [paper](http://redwood-data.org/indoor/) + + Specifically, the noise model defined in `habitat_sim/sensors/noise_models/redwood_depth_noise_model.py` and `src/esp/sensor/RedwoodNoiseModel.*` + + ## License Habitat-Sim is MIT licensed. See the LICENSE file for details. diff --git a/docs/conf.py b/docs/conf.py index 2633ff47bb..6b2508f970 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -30,7 +30,7 @@ FAVICON = "habitat-blue.png" MAIN_PROJECT_URL = "/" INPUT_MODULES = [habitat_sim] -INPUT_DOCS = ["docs.rst", "gfx.rst"] +INPUT_DOCS = ["docs.rst", "gfx.rst", "noise_models.rst"] INPUT_PAGES = [ "pages/index.rst", "pages/new-actions.rst", diff --git a/docs/noise_models.rst b/docs/noise_models.rst new file mode 100644 index 0000000000..453c8a8610 --- /dev/null +++ b/docs/noise_models.rst @@ -0,0 +1,55 @@ +.. py:module:: habitat_sim.sensors.noise_models + :summary: Library of Sensor noise models + + + A library of noise models to close the gap between simulated observations + and observations from real sensors. + + A noise model can be applied to a sensor by specifying the name of the noise + model in the `sensor.SensorSpec.noise_model` feild. + Arguments can be passed to the noise model constructor as keyword arguments using + the `sensor.SensorSpec.noise_model_kwargs` field. For instance, to use the `RedwoodDepthNoiseModel` + with a ``noise_multiplier`` of 5 + + .. code:: py + + sensor_spec.noise_model = "RedwoodDepthNoiseModel" + sensor_spec.noise_model_kwargs = dict(noise_multiplier=5) + + + + These noise models are commonly the result of contributions from various research projects. + If you use a noise model in your research, please cite the relevant work specified by the docummentation + + + **Depth Noise Models** + + * Redwood Noise Model for PrimSense depth cameras: `RedwoodDepthNoiseModel` + +.. py:class:: habitat_sim.sensors.noise_models.NoSensorNoiseModel + :summary: No noise noise model. Simply returns a copy of the input + + Accessible from the registry under the name ``"None"`` + +.. py:class:: habitat_sim.sensors.noise_models.RedwoodDepthNoiseModel + :summary: Redwood Noise Model for PrimSense depth cameras + + Accessible from the registry under the name ``"RedwoodDepthNoiseModel"`` + + Implements the noise model provided by http://redwood-data.org/indoor/dataset.html + + If you use this noise model, please cite:: + + .. code:: bibtex + + @inproceedings{choi2015robust, + title={Robust reconstruction of indoor scenes}, + author={Choi, Sungjoon and Zhou, Qian-Yi and Koltun, Vladlen}, + booktitle={Proceedings of the IEEE Conference on Computer Vision and + Pattern Recognition}, pages={5556--5565}, year={2015} + } + + +.. py:function:: habitat_sim.sensors.noise_models.RedwoodDepthNoiseModel.__init__ + :param gpu_device_id: The ID of CUDA device to use (only applicable if habitat-sim was built with ``--with-cuda``) + :param noise_multiplier: Multipler for the Gaussian random-variables. This reduces or increases the amount of noise diff --git a/habitat_sim/agent/agent.py b/habitat_sim/agent/agent.py index 6e5d82bf85..536fe409be 100644 --- a/habitat_sim/agent/agent.py +++ b/habitat_sim/agent/agent.py @@ -12,7 +12,7 @@ import habitat_sim.bindings as hsim import habitat_sim.errors -from habitat_sim.sensors import SensorSuite +from habitat_sim.sensors.sensor_suite import SensorSuite from habitat_sim.utils.common import ( quat_from_coeffs, quat_from_magnum, diff --git a/habitat_sim/registry.py b/habitat_sim/registry.py index 41c600a3f1..d80fb88b49 100644 --- a/habitat_sim/registry.py +++ b/habitat_sim/registry.py @@ -8,8 +8,6 @@ import re from typing import Optional, Type -from habitat_sim.agent.controls import SceneNodeControl - __all__ = ["registry"] @@ -37,12 +35,12 @@ class _Registry: @classmethod def register_move_fn( cls, - controller: Optional[Type[SceneNodeControl]] = None, + controller: Optional[Type] = None, *, name: Optional[str] = None, body_action: bool = None, ): - r"""Registers a new control with Habitat-Sim. Registered conrtols can + r"""Registers a new control with Habitat-Sim. Registered controls can then be retrieved via `get_move_fn()` See `new-actions `_ for an example of how to add new actions @@ -62,6 +60,7 @@ def register_move_fn( assert ( body_action is not None ), "body_action must be explicitly set to True or False" + from habitat_sim.agent.controls.controls import SceneNodeControl def _wrapper(controller: Type[SceneNodeControl]): assert issubclass( @@ -79,17 +78,54 @@ def _wrapper(controller: Type[SceneNodeControl]): else: return _wrapper(controller) + @classmethod + def register_noise_model( + cls, noise_model: Optional[Type] = None, *, name: Optional[str] = None + ): + r"""Registers a new sensor noise model with Habitat-Sim + + :param noise_model: The class of the noise model to register + If `None`, will return a wrapper for use with decorator syntax + :param name: The name to register the noise model with + If `None`, will register with the name of the noise_model + """ + from habitat_sim.sensors.noise_models.sensor_noise_model import SensorNoiseModel + + def _wrapper(noise_model: Type[SensorNoiseModel]): + assert issubclass( + noise_model, SensorNoiseModel + ), "All noise_models must inherit from habitat_sim.sensor.SensorNoiseModel" + + cls._mapping["sensor_noise_model"][ + noise_model.__name__ if name is None else name + ] = noise_model + + return noise_model + + if noise_model is None: + return _wrapper + else: + return _wrapper(noise_model) + @classmethod def _get_impl(cls, _type, name): return cls._mapping[_type].get(name, None) @classmethod - def get_move_fn(cls, name: str) -> SceneNodeControl: + def get_move_fn(cls, name: str): r"""Retrieve the move_fn register under ``name`` :param name: The name provided to `register_move_fn` """ return cls._get_impl("move_fn", name) + @classmethod + def get_noise_model(cls, name: str): + r"""Retrieve the noise_model registered under ``name`` + + :param name: The name provided to `register_noise_model` + """ + return cls._get_impl("sensor_noise_model", name) + registry = _Registry() diff --git a/habitat_sim/sensors/__init__.py b/habitat_sim/sensors/__init__.py index 249bc99396..45992bc9ec 100644 --- a/habitat_sim/sensors/__init__.py +++ b/habitat_sim/sensors/__init__.py @@ -4,6 +4,8 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from habitat_sim.sensors import noise_models + from .sensor_suite import SensorSuite -__all__ = ["SensorSuite"] +__all__ = ["SensorSuite", "noise_models"] diff --git a/habitat_sim/sensors/noise_models/__init__.py b/habitat_sim/sensors/noise_models/__init__.py new file mode 100644 index 0000000000..bddc21826d --- /dev/null +++ b/habitat_sim/sensors/noise_models/__init__.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any, Dict + +from habitat_sim.registry import registry +from habitat_sim.sensors.noise_models.no_noise_model import NoSensorNoiseModel +from habitat_sim.sensors.noise_models.redwood_depth_noise_model import ( + RedwoodDepthNoiseModel, +) +from habitat_sim.sensors.noise_models.sensor_noise_model import SensorNoiseModel + + +def make_sensor_noise_model(name: str, kwargs: Dict[str, Any]) -> SensorNoiseModel: + r"""Constructs a noise model using the given name and keyword arguments + + :param name: The name of the noise model in the `habitat_sim.registry` + :param kwargs: The keyword arguments to be passed to the constructor of the noise model + """ + + model = registry.get_noise_model(name) + assert model is not None, "Could not find a noise model for name '{}'".format(name) + + return model(**kwargs) + + +__all__ = [ + "make_sensor_noise_model", + "SensorNoiseModel", + "RedwoodDepthNoiseModel", + "NoSensorNoiseModel", +] diff --git a/habitat_sim/sensors/noise_models/data/redwood-depth-dist-model.npy b/habitat_sim/sensors/noise_models/data/redwood-depth-dist-model.npy new file mode 100644 index 0000000000..9649064013 Binary files /dev/null and b/habitat_sim/sensors/noise_models/data/redwood-depth-dist-model.npy differ diff --git a/habitat_sim/sensors/noise_models/no_noise_model.py b/habitat_sim/sensors/noise_models/no_noise_model.py new file mode 100644 index 0000000000..93a8200f0e --- /dev/null +++ b/habitat_sim/sensors/noise_models/no_noise_model.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np + +from habitat_sim.registry import registry +from habitat_sim.sensor import SensorType +from habitat_sim.sensors.noise_models.sensor_noise_model import SensorNoiseModel + +try: + import torch +except ImportError: + torch = None + + +@registry.register_noise_model(name="None") +class NoSensorNoiseModel(SensorNoiseModel): + @staticmethod + def is_valid_sensor_type(sensor_type: SensorType) -> bool: + return True + + def apply(self, x): + if isinstance(x, np.ndarray): + return x.copy() + elif torch is not None and torch.is_tensor(x): + return x.clone() + else: + return x diff --git a/habitat_sim/sensors/noise_models/redwood_depth_noise_model.py b/habitat_sim/sensors/noise_models/redwood_depth_noise_model.py new file mode 100644 index 0000000000..86e9fc826a --- /dev/null +++ b/habitat_sim/sensors/noise_models/redwood_depth_noise_model.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os.path as osp + +import attr +import numba +import numpy as np + +from habitat_sim.bindings import cuda_enabled +from habitat_sim.registry import registry +from habitat_sim.sensor import SensorType +from habitat_sim.sensors.noise_models.sensor_noise_model import SensorNoiseModel + +if cuda_enabled: + from habitat_sim._ext.habitat_sim_bindings import RedwoodNoiseModelGPUImpl + import torch + + +# Read about the noise model here: http://www.alexteichman.com/octo/clams/ +# Original source code: http://redwood-data.org/indoor/data/simdepth.py +@numba.jit(nopython=True) +def _undistort(x, y, z, model): + i2 = int((z + 1) / 2) + i1 = int(i2 - 1) + a = (z - (i1 * 2.0 + 1.0)) / 2.0 + x = x // 8 + y = y // 6 + f = (1 - a) * model[y, x, min(max(i1, 0), 4)] + a * model[y, x, min(i2, 4)] + + if f < 1e-5: + return 0 + else: + return z / f + + +@numba.jit(nopython=True, parallel=True) +def _simulate(gt_depth, model, noise_multiplier): + noisy_depth = np.empty_like(gt_depth) + + H, W = gt_depth.shape + ymax, xmax = H - 1, W - 1 + + rand_nums = np.random.randn(H, W, 3).astype(np.float32) + for j in range(H): + for i in range(W): + y = int( + min(max(j + rand_nums[j, i, 0] * 0.25 * noise_multiplier, 0.0), ymax) + + 0.5 + ) + x = int( + min(max(i + rand_nums[j, i, 1] * 0.25 * noise_multiplier, 0.0), xmax) + + 0.5 + ) + + # Downsample + d = gt_depth[y - y % 2, x - x % 2] + # If the depth is greater than 10, the sensor will just return 0 + if d >= 10.0: + noisy_depth[j, i] = 0.0 + else: + # Distort + # The noise model was originally made for a 640x480 sensor, + # so re-map our arbitrarily sized sensor to that size! + undistorted_d = _undistort( + int(x / xmax * 639.0 + 0.5), int(y / ymax * 479.0 + 0.5), d, model + ) + + if undistorted_d == 0.0: + noisy_depth[j, i] = 0.0 + else: + denom = round( + ( + 35.130 / undistorted_d + + rand_nums[j, i, 2] * 0.027778 * noise_multiplier + ) + * 8.0 + ) + if denom <= 1e-5: + noisy_depth[j, i] = 0.0 + else: + noisy_depth[j, i] = 35.130 * 8.0 / denom + + return noisy_depth + + +@attr.s(auto_attribs=True) +class RedwoodNoiseModelCPUImpl: + model: np.ndarray + noise_multiplier: float + + def __attrs_post_init__(self): + self.model = self.model.reshape(self.model.shape[0], -1, 4) + + def simulate(self, gt_depth): + return _simulate(gt_depth, self.model, self.noise_multiplier) + + +@registry.register_noise_model +@attr.s(auto_attribs=True, kw_only=True) +class RedwoodDepthNoiseModel(SensorNoiseModel): + noise_multiplier: float = 1.0 + + def __attrs_post_init__(self): + dist = np.load( + osp.join(osp.dirname(__file__), "data", "redwood-depth-dist-model.npy") + ) + + if cuda_enabled: + self._impl = RedwoodNoiseModelGPUImpl( + dist, self.gpu_device_id, self.noise_multiplier + ) + else: + self._impl = RedwoodNoiseModelCPUImpl(dist, self.noise_multiplier) + + @staticmethod + def is_valid_sensor_type(sensor_type: SensorType) -> bool: + return sensor_type == SensorType.DEPTH + + def simulate(self, gt_depth): + if cuda_enabled: + if torch.is_tensor(gt_depth): + noisy_depth = torch.empty_like(gt_depth) + rows, cols = gt_depth.size() + self._impl.simulate_from_gpu( + gt_depth.data_ptr(), rows, cols, noisy_depth.data_ptr() + ) + return noisy_depth + else: + return self._impl.simulate_from_cpu(gt_depth) + else: + return self._impl.simulate(gt_depth) + + def apply(self, gt_depth): + r"""Alias of `simulate()` to conform to base-class and expected API + """ + return self.simulate(gt_depth) diff --git a/habitat_sim/sensors/noise_models/sensor_noise_model.py b/habitat_sim/sensors/noise_models/sensor_noise_model.py new file mode 100644 index 0000000000..02f7d2cbd1 --- /dev/null +++ b/habitat_sim/sensors/noise_models/sensor_noise_model.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 + +# Copyright (c) Facebook, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import abc +from typing import Dict, Optional, Type + +import attr + +from habitat_sim.sensor import SensorType + + +@attr.s(auto_attribs=True, kw_only=True) +class SensorNoiseModel(abc.ABC): + r"""Base class for all sensor noise models + + :param gpu_device_id: The ID of the CUDA device to use (only applicable to + noise models that proivde a CUDA implementation and, generally, if habitat-sim was build with ``--with-cuda``) + """ + gpu_device_id: Optional[int] = None + + @staticmethod + @abc.abstractmethod + def is_valid_sensor_type(sensor_type: SensorType) -> bool: + r"""Used to determine whether or not the noise model + is applicable to the sensor type + + :return: True if this noise model can be applied to this sensor input type + """ + pass + + @abc.abstractmethod + def apply(self, sensor_observation): + r"""Applies the noise model to the sensor observation + + :param sensor_observation: The clean sensor observation. + Should not be modified. + + :return: The sensor observation with noise applied. + """ + pass + + def __call__(self, sensor_observation): + r"""Alias of `apply()` + """ + return self.apply(sensor_observation) diff --git a/habitat_sim/simulator.py b/habitat_sim/simulator.py index 6b6508899e..73e3973bc8 100644 --- a/habitat_sim/simulator.py +++ b/habitat_sim/simulator.py @@ -18,6 +18,7 @@ from habitat_sim.logging import logger from habitat_sim.nav import GreedyGeodesicFollower from habitat_sim.physics import MotionType +from habitat_sim.sensors.noise_models import make_sensor_noise_model from habitat_sim.utils.common import quat_from_angle_axis torch = None @@ -348,12 +349,18 @@ def __init__(self, sim, agent, sensor_id): dtype=np.uint8, ) + noise_model_kwargs = self._spec.noise_model_kwargs + self._noise_model = make_sensor_noise_model( + self._spec.noise_model, + {"gpu_device_id": self._sim.gpu_device, **noise_model_kwargs}, + ) + assert self._noise_model.is_valid_sensor_type( + self._spec.sensor_type + ), "Noise model '{}' is not valid for sensor '{}'".format( + self._spec.noise_model, self._spec.uuid + ) + def draw_observation(self): - # draw the scene with the visual sensor: - # it asserts the sensor is a visual sensor; - # internally it will set the camera parameters (from the sensor) to the - # default render camera in the scene so that - # it has correct modelview matrix, projection matrix to render the scene # sanity check: # see if the sensor is attached to a scene graph, otherwise it is invalid, @@ -400,7 +407,7 @@ def get_observation(self): else: tgt.read_frame_rgba_gpu(self._buffer.data_ptr()) - return self._buffer.flip(0).clone() + obs = self._buffer.flip(0) else: size = self._sensor_object.framebuffer_size @@ -421,7 +428,9 @@ def get_observation(self): ) ) - return np.flip(self._buffer, axis=0).copy() + obs = np.flip(self._buffer, axis=0) + + return self._noise_model(obs) def close(self): self._sim = None diff --git a/pyproject.toml b/pyproject.toml index 8de4d0a614..6a5a79de57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.isort] skip_glob = ["*/deps/*", "*/build/*", "*/obselete/*"] -known_third_party = ["PIL", "attr", "conf", "demo_runner", "hypothesis", "magnum", "numpy", "pytest", "quaternion", "scipy", "settings", "setuptools", "tqdm"] +known_third_party = ["PIL", "attr", "conf", "demo_runner", "hypothesis", "magnum", "numba", "numpy", "pytest", "quaternion", "scipy", "settings", "setuptools", "tqdm"] multi_line_output = 3 force_grid_wrap = false line_length = 88 diff --git a/setup.py b/setup.py index fddf6f1afc..8116b93970 100644 --- a/setup.py +++ b/setup.py @@ -377,6 +377,7 @@ def load(filename): # add custom build_ext command cmdclass=dict(build_ext=CMakeBuild), zip_safe=False, + include_package_data=True, ) pymagnum_build_dir = osp.join( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 51a32c74e7..d8b62933f7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -4,7 +4,14 @@ cmake_minimum_required(VERSION 3.10) -project(esp) +option(BUILD_WITH_CUDA "Build Habitat-Sim with CUDA features enabled -- Requires CUDA" OFF) + +if(BUILD_WITH_CUDA) + project(esp LANGUAGES C CXX CUDA) + find_library(CUDART_LIBRARY cudart ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}) +else() + project(esp LANGUAGES C CXX) +endif() if(MSVC) add_definitions(/DNOMINMAX) @@ -24,7 +31,6 @@ option(BUILD_PYTHON_BINDINGS "Whether to build python bindings" ON) option(BUILD_DATATOOL "Whether to build datatool utility binary" ON) option(BUILD_PTEX_SUPPORT "Whether to build ptex mesh support" ON) option(BUILD_GUI_VIEWERS "Whether to build GUI viewer utility binary" OFF) -option(BUILD_WITH_CUDA "Build Habitat-Sim with CUDA features enabled -- Requires CUDA" OFF) option(BUILD_WITH_BULLET "Build Habitat-Sim with Bullet physics enabled -- Requires Bullet" OFF) option(BUILD_TEST "Build test binaries" OFF) option(USE_SYSTEM_ASSIMP "Use system Assimp instead of a bundled submodule" OFF) diff --git a/src/esp/bindings/bindings.cpp b/src/esp/bindings/bindings.cpp index 7f69561a97..6426854862 100644 --- a/src/esp/bindings/bindings.cpp +++ b/src/esp/bindings/bindings.cpp @@ -23,6 +23,10 @@ using namespace py::literals; #include "esp/sensor/PinholeCamera.h" #include "esp/sensor/Sensor.h" +#ifdef ESP_BUILD_WITH_CUDA +#include "esp/sensor/RedwoodNoiseModel.h" +#endif + #include #include #include @@ -354,8 +358,7 @@ PYBIND11_MODULE(habitat_sim_bindings, m) { .value("DEPTH", SensorType::DEPTH) .value("SEMANTIC", SensorType::SEMANTIC); - // ==== SensorSpec ==== - py::class_(m, "SensorSpec") + py::class_(m, "SensorSpec", py::dynamic_attr()) .def(py::init(&SensorSpec::create<>)) .def_readwrite("uuid", &SensorSpec::uuid) .def_readwrite("sensor_type", &SensorSpec::sensorType) @@ -368,6 +371,19 @@ PYBIND11_MODULE(habitat_sim_bindings, m) { .def_readwrite("encoding", &SensorSpec::encoding) .def_readwrite("gpu2gpu_transfer", &SensorSpec::gpu2gpuTransfer) .def_readwrite("observation_space", &SensorSpec::observationSpace) + .def_readwrite("noise_model", &SensorSpec::noiseModel) + .def_property( + "noise_model_kwargs", + [](SensorSpec& self) -> py::dict { + py::handle handle = py::cast(self); + if (!py::hasattr(handle, "__noise_model_kwargs")) { + py::setattr(handle, "__noise_model_kwargs", py::dict()); + } + return py::getattr(handle, "__noise_model_kwargs"); + }, + [](SensorSpec& self, py::dict v) { + py::setattr(py::cast(self), "__noise_model_kwargs", v); + }) .def("__eq__", [](const SensorSpec& self, const SensorSpec& other) -> bool { return self == other; @@ -443,6 +459,20 @@ PYBIND11_MODULE(habitat_sim_bindings, m) { .def_property_readonly("framebuffer_size", &Sensor::framebufferSize) .def_property_readonly("render_target", &Sensor::renderTarget); +#ifdef ESP_BUILD_WITH_CUDA + py::class_( + m, "RedwoodNoiseModelGPUImpl") + .def(py::init(&RedwoodNoiseModelGPUImpl::create_unique< + const Eigen::Ref&, int, float>)) + .def("simulate_from_cpu", &RedwoodNoiseModelGPUImpl::simulateFromCPU) + .def("simulate_from_gpu", [](RedwoodNoiseModelGPUImpl& self, + std::size_t devDepth, const int rows, + const int cols, std::size_t devNoisyDepth) { + self.simulateFromGPU(reinterpret_cast(devDepth), rows, + cols, reinterpret_cast(devNoisyDepth)); + }); +#endif + // ==== PinholeCamera (subclass of Sensor) ==== py::class_, diff --git a/src/esp/core/esp.h b/src/esp/core/esp.h index 69f48e72da..72a41f2cf4 100644 --- a/src/esp/core/esp.h +++ b/src/esp/core/esp.h @@ -31,6 +31,8 @@ typedef Matrix Vector4uc; typedef Matrix Vector4ui; typedef Matrix Vector4ul; +typedef Matrix RowMatrixXf; + //! Eigen JSON string format specification static const IOFormat kJsonFormat(StreamPrecision, DontAlignCols, diff --git a/src/esp/gfx/CMakeLists.txt b/src/esp/gfx/CMakeLists.txt index 3d162a7fb1..89be0fccf7 100644 --- a/src/esp/gfx/CMakeLists.txt +++ b/src/esp/gfx/CMakeLists.txt @@ -77,16 +77,15 @@ target_include_directories(gfx ) if(BUILD_WITH_CUDA) - find_package(CUDA 8.0 REQUIRED) target_include_directories(gfx PRIVATE - ${CUDA_INCLUDE_DIRS} + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} ${CMAKE_CURRENT_LIST_DIR}/cuda_helpers ) target_link_libraries(gfx PUBLIC - ${CUDA_LIBRARIES} + ${CUDART_LIBRARY} ) endif() diff --git a/src/esp/sensor/CMakeLists.txt b/src/esp/sensor/CMakeLists.txt index 921ff483d0..52dd8f89f2 100644 --- a/src/esp/sensor/CMakeLists.txt +++ b/src/esp/sensor/CMakeLists.txt @@ -1,13 +1,50 @@ -add_library(sensor STATIC +set(sensor_SOURCES PinholeCamera.cpp PinholeCamera.h Sensor.cpp Sensor.h ) + +if(BUILD_WITH_CUDA) + list(APPEND sensor_SOURCES + RedwoodNoiseModel.cpp + RedwoodNoiseModel.h + ) +endif() + + +add_library(sensor STATIC ${sensor_SOURCES}) + target_link_libraries(sensor PUBLIC core gfx scene ) + +if(BUILD_WITH_CUDA) + add_library(noise_model_kernels STATIC + RedwoodNoiseModel.cu + RedwoodNoiseModel.cuh + ) + target_link_libraries(noise_model_kernels + PUBLIC + ${CUDART_LIBRARY} + ) + target_include_directories(noise_model_kernels + PRIVATE + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ) + target_compile_features(noise_model_kernels PUBLIC cxx_std_11) + + + target_link_libraries(sensor + PUBLIC + noise_model_kernels + ) + target_include_directories(sensor + PRIVATE + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ) +endif() diff --git a/src/esp/sensor/RedwoodNoiseModel.cpp b/src/esp/sensor/RedwoodNoiseModel.cpp new file mode 100644 index 0000000000..f7c1126f36 --- /dev/null +++ b/src/esp/sensor/RedwoodNoiseModel.cpp @@ -0,0 +1,78 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include "RedwoodNoiseModel.h" + +namespace esp { +namespace sensor { + +namespace { + +struct CudaDeviceContext { + explicit CudaDeviceContext(const int deviceId) { + cudaGetDevice(¤tDevice); + if (deviceId != currentDevice) { + cudaSetDevice(deviceId); + setDevice = true; + } + } + + ~CudaDeviceContext() { + if (setDevice) + cudaSetDevice(currentDevice); + } + + private: + bool setDevice = false; + int currentDevice = -1; +}; + +} // namespace + +RedwoodNoiseModelGPUImpl::RedwoodNoiseModelGPUImpl( + const Eigen::Ref model, + const int gpuDeviceId, + const float noiseMultiplier) + : gpuDeviceId_{gpuDeviceId}, noiseMultiplier_{noiseMultiplier} { + CudaDeviceContext ctx{gpuDeviceId_}; + + cudaMalloc(&devModel_, model.rows() * model.cols() * sizeof(float)); + cudaMemcpy(devModel_, model.data(), + model.rows() * model.cols() * sizeof(float), + cudaMemcpyHostToDevice); + curandStates_ = impl::getCurandStates(); +} + +RedwoodNoiseModelGPUImpl::~RedwoodNoiseModelGPUImpl() { + CudaDeviceContext ctx{gpuDeviceId_}; + + if (devModel_ != nullptr) + cudaFree(devModel_); + impl::freeCurandStates(curandStates_); +} + +Eigen::RowMatrixXf RedwoodNoiseModelGPUImpl::simulateFromCPU( + const Eigen::Ref depth) { + CudaDeviceContext ctx{gpuDeviceId_}; + + Eigen::RowMatrixXf noisyDepth(depth.rows(), depth.cols()); + + impl::simulateFromCPU(depth.data(), depth.rows(), depth.cols(), devModel_, + curandStates_, noiseMultiplier_, noisyDepth.data()); + return noisyDepth; +} + +void RedwoodNoiseModelGPUImpl::simulateFromGPU(const float* devDepth, + const int rows, + const int cols, + float* devNoisyDepth) { + CudaDeviceContext ctx{gpuDeviceId_}; + impl::simulateFromGPU(devDepth, rows, cols, devModel_, curandStates_, + noiseMultiplier_, devNoisyDepth); +} + +} // namespace sensor +} // namespace esp diff --git a/src/esp/sensor/RedwoodNoiseModel.cu b/src/esp/sensor/RedwoodNoiseModel.cu new file mode 100644 index 0000000000..40db450df2 --- /dev/null +++ b/src/esp/sensor/RedwoodNoiseModel.cu @@ -0,0 +1,185 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +#include "RedwoodNoiseModel.cuh" + +#include + +#include +#include + +namespace { +const int MODEL_N_DIMS = 5; +const int MODEL_N_COLS = 80; + +// Read about the noise model here: http://www.alexteichman.com/octo/clams/ +// Original source code: http://redwood-data.org/indoor/data/simdepth.py +__device__ float undistort(const int _x, + const int _y, + const float z, + const float* __restrict__ model) { + const int i2 = (z + 1) / 2; + const int i1 = i2 - 1; + const float a = (z - (i1 * 2 + 1)) / 2.0f; + const int x = _x / 8; + const int y = _y / 6; + + const float f = + (1 - a) * + model[(y * MODEL_N_COLS + x) * MODEL_N_DIMS + min(max(i1, 0), 4)] + + a * model[(y * MODEL_N_COLS + x) * MODEL_N_DIMS + min(i2, 4)]; + + if (f <= 1e-5f) + return 0; + else + return z / f; +} + +__global__ void redwoodNoiseModelKernel(const float* __restrict__ depth, + const int H, + const int W, + curandState_t* states, + const float* __restrict__ model, + const float noiseMultiplier, + float* __restrict__ noisyDepth) { + const int TID = threadIdx.x; + const int BID = blockIdx.x; + + // curandStates are thread-safe, so all threads in a block share the same + // state. They are NOT block safe however + curandState_t curandState = states[BID]; + + const float ymax = H - 1; + const float xmax = W - 1; + + for (int j = BID; j < H; j += gridDim.x) { + for (int i = TID; i < W; i += blockDim.x) { + // Shuffle pixels + const int y = + min(max(j + curand_normal(&curandState) * 0.25f * noiseMultiplier, + 0.0f), + ymax) + + 0.5f; + const int x = + min(max(i + curand_normal(&curandState) * 0.25f * noiseMultiplier, + 0.0f), + xmax) + + 0.5f; + + // downsample + const float d = depth[(y - y % 2) * W + x - x % 2]; + // If depth is greater than 10m, the sensor will just return a zero + if (d >= 10.0f) { + noisyDepth[j * W + i] = 0.0f; + } else { + // Distortion + // The noise model was originally made for a 640x480 sensor, + // so re-map our arbitrarily sized sensor to that size! + const float undistorted_d = + undistort(x / xmax * 639.0f, y / ymax * 479.0f, d, model); + + // quantization and high freq noise + if (undistorted_d == 0.0f) + noisyDepth[j * W + i] = 0.0f; + else { + const float denom = (round(35.130f / undistorted_d + + curand_normal(&curandState) * 0.027778f * + noiseMultiplier) * + 8.0f); + noisyDepth[j * W + i] = denom > 1e-5 ? (35.130f * 8.0f / denom) : 0.0; + } + } + } + } +} + +__global__ void curandStatesSetupKernel(curandState_t* state, int seed, int n) { + int id = threadIdx.x + blockIdx.x * 64; + if (id < n) { + curand_init(seed, id + 1, 0, &state[id]); + } +} + +} // namespace + +namespace esp { +namespace sensor { +namespace impl { + +struct CurandStates { + CurandStates() : devStates(0), n_blocks_(0) {} + void alloc(const int n_blocks) { + if (n_blocks > n_blocks_) { + release(); + cudaMalloc(&devStates, n_blocks * sizeof(curandState_t)); + curandStatesSetupKernel<<>>( + devStates, rand(), n_blocks); + n_blocks_ = n_blocks; + } + } + + void release() { + if (devStates != 0) { + cudaFree(devStates); + devStates = 0; + } + } + + ~CurandStates() { release(); } + + curandState_t* devStates; + + private: + int n_blocks_; +}; + +CurandStates* getCurandStates() { + return new CurandStates(); +} +void freeCurandStates(CurandStates* curandStates) { + if (curandStates != 0) + delete curandStates; +} + +void simulateFromGPU(const float* __restrict__ devDepth, + const int H, + const int W, + const float* __restrict__ devModel, + CurandStates* curandStates, + const float noiseMultiplier, + float* __restrict__ devNoisyDepth) { + const int n_threads = std::min(std::max(W / 4, 1), 256); + const int n_blocks = std::max(H / 8, 1); + + curandStates->alloc(n_blocks); + redwoodNoiseModelKernel<<>>( + devDepth, H, W, curandStates->devStates, devModel, noiseMultiplier, + devNoisyDepth); +} + +void simulateFromCPU(const float* __restrict__ depth, + const int H, + const int W, + const float* __restrict__ devModel, + CurandStates* curandStates, + const float noiseMultiplier, + float* __restrict__ noisyDepth) { + float *devDepth, *devNoisyDepth; + cudaMalloc(&devDepth, H * W * sizeof(float)); + cudaMalloc(&devNoisyDepth, H * W * sizeof(float)); + + cudaMemcpy(devDepth, depth, H * W * sizeof(float), cudaMemcpyHostToDevice); + + simulateFromGPU(devDepth, H, W, devModel, curandStates, noiseMultiplier, + devNoisyDepth); + + cudaMemcpy(noisyDepth, devNoisyDepth, H * W * sizeof(float), + cudaMemcpyDeviceToHost); + + cudaFree(devNoisyDepth); + cudaFree(devDepth); +} +} // namespace impl +} // namespace sensor +} // namespace esp diff --git a/src/esp/sensor/RedwoodNoiseModel.cuh b/src/esp/sensor/RedwoodNoiseModel.cuh new file mode 100644 index 0000000000..51e1140f4c --- /dev/null +++ b/src/esp/sensor/RedwoodNoiseModel.cuh @@ -0,0 +1,34 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +namespace esp { +namespace sensor { +namespace impl { + +struct CurandStates; + +CurandStates* getCurandStates(); + +void freeCurandStates(CurandStates* curandStates); + +void simulateFromCPU(const float* __restrict__ depth, + const int H, + const int W, + const float* __restrict__ devModel, + CurandStates* curandStates, + const float noiseMultiplier, + float* __restrict__ noisyDepth); + +void simulateFromGPU(const float* __restrict__ devDepth, + const int H, + const int W, + const float* __restrict__ devModel, + CurandStates* curandStates, + const float noiseMultiplier, + float* __restrict__ devNoisyDepth); +} // namespace impl +} // namespace sensor +} // namespace esp diff --git a/src/esp/sensor/RedwoodNoiseModel.h b/src/esp/sensor/RedwoodNoiseModel.h new file mode 100644 index 0000000000..e0c47b0158 --- /dev/null +++ b/src/esp/sensor/RedwoodNoiseModel.h @@ -0,0 +1,87 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. +// +#pragma once + +#include "esp/core/esp.h" +#include "esp/core/random.h" + +#include "RedwoodNoiseModel.cuh" + +namespace esp { +namespace sensor { + +/** + * Provides a CUDA/GPU implementation of the Redwood Noise Model for PrimSense + Depth sensors + * provided here: http://redwood-data.org/indoor/dataset.html + * + * Please cite the following work if you use this noise model + * @verbatim +@inproceedings{choi2015robust, + title={Robust reconstruction of indoor scenes}, + author={Choi, Sungjoon and Zhou, Qian-Yi and Koltun, Vladlen}, + booktitle={Proceedings of the IEEE Conference on Computer Vision and + Pattern Recognition}, pages={5556--5565}, year={2015} +} + @endverbatim + */ +struct RedwoodNoiseModelGPUImpl { + /** + * @brief Constructor + * @param model The distortion model from + * http://redwood-data.org/indoor/data/dist-model.txt + * The 3rd dimension is assumed to have been + * flattened into the second + * @param gpuDeviceId The CUDA device ID to use + * @param noiseMultiplier Multiplier for the Gaussian random-variables. This + * can be used to increase or decrease the noise + * level + */ + RedwoodNoiseModelGPUImpl(const Eigen::Ref model, + const int gpuDeviceId, + const float noiseMultiplier); + + /** + * @brief Simulates noisy depth from clean depth. The input is assumed to be + * on the CPU and the output will be on the CPU. If the input isn't on the + * CPU, bad things happen, segfaults happen. + * + * @param[in] depth Clean depth, i.e. depth from habitat's depth shader + * @return Simulated noisy depth + */ + Eigen::RowMatrixXf simulateFromCPU( + const Eigen::Ref depth); + + /** + * @brief Similar to @ref simulateFromCPU() but the input and output are + * assumed to be on the GPU. If they aren't, bad things happen, segfaults + * happen. + * + * @param[in] devDepth Device pointer to the clean depth + * Assumed to be a continguous array in row-major + * order + * @param[in] rows The number of rows in the depth image + * @param[in] cols The number of columns + * @param[out] devNoisyDepth Device pointer to the memory to write the noisy + * depth + */ + void simulateFromGPU(const float* devDepth, + const int rows, + const int cols, + float* devNoisyDepth); + + ~RedwoodNoiseModelGPUImpl(); + + private: + const int gpuDeviceId_; + const float noiseMultiplier_; + float* devModel_ = nullptr; + impl::CurandStates* curandStates_ = nullptr; + + ESP_SMART_POINTERS(RedwoodNoiseModelGPUImpl) +}; + +} // namespace sensor +} // namespace esp diff --git a/src/esp/sensor/Sensor.cpp b/src/esp/sensor/Sensor.cpp index a397dfb912..3724fa1088 100644 --- a/src/esp/sensor/Sensor.cpp +++ b/src/esp/sensor/Sensor.cpp @@ -68,7 +68,7 @@ bool operator==(const SensorSpec& a, const SensorSpec& b) { a.position == b.position && a.orientation == b.orientation && a.resolution == b.resolution && a.channels == b.channels && a.encoding == b.encoding && a.observationSpace == b.observationSpace && - a.gpu2gpuTransfer == b.gpu2gpuTransfer; + a.noiseModel == b.noiseModel && a.gpu2gpuTransfer == b.gpu2gpuTransfer; } bool operator!=(const SensorSpec& a, const SensorSpec& b) { return !(a == b); diff --git a/src/esp/sensor/Sensor.h b/src/esp/sensor/Sensor.h index 8dcdb0602e..749bb61df5 100644 --- a/src/esp/sensor/Sensor.h +++ b/src/esp/sensor/Sensor.h @@ -55,9 +55,11 @@ struct SensorSpec { std::string encoding = "rgba_uint8"; // description of Sensor observation space as gym.spaces.Dict() std::string observationSpace = ""; + std::string noiseModel = "None"; bool gpu2gpuTransfer = false; ESP_SMART_POINTERS(SensorSpec) }; + bool operator==(const SensorSpec& a, const SensorSpec& b); bool operator!=(const SensorSpec& a, const SensorSpec& b); diff --git a/tests/test_sensors.py b/tests/test_sensors.py index ab9a494660..f8fc4cd8c4 100644 --- a/tests/test_sensors.py +++ b/tests/test_sensors.py @@ -17,6 +17,45 @@ from examples.settings import make_cfg from habitat_sim.utils.common import quat_from_coeffs + +def _render_and_load_gt(sim, scene, sensor_type, gpu2gpu): + gt_data_pose_file = osp.abspath( + osp.join( + osp.dirname(__file__), + "gt_data", + "{}-state.json".format(osp.basename(osp.splitext(scene)[0])), + ) + ) + with open(gt_data_pose_file, "r") as f: + render_state = json.load(f) + state = habitat_sim.AgentState() + state.position = render_state["pos"] + state.rotation = quat_from_coeffs(render_state["rot"]) + + sim.initialize_agent(0, state) + obs = sim.step("move_forward") + + assert sensor_type in obs, f"{sensor_type} not in obs" + + gt_obs_file = osp.abspath( + osp.join( + osp.dirname(__file__), + "gt_data", + "{}-{}.npy".format(osp.basename(osp.splitext(scene)[0]), sensor_type), + ) + ) + gt = np.load(gt_obs_file) + + if gpu2gpu: + import torch + + for k, v in obs.items(): + if torch.is_tensor(v): + obs[k] = v.cpu().numpy() + + return obs, gt + + _test_scenes = [ osp.abspath( osp.join( @@ -79,39 +118,7 @@ def test_sensors(scene, has_sem, sensor_type, gpu2gpu, sim, make_cfg_settings): sim.reconfigure(cfg) - gt_data_pose_file = osp.abspath( - osp.join( - osp.dirname(__file__), - "gt_data", - "{}-state.json".format(osp.basename(osp.splitext(scene)[0])), - ) - ) - with open(gt_data_pose_file, "r") as f: - render_state = json.load(f) - state = habitat_sim.AgentState() - state.position = render_state["pos"] - state.rotation = quat_from_coeffs(render_state["rot"]) - - sim.initialize_agent(0, state) - obs = sim.step("move_forward") - - assert sensor_type in obs, f"{sensor_type} not in obs" - - gt_obs_file = osp.abspath( - osp.join( - osp.dirname(__file__), - "gt_data", - "{}-{}.npy".format(osp.basename(osp.splitext(scene)[0]), sensor_type), - ) - ) - gt = np.load(gt_obs_file) - - if gpu2gpu: - import torch - - for k, v in obs.items(): - if torch.is_tensor(v): - obs[k] = v.cpu().numpy() + obs, gt = _render_and_load_gt(sim, scene, sensor_type, gpu2gpu) # Different GPUs and different driver version will produce slightly different images assert np.linalg.norm( @@ -134,3 +141,33 @@ def test_smoke_no_sensors(make_cfg_settings): cfg = make_cfg(make_cfg_settings) cfg.agents[0].sensor_specifications = [] sims.append(habitat_sim.Simulator(cfg)) + + +@pytest.mark.gfxtest +@pytest.mark.parametrize( + "scene,gpu2gpu", itertools.product(_test_scenes, [True, False]) +) +def test_smoke_redwood_noise(scene, gpu2gpu, sim, make_cfg_settings): + if not osp.exists(scene): + pytest.skip("Skipping {}".format(scene)) + + if not habitat_sim.cuda_enabled and gpu2gpu: + pytest.skip("Skipping GPU->GPU test") + + make_cfg_settings = {k: v for k, v in make_cfg_settings.items()} + make_cfg_settings["depth_sensor"] = True + make_cfg_settings["color_sensor"] = False + make_cfg_settings["semantic_sensor"] = False + make_cfg_settings["scene"] = scene + hsim_cfg = make_cfg(make_cfg_settings) + hsim_cfg.agents[0].sensor_specifications[0].noise_model = "RedwoodDepthNoiseModel" + for sensor_spec in hsim_cfg.agents[0].sensor_specifications: + sensor_spec.gpu2gpu_transfer = gpu2gpu + + sim.reconfigure(hsim_cfg) + + obs, gt = _render_and_load_gt(sim, scene, "depth_sensor", gpu2gpu) + + assert np.linalg.norm( + obs["depth_sensor"].astype(np.float) - gt.astype(np.float) + ) > 1.5e-2 * np.linalg.norm(gt.astype(np.float)), f"Incorrect {sensor_type} output"