merged bayesian stuff

Author: gerni17

.gitignore

@@ -6,3 +6,4 @@
 .run/
 docs/build 
 _build
+

elevation_mapping_cupy/script/elevation_mapping_cupy/custom_kernels.py

@@ -694,6 +694,50 @@ def sum_kernel(
     )
     return sum_kernel
 
+def alpha_kernel(
+        resolution,
+        width,
+        height,
+    ):
+    # input the list of layers, amount of channels can slo be input through kernel
+    alpha_kernel = cp.ElementwiseKernel(
+        in_params="raw U p, raw W pcl_chan, raw W map_lay, raw W pcl_channels",
+        out_params="raw U newmap",
+        preamble=string.Template(
+            """
+                __device__ int get_map_idx(int idx, int layer_n) {
+                    const int layer = ${width} * ${height};
+                    return layer * layer_n + idx;
+                }
+            """
+        ).substitute(resolution=resolution,width=width, height=height),
+        operation=string.Template(
+            """
+            U idx = p[i * pcl_channels[0]];
+            U valid = p[i * pcl_channels[0] + 1];
+            U inside = p[i * pcl_channels[0] + 2];
+            if (valid) {
+                if (inside) {
+                    U theta_max = 0;
+                    W arg_max = 0;
+                    for ( W it=0;it<pcl_channels[1];it++){
+                    U theta = p[i * pcl_channels[0] + pcl_chan[it]];
+                        if (theta >=theta_max){
+                            arg_max = map_lay[it];
+                            theta_max = theta;
+                        }
+                    }
+                    atomicAdd(&newmap[get_map_idx(idx, arg_max)], 1.0);
+                }
+            }
+            """
+        ).substitute(
+        ),
+        name="alpha_kernel",
+    )
+    return alpha_kernel
+
+
 
 def average_kernel(
         width,
@@ -875,7 +919,6 @@ def color_average_kernel(
     )
     return color_average_kernel
 
-# TODO: in future add more kernels e.g. bayesian kernel
 
 if __name__ == "__main__":
     for i in range(10):

elevation_mapping_cupy/script/elevation_mapping_cupy/plugins/semantic_filter.py

@@ -24,22 +24,35 @@ def __init__(
         self,
         cell_n: int = 100,
         classes: list = ["person", "grass"],
-        colors: list = [
-            [0, 255, 0],
-            [0, 255, 255],
-            [0, 0, 255],
-            [255, 255, 0],
-            [255, 0, 0],
-        ],
         **kwargs,
     ):
         super().__init__()
         self.indices = []
         self.classes = classes
-        color_classes = np.array(colors)
-        self.color_encoding = self.transform_color(color_classes)
+        self.color_encoding = self.transform_color()
 
-    def transform_color(self, color_classes):
+    def color_map(self, N=256, normalized=False):
+        def bitget(byteval, idx):
+            return (byteval & (1 << idx)) != 0
+
+        dtype = "float32" if normalized else "uint8"
+        cmap = np.zeros((N+1, 3), dtype=dtype)
+        for i in range(N+1):
+            r = g = b = 0
+            c = i
+            for j in range(8):
+                r = r | (bitget(c, 0) << 7 - j)
+                g = g | (bitget(c, 1) << 7 - j)
+                b = b | (bitget(c, 2) << 7 - j)
+                c = c >> 3
+
+            cmap[i] = np.array([r, g, b])
+
+        cmap = cmap / 255 if normalized else cmap
+        return cmap[1:]
+
+    def transform_color(self):
+        color_classes = self.color_map(len(self.classes))
         r = np.asarray(color_classes[:, 0], dtype=np.uint32)
         g = np.asarray(color_classes[:, 1], dtype=np.uint32)
         b = np.asarray(color_classes[:, 2], dtype=np.uint32)
@@ -59,9 +72,7 @@ def __call__(
         # get indices of all layers that
         layer_indices = cp.array([], dtype=cp.int32)
         for it, fusion_alg in enumerate(semantic_map.param.fusion_algorithms):
-            if fusion_alg == "class_average" and (
-                semantic_map.param.additional_layers[it] in self.classes
-            ):
+            if (fusion_alg in ["class_bayesian","class_average"]):
                 layer_indices = cp.append(layer_indices, it).astype(cp.int32)
 
         # check which has the highest value

elevation_mapping_cupy/script/elevation_mapping_cupy/semantic_map.py

@@ -8,6 +8,7 @@
     color_average_kernel,
     average_kernel,
     class_average_kernel,
+    alpha_kernel,
 )
 
 xp = cp
@@ -74,7 +75,7 @@ def compile_kernels(self) -> None:
                 self.param.cell_n, self.param.cell_n
             )
         if "class_average" in self.unique_fusion:
-            print("Initialize fusion kernel")
+            print("Initialize class average kernel")
             self.sum_kernel = sum_kernel(
                 self.param.resolution,
                 self.param.cell_n,
@@ -84,6 +85,13 @@ def compile_kernels(self) -> None:
                 self.param.cell_n,
                 self.param.cell_n,
             )
