Rgb noise models

habitat_sim/sensors/noise_models/__init__.py

@@ -7,11 +7,17 @@
 from typing import Any, Dict
 
 from habitat_sim.registry import registry
+from habitat_sim.sensors.noise_models.gaussian_noise_model import GaussianNoiseModel
 from habitat_sim.sensors.noise_models.no_noise_model import NoSensorNoiseModel
+from habitat_sim.sensors.noise_models.poisson_noise_model import PoissonNoiseModel
 from habitat_sim.sensors.noise_models.redwood_depth_noise_model import (
     RedwoodDepthNoiseModel,
 )
+from habitat_sim.sensors.noise_models.salt_and_pepper_noise_model import (
+    SaltAndPepperNoiseModel,
+)
 from habitat_sim.sensors.noise_models.sensor_noise_model import SensorNoiseModel
+from habitat_sim.sensors.noise_models.speckle_noise_model import SpeckleNoiseModel
 
 
 def make_sensor_noise_model(name: str, kwargs: Dict[str, Any]) -> SensorNoiseModel:
@@ -32,4 +38,8 @@ def make_sensor_noise_model(name: str, kwargs: Dict[str, Any]) -> SensorNoiseMod
     "SensorNoiseModel",
     "RedwoodDepthNoiseModel",
     "NoSensorNoiseModel",
+    "GaussianNoiseModel",
+    "SaltAndPepperNoiseModel",
+    "PoissonNoiseModel",
+    "SpeckleNoiseModel",
 ]

habitat_sim/sensors/noise_models/gaussian_noise_model.py

@@ -0,0 +1,58 @@
+#!/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 attr
+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
+
+
+def _simulate(image, intensity_constant, mean, sigma):
+    image = image / 255.0
+
+    noise = np.random.normal(mean, sigma, image.shape)
+    noisy_rgb = np.clip((image + (noise * intensity_constant)), 0, 1)
+
+    noisy_rgb = (noisy_rgb * 255).astype(np.uint8)
+
+    return noisy_rgb
+
+
+@attr.s(auto_attribs=True)
+class GaussianNoiseModelCPUImpl:
+    intensity_constant: float
+    mean: int
+    sigma: int
+
+    def simulate(self, image):
+        return _simulate(image, self.intensity_constant, self.mean, self.sigma)
+
+
+@registry.register_noise_model
+@attr.s(auto_attribs=True, kw_only=True)
+class GaussianNoiseModel(SensorNoiseModel):
+    intensity_constant: float = 0.2
+    mean: int = 0
+    sigma: int = 1
+
+    def __attrs_post_init__(self):
+        self._impl = GaussianNoiseModelCPUImpl(
+            self.intensity_constant, self.mean, self.sigma
+        )
+
+    @staticmethod
+    def is_valid_sensor_type(sensor_type: SensorType) -> bool:
+        return sensor_type == SensorType.COLOR
+
+    def simulate(self, image):
+        return self._impl.simulate(image)
+
+    def apply(self, image):
+        r"""Alias of `simulate()` to conform to base-class and expected API
+        """
+        return self.simulate(image)

habitat_sim/sensors/noise_models/poisson_noise_model.py

@@ -0,0 +1,52 @@
+#!/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 attr
+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
+
+
+def _simulate(image):
+
+    image = image / 255.0
+
+    values = len(np.unique(image))
+    values = 2 ** np.ceil(np.log2(values))
+
+    noisy = np.random.poisson(image * values) / float(values)
+    noisy_rgb = np.clip(noisy, 0, 1)
+    noisy_rgb = (noisy_rgb * 255).astype(np.uint8)
+
+    return noisy_rgb
+
+
+@attr.s(auto_attribs=True)
+class PoissonNoiseModelCPUImpl:
+    @staticmethod
+    def simulate(image):
+        return _simulate(image)
+
+
+@registry.register_noise_model
+@attr.s(auto_attribs=True, kw_only=True)
+class PoissonNoiseModel(SensorNoiseModel):
+    def __attrs_post_init__(self):
+        self._impl = PoissonNoiseModelCPUImpl()
+
+    @staticmethod
+    def is_valid_sensor_type(sensor_type: SensorType) -> bool:
+        return sensor_type == SensorType.COLOR
+
+    def simulate(self, image):
+        return self._impl.simulate(image)
+
+    def apply(self, image):
+        r"""Alias of `simulate()` to conform to base-class and expected API
+        """
+        return self.simulate(image)

habitat_sim/sensors/noise_models/salt_and_pepper_noise_model.py

