From 764288d3ac1acc943df39d4f236a4b3713d69d9d Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Thu, 26 Aug 2021 13:37:09 +0200
Subject: [PATCH] Apply noise to the synthetic depth images

---
 active_grasp/__init__.py |  2 ++
 scripts/bt_sim_node.py   | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/active_grasp/__init__.py b/active_grasp/__init__.py
index f2a2e42..3e66aed 100644
--- a/active_grasp/__init__.py
+++ b/active_grasp/__init__.py
@@ -1,7 +1,9 @@
 from .policy import register
 from .baselines import *
+from .nbv import NextBestView
 
 register("initial-view", InitialView)
 register("top-view", TopView)
 register("top-trajectory", TopTrajectory)
 register("circular-trajectory", CircularTrajectory)
+register("nbv", NextBestView)
diff --git a/scripts/bt_sim_node.py b/scripts/bt_sim_node.py
index bf51b5a..93349e2 100755
--- a/scripts/bt_sim_node.py
+++ b/scripts/bt_sim_node.py
@@ -7,8 +7,8 @@ from geometry_msgs.msg import PoseStamped
 import numpy as np
 import rospy
 from sensor_msgs.msg import JointState, Image, CameraInfo
+import skimage.transform
 from threading import Thread
-import tf2_ros
 
 from active_grasp.bbox import to_bbox_msg
 from active_grasp.srv import *
@@ -199,11 +199,22 @@ class CameraPlugin(Plugin):
         self.info_pub.publish(msg)
 
         _, depth, _ = self.camera.get_image()
+        depth = apply_noise(depth)
+
         msg = self.cv_bridge.cv2_to_imgmsg(depth)
         msg.header.stamp = stamp
         self.depth_pub.publish(msg)
 
 
+def apply_noise(img, k=1000, theta=0.001, sigma=0.005, l=4.0):
+    # Multiplicative and additive noise
+    img *= np.random.gamma(k, theta)
+    h, w = img.shape
+    noise = np.random.randn(int(h / l), int(w / l)) * sigma
+    img += skimage.transform.resize(noise, img.shape, order=1, mode="constant")
+    return img
+
+
 def main():
     rospy.init_node("bt_sim")
     server = BtSimNode()