+        if "class_bayesian" in self.unique_fusion:
+            print("Initialize class bayesian kernel")
+            self.alpha_kernel = alpha_kernel(
+                self.param.resolution,
+                self.param.cell_n,
+                self.param.cell_n,
+            )
 
     def get_fusion_of_pcl(self, channels: List[str]) -> List[str]:
         """Get all fusion algorithms that need to be applied to a specific pointcloud
@@ -99,7 +107,7 @@ def get_fusion_of_pcl(self, channels: List[str]) -> List[str]:
         return fusion_list
 
     def get_indices_fusion(self, pcl_channels: List[str], fusion_alg: str):
-        """Computes the indices that are used for the additional kernel update of the pcl and the elmap.
+        """Computes the indices that are used for the additional kernel update of the popintcloud and the semantic map.
 
         Args:
             pcl_channels (List[str]):
@@ -172,6 +180,25 @@ def update_layers(self, points_all, channels, R, t, elevation_map):
                 size=(self.param.cell_n * self.param.cell_n),
             )
 
+        if "class_bayesian" in additional_fusion:
+            pcl_ids, layer_ids = self.get_indices_fusion(channels, "class_bayesian")
+            # alpha sum get points as input and calculate for each point to what cell it belongs and then
+            # adds to the right channel a one
+            self.alpha_kernel(
+                points_all,
+                pcl_ids,
+                layer_ids,
+                cp.array([points_all.shape[1], pcl_ids.shape[0]], dtype=np.int32),
+                self.new_map,
+                size=(points_all.shape[0]),
+            )
+            # calculate new thetas
+            sum_alpha = cp.sum(self.new_map[layer_ids], axis=0)
+            self.map[layer_ids] = self.new_map[layer_ids] / cp.expand_dims(
+                sum_alpha, axis=0
+            )
+            # assert  cp.unique(cp.sum(self.map[layer_ids], axis=0)) equal to zero or to nan
+
         if "color" in additional_fusion:
             pcl_ids, layer_ids = self.get_indices_fusion(channels, "color")
             if self.color_map is None:

elevation_mapping_cupy/script/elevation_mapping_cupy/tests/test_elevation_mapping.py

@@ -25,6 +25,8 @@ def elmap_ex(add_lay, fusion_alg):
     [
         (["feat_0", "feat_1", "rgb"], ["average", "average", "color"]),
         (["feat_0", "feat_1"], ["average", "average"]),
+        (["feat_0", "feat_1"], ["class_average", "class_average"]),
+        (["feat_0", "feat_1"], ["class_bayesian", "class_bayesian"]),
     ],
 )
 class TestElevationMap:

elevation_mapping_cupy/script/elevation_mapping_cupy/tests/test_semantic_map.py

@@ -26,6 +26,9 @@ def semmap_ex(add_lay, fusion_alg):
     [(["feat_0", "feat_1"], ["average", "average"],["feat_0"]),
     (["feat_0", "feat_1"], ["average", "average"],[]),
     (["feat_0", "feat_1", "rgb"], ["average", "average", "color"],["rgb", "feat_0"]),
+    (["feat_0", "feat_1", "rgb"], ["class_average", "average", "color"], ["rgb", "feat_0"]),
+    (["feat_0", "feat_1", "rgb"], ["class_bayesian", "average", "color"], ["rgb", "feat_0"]),
+     (["feat_0", "feat_1", "rgb"], ["class_bayesian", "average", "color"], ["rgb", "feat_0", "feat_1"]),
      ],
 )
 def test_fusion_of_pcl(semmap_ex, channels):
@@ -42,4 +45,4 @@ def test_fusion_of_pcl(semmap_ex, channels):
 )
 @pytest.mark.parametrize("channels", [["rgb"], ["rgb", "feat_0"], []])
 def test_indices_fusion(semmap_ex, channels,fusion_alg):
-    pcl_indices, layer_indices = semmap_ex.get_indices_fusion(pcl_channels=channels,fusion_alg=fusion_alg[0])
+    pcl_indices, layer_indices = semmap_ex.get_indices_fusion(pcl_channels=channels, fusion_alg=fusion_alg[0])

requirements.txt

@@ -4,7 +4,8 @@ scipy
 dataclasses
 ruamel.yaml
 opencv-python
-simple-parsing
+simple_parsing
+scikit-image
 
 ###### Requirements with Version Specifiers ######`
 shapely==1.7.1

sensor_processing/semantic_pointcloud/script/semantic_pointcloud/pointcloud_node.py

@@ -4,10 +4,12 @@
 import numpy as np
 import cupy as cp
 import ros_numpy
+import matplotlib.pyplot as plt
+from skimage.io import imshow
 
 import message_filters
 from sensor_msgs.msg import Image, CameraInfo
-from sensor_msgs.msg import PointCloud2
+from sensor_msgs.msg import PointCloud2, Image
 from cv_bridge import CvBridge
 
 from semantic_pointcloud.pointcloud_parameters import PointcloudParameter
@@ -67,6 +69,52 @@ def __init__(self, sensor_name):
         ts.registerCallback(self.image_callback)
 
         self.pcl_pub = rospy.Publisher(self.param.topic_name, PointCloud2, queue_size=2)
+        if self.param.publish_segmentation_image:
+            self.seg_pub = rospy.Publisher(
+                self.param.segmentation_image_topic, Image, queue_size=2
+            )
+        self.semseg_color_map = self.color_map(
+            len(list(self.segmentation_channels.keys()))
+        )
+        self.color_map_viz()
+
+    def color_map(self, N=256, normalized=False):
+        def bitget(byteval, idx):
+            return (byteval & (1 << idx)) != 0
+
+        dtype = "float32" if normalized else "uint8"
+        cmap = np.zeros((N + 1, 3), dtype=dtype)
+        for i in range(N + 1):
+            r = g = b = 0
+            c = i
+            for j in range(8):
+                r = r | (bitget(c, 0) << 7 - j)
+                g = g | (bitget(c, 1) << 7 - j)
+                b = b | (bitget(c, 2) << 7 - j)
+                c = c >> 3
+
+            cmap[i] = np.array([r, g, b])
+
+        cmap = cmap / 255 if normalized else cmap
+        return cmap[1:]
+
+    def color_map_viz(self):
+        # labels = self.semantic_model["model"].get_classes()
+        labels = list(self.segmentation_channels.keys())
+        nclasses = len(labels)
+        row_size = 50
+        col_size = 500
+        cmap = self.semseg_color_map
+        array = np.empty(
+            (row_size * (nclasses), col_size, cmap.shape[1]), dtype=cmap.dtype
+        )
+        for i in range(nclasses):
+            array[i * row_size : i * row_size + row_size, :] = cmap[i]
+        imshow(array)
+        plt.yticks([row_size * i + row_size / 2 for i in range(nclasses)], labels)
+        plt.xticks([])
+        plt.show()
+        # plt.savefig("./labels.png")
 
     def initialize_semantics(self):
         if self.param.semantic_segmentation:
@@ -81,10 +129,12 @@ def initialize_semantics(self):
             )
             indices = []
             channels = []
-            for chan in self.param.channels:
+            for it, chan in enumerate(self.param.channels):
                 if chan in [cls for cls in list(class_to_idx.keys())]:
                     indices.append(class_to_idx[chan])
                     channels.append(chan)
+                elif self.param.fusion[it] in ["class_average", "class_bayesian"]:
+                    print(chan, " is not in the semantic segmentation model.")
             self.segmentation_channels = dict(zip(channels, indices))
         if self.param.feature_extractor:
             self.feature_channels = []
@@ -111,6 +161,7 @@ def cam_info_callback(self, msg):
         self.P = cp.resize(a, (3, 4))
         self.height = msg.height
         self.width = msg.width
+        self.header = msg.header
 
     def image_callback(self, rgb_msg, depth_msg, confidence_msg=None):
         confidence = None
@@ -190,13 +241,26 @@ def perform_segmentation(self, image, points, u, v):
         values = mask[:, v.get(), u.get()].cpu().detach().numpy()
         for it, channel in enumerate(self.segmentation_channels.keys()):
             points[channel] = values[it]
+        if self.param.publish_segmentation_image:
+            self.publish_segmentation_image(mask)
 
     def extract_features(self, image, points, u, v):
         prediction = self.feature_extractor["model"](image)
         values = prediction[:, v.get(), u.get()].cpu().detach().numpy()
         for it, channel in enumerate(self.feature_channels):
             points[channel] = values[it]
 
+    def publish_segmentation_image(self, probabilities):
+        probabilities = cp.asarray(probabilities)
+        colors = cp.asarray(self.semseg_color_map)
+        assert colors.ndim == 2 and colors.shape[1] == 3
+        img = cp.argmax(probabilities, axis=0)
+        img = colors[img].astype(cp.uint8)  # N x H x W x 3
+        img = img.get()
+        seg_msg = self.cv_bridge.cv2_to_imgmsg(img, encoding="rgb8")
+        seg_msg.header.frame_id = self.header.frame_id
+        self.seg_pub.publish(seg_msg)
+
 
 if __name__ == "__main__":
     sensor_name = sys.argv[1]

sensor_processing/semantic_pointcloud/script/semantic_pointcloud/pointcloud_parameters.py

@@ -25,6 +25,8 @@ class PointcloudParameter(Serializable):
 
     semantic_segmentation: bool = True
     segmentation_model: str = "detectron_coco_panoptic_fpn_R_101_3x"
+    publish_segmentation_image: bool = True
+    segmentation_image_topic: str = "/semantic_pointcloud/sem_seg"
 
     cam_info_topic: str = "/zed2i/zed_node/depth/camera_info"
     image_topic: str = "/zed2i/zed_node/left/image_rect_color"