@@ -0,0 +1,59 @@
+#!/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 attr
+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
+
+
+def _simulate(image, s_vs_p, amount):
+    noisy_rgb = np.copy(image)
+    # Salt
+    num_salt = np.ceil(amount * image.size * s_vs_p)
+    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape]
+    noisy_rgb[tuple(coords)] = 1
+
+    # Pepper
+    num_pepper = np.ceil(amount * image.size * (1.0 - s_vs_p))
+    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape]
+    noisy_rgb[tuple(coords)] = 0
+
+    return noisy_rgb
+
+
+@attr.s(auto_attribs=True)
+class SaltAndPepperNoiseModelCPUImpl:
+    s_vs_p: float
+    amount: float
+
+    def simulate(self, image):
+        return _simulate(image, self.s_vs_p, self.amount)
+
+
+@registry.register_noise_model
+@attr.s(auto_attribs=True, kw_only=True)
+class SaltAndPepperNoiseModel(SensorNoiseModel):
+    s_vs_p: float = 0.5
+    amount: float = 0.05
+
+    def __attrs_post_init__(self):
+        self._impl = SaltAndPepperNoiseModelCPUImpl(self.s_vs_p, self.amount)
+
+    @staticmethod
+    def is_valid_sensor_type(sensor_type: SensorType) -> bool:
+        return sensor_type == SensorType.COLOR
+
+    def simulate(self, image):
+        return self._impl.simulate(image)
+
+    def apply(self, image):
+        r"""Alias of `simulate()` to conform to base-class and expected API
+        """
+        return self.simulate(image)

habitat_sim/sensors/noise_models/speckle_noise_model.py

@@ -0,0 +1,58 @@
+#!/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 attr
+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
+
+
+def _simulate(image, intensity_constant, mean, sigma):
+    image = image / 255.0
+
+    noise = np.random.normal(mean, sigma, image.shape)
+    noisy_rgb = np.clip((image + (image * noise * intensity_constant)), 0, 1)
+
+    noisy_rgb = (noisy_rgb * 255).astype(np.uint8)
+
+    return noisy_rgb
+
+
+@attr.s(auto_attribs=True)
+class SpeckleNoiseModelCPUImpl:
+    intensity_constant: float
+    mean: int
+    sigma: int
+
+    def simulate(self, image):
+        return _simulate(image, self.intensity_constant, self.mean, self.sigma)
+
+
+@registry.register_noise_model
+@attr.s(auto_attribs=True, kw_only=True)
+class SpeckleNoiseModel(SensorNoiseModel):
+    intensity_constant: float = 0.2
+    mean: int = 0
+    sigma: int = 1
+
+    def __attrs_post_init__(self):
+        self._impl = SpeckleNoiseModelCPUImpl(
+            self.intensity_constant, self.mean, self.sigma
+        )
+
+    @staticmethod
+    def is_valid_sensor_type(sensor_type: SensorType) -> bool:
+        return sensor_type == SensorType.COLOR
+
+    def simulate(self, image):
+        return self._impl.simulate(image)
+
+    def apply(self, image):
+        r"""Alias of `simulate()` to conform to base-class and expected API
+        """
+        return self.simulate(image)

tests/test_sensors.py

@@ -171,3 +171,35 @@ def test_smoke_redwood_noise(scene, gpu2gpu, sim, make_cfg_settings):
     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"
+
+
+@pytest.mark.gfxtest
+@pytest.mark.parametrize("scene", _test_scenes)
+@pytest.mark.parametrize(
+    "model_name",
+    [
+        "SpeckleNoiseModel",
+        "GaussianNoiseModel",
+        "SaltAndPepperNoiseModel",
+        "PoissonNoiseModel",
+    ],
+)
+def test_rgb_noise(scene, model_name, sim, make_cfg_settings):
+    if not osp.exists(scene):
+        pytest.skip("Skipping {}".format(scene))
+
+    make_cfg_settings = {k: v for k, v in make_cfg_settings.items()}
+    make_cfg_settings["depth_sensor"] = False
+    make_cfg_settings["color_sensor"] = True
+    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 = model_name
+
+    sim.reconfigure(hsim_cfg)
+
+    obs, gt = _render_and_load_gt(sim, scene, "color_sensor", False)
+
+    assert np.linalg.norm(
+        obs["color_sensor"].astype(np.float) - gt.astype(np.float)
+    ) > 1.5e-2 * np.linalg.norm(gt.astype(np.float)), f"Incorrect {sensor_type} output"