origin

scripts/ros/depth_value_pub.py

@@ -0,0 +1,54 @@
+""" Node to publish max depth value when gel is deformed """
+import os
+import hydra
+import rospy
+from std_msgs.msg import Float32
+
+from digit_depth.third_party import geom_utils
+from digit_depth.digit.digit_sensor import DigitSensor
+from digit_depth.train.mlp_model import MLP
+from digit_depth.train.prepost_mlp import *
+seed = 42
+torch.seed = seed
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
+
+@hydra.main(config_path=f"{BASE_PATH}/config", config_name="rgb_to_normal.yaml", version_base=None)
+def print_depth(cfg):
+    model = torch.load(cfg.model_path).to(device)
+    model.eval()
+    # setup digit sensor
+    digit = DigitSensor(30, "QVGA", cfg.sensor.serial_num)
+    digit_call = digit()
+    pub = rospy.Publisher('chatter', Float32, queue_size=1)
+    rospy.init_node('depth', anonymous=True)
+    try:
+        while not rospy.is_shutdown():
+            frame = digit_call.get_frame()
+            img_np = preproc_mlp(frame)
+            img_np = model(img_np).detach().cpu().numpy()
+            img_np, normal_img = post_proc_mlp(img_np)
+
+            # get gradx and grady
+            gradx_img, grady_img = geom_utils._normal_to_grad_depth(img_normal=img_np, gel_width=cfg.sensor.gel_width,
+                                                                    gel_height=cfg.sensor.gel_height,bg_mask=None)
+            # reconstruct depth
+            img_depth = geom_utils._integrate_grad_depth(gradx_img, grady_img, boundary=None, bg_mask=None,max_depth=0.0237)
+            img_depth = img_depth.detach().cpu().numpy().flatten()
+
+            # get max depth value
+            max_depth = np.min(img_depth)
+            rospy.loginfo(f"max:{max_depth}")
+            img_depth_calibrated = np.abs((max_depth - 0.02362))
+
+            # publish max depth value
+            pub.publish(Float32(img_depth_calibrated*10000))  # convert to mm
+
+    except KeyboardInterrupt:
+        print("Shutting down")
+        digit().disconnect()
+
+
+if __name__ == "__main__":
+    rospy.loginfo("starting...")
+    print_depth()

scripts/ros/digit_image_pub.py

@@ -0,0 +1,52 @@
+""" ROS image publisher for DIGIT sensor """
+
+import argparse
+# OpenCV
+import cv2
+from PIL import Image as Im
+from sensor_msgs.msg import Image
+from cv_bridge import CvBridge
+
+# Ros libraries
+import roslib
+import rospy
+
+# Ros Messages
+from sensor_msgs.msg import CompressedImage
+from sensor_msgs.msg import std_msgs
+from digit_depth.digit.digit_sensor import DigitSensor
+
+
+class ImageFeature:
+
+    def __init__(self):
+        # topic where we publish
+
+        self.image_pub = rospy.Publisher("/output/image_raw/compressed",
+                                         CompressedImage, queue_size=10)
+        self.br = CvBridge()
+
+
+def rgb_pub(digit_sensor: DigitSensor):
+    # Initializes and cleanup ros node
+    ic = ImageFeature()
+    rospy.init_node('image_feature', anonymous=True)
+    digit_call = digit_sensor()
+    br = CvBridge()
+    while True:
+        frame = digit_call.get_frame()
+        msg = br.cv2_to_compressed_imgmsg(frame, "png")
+        ic.image_pub.publish(msg)
+        rospy.loginfo("published ...")
+        if cv2.waitKey(1) == 27:
+            break
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--fps", type=int, default=30)
+    argparser.add_argument("--resolution", type=str, default="QVGA")
+    argparser.add_argument("--serial_num", type=str, default="D00001")
+    args, unknown = argparser.parse_known_args()
+    digit = DigitSensor(args.fps, args.resolution, args.serial_num)
+    rgb_pub(digit)