Show class colour and filter points by class in segmentation model

labelCloud/control/controller.py

@@ -52,12 +52,16 @@ def startup(self, view: "GUI") -> None:
         # Read labels from folders
         self.pcd_manager.read_pointcloud_folder()
         self.next_pcd(save=False)
+        if LabelConfig().type == LabelingMode.SEMANTIC_SEGMENTATION:
+            self.pcd_manager.populate_segmentation_list()
 
     def loop_gui(self) -> None:
         """Function collection called during each event loop iteration."""
         self.set_crosshair()
         self.set_selected_side()
         self.view.gl_widget.updateGL()
+        if LabelConfig().type == LabelingMode.SEMANTIC_SEGMENTATION:
+            self.pcd_manager.loop_seg_list_check_state()
 
     # POINT CLOUD METHODS
     def next_pcd(self, save: bool = True) -> None:

labelCloud/control/pcd_manager.py

@@ -10,12 +10,17 @@
 import numpy as np
 import open3d as o3d
 import pkg_resources
+from PyQt5 import QtCore
+from PyQt5.QtWidgets import QCheckBox, QLabel
 
-from ..definitions import LabelingMode, Point3D
+from ..definitions import LabelingMode
+from ..definitions.types import Point3D
 from ..io.labels.config import LabelConfig
 from ..io.pointclouds import BasePointCloudHandler, Open3DHandler
 from ..model import BBox, Perspective, PointCloud
+from ..utils.color import rgb_to_hex
 from ..utils.logger import blue, green, print_column
+from ..view.startup.color_button import ColorButton
 from .config_manager import config
 from .label_manager import LabelManager
 
@@ -263,7 +268,8 @@ def rotate_pointcloud(
 
     def assign_point_label_in_box(self, box: BBox) -> None:
         assert self.pointcloud is not None
-        points = self.pointcloud.points
+        points = self.pointcloud.points.copy()
+        points[~self.pointcloud.visible] = np.finfo(np.float32).max
         points_inside = box.is_inside(points)
 
         # Relabel the points if its inside the box
@@ -306,3 +312,53 @@ def update_pcd_infos(self, pointcloud_label: Optional[str] = None) -> None:
         else:
             self.view.button_next_pcd.setEnabled(True)
             self.view.button_prev_pcd.setEnabled(True)
+
+    def populate_segmentation_list(self) -> None:
+        assert self.pointcloud is not None
+        assert self.pointcloud.labels is not None
+        self.seg_list_label: List[QLabel] = []
+        self.seg_list_check_box: List[QCheckBox] = []
+        self.seg_list_check_state: List[QtCore.Qt.CheckState] = []
+        for idx, label_class in enumerate(LabelConfig().classes, start=1):
+            self.seg_list_label.append(QLabel(label_class.name))
+            check_box = QCheckBox()
+            check_box.setCheckState(QtCore.Qt.Checked)
+            color_button = ColorButton(
+                color=rgb_to_hex(label_class.color), changeable=False
+            )
+            self.seg_list_check_box.append(check_box)
+            self.seg_list_check_state.append(self.seg_list_check_box[-1].checkState())
+            self.view.segmentation_list.addWidget(self.seg_list_label[-1], idx, 0)
+            self.view.segmentation_list.addWidget(color_button, idx, 1)
+            self.view.segmentation_list.addWidget(self.seg_list_check_box[-1], idx, 2)
+
+    def loop_seg_list_check_state(self):
+        curr_checked_status = self.seg_list_check_state.copy()
+        any_changed = False
+
+        move_back = []
+        move_away = []
+        for idx, (box, prev_status) in enumerate(
+            zip(
+                self.seg_list_check_box,
+                self.seg_list_check_state,
+            )
+        ):
+            if box.checkState() != prev_status:
+                any_changed = True
+                curr_checked_status[idx] = box.checkState()
+                changed_id = LabelConfig().classes[idx].id
+                if box.checkState() == QtCore.Qt.Checked:
+                    move_back.append(changed_id)
+                else:
+                    move_away.append(changed_id)
+
+        if any_changed:
+            points_move_away = (
+                np.isin(self.pointcloud.labels, move_away) if move_away else None
+            )
+            points_move_back = (
+                np.isin(self.pointcloud.labels, move_back) if move_back else None
+            )
+            self.pointcloud.update_position_vbo(points_move_away, points_move_back)
+            self.seg_list_check_state = curr_checked_status

labelCloud/model/point_cloud.py

@@ -65,6 +65,8 @@ def __init__(
             self.validate_segmentation_label()
             self.mix_ratio = config.getfloat("POINTCLOUD", "label_color_mix_ratio")
 
+        self.visible = np.ones((self.points.shape[0],), dtype=np.bool_)
+
         self.vbo = None
         self.center: Point3D = tuple(np.sum(points[:, i]) / len(points) for i in range(3))  # type: ignore
         self.pcd_mins: npt.NDArray[np.float32] = np.amin(points, axis=0)
@@ -233,11 +235,51 @@ def color_with_label(self) -> bool:
     def has_label(self) -> bool:
         return self.labels is not None
 
+    def update_position_vbo(
+        self,
+        points_move_away: Optional[npt.NDArray[np.bool_]],
+        points_move_back: Optional[npt.NDArray[np.bool_]],
+    ):
+        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.position_vbo)
+        # Move points to super far away
+        if points_move_away is not None and np.any(points_move_away):
+            move_away_idxs = np.where(points_move_away)[0]
+            self.visible[move_away_idxs] = False
+            arrays = consecutive(move_away_idxs)
+            stride = self.points.shape[1] * SIZE_OF_FLOAT
+            for arr in arrays:
+                super_far: npt.NDArray[np.float32] = (
+                    np.ones((arr.shape[0], 3), dtype=np.float32)
+                    * np.finfo(np.float32).max
+                )
+                # partially update label_vbo from positions arr[0] to arr[-1]
+                GL.glBufferSubData(
+                    GL.GL_ARRAY_BUFFER,
+                    offset=arr[0] * stride,
+                    size=super_far.nbytes,
+                    data=super_far,
+                )
+        # Move points back
+        if points_move_back is not None and np.any(points_move_back):
+            move_back_idxs = np.where(points_move_back)[0]
+            self.visible[move_back_idxs] = True
+            arrays = consecutive(move_back_idxs)
+            stride = self.points.shape[1] * SIZE_OF_FLOAT
+
+            for arr in arrays:
+                points: npt.NDArray[np.float32] = self.points[arr]
+                GL.glBufferSubData(
+                    GL.GL_ARRAY_BUFFER,
+                    offset=arr[0] * stride,
+                    size=points.nbytes,
+                    data=points,
+                )
+
     def update_selected_points_in_label_vbo(
         self, points_inside: npt.NDArray[np.bool_]
     ) -> None:
-        """Send the selected updated label colors to label vbo. This function
-        assumes the `self.label_colors[points_inside]` have been altered.
+        """Send the selected updated label colors to `self.label_vbo`. This
+        function assumes the `self.label_colors[points_inside]` have been altered.
         This function only partially updates the label vbo to minimise the
         data sent to gpu. It leverages `glBufferSubData` method to perform
         partial update and `consecutive` method to find consecutive indexes

labelCloud/resources/interfaces/interface.ui

@@ -1460,6 +1460,97 @@
         </item>
        </layout>
       </item>
+      <item>
+       <widget class="QGroupBox" name="segmentation_list_group">
+        <property name="font">
+         <font>
+          <family>DejaVu Sans,Arial</family>
+          <weight>75</weight>
+          <bold>true</bold>
+         </font>
+        </property>
+        <property name="title">
+          <string>Segmentation Controls</string>
+        </property>
+        <layout class="QGridLayout" name="segmentation_list">
+        <property name="sizeConstraint">
+        <enum>QLayout::SetDefaultConstraint</enum>
+        </property>
+        <property name="topMargin">
+        <number>0</number>
+        </property>
+        <property name="verticalSpacing">
+        <number>6</number>
+        </property>
+         <item row="0" column="0">
+          <widget class="QLabel" name="segmentation_control_class_list">
+           <property name="sizePolicy">
+            <sizepolicy hsizetype="Preferred" vsizetype="Maximum">
+             <horstretch>0</horstretch>
+             <verstretch>0</verstretch>
+            </sizepolicy>
+           </property>
+           <property name="font">
+            <font>
+             <family>DejaVu Sans,Arial</family>
+             <pointsize>12</pointsize>
+             <weight>50</weight>
+             <bold>true</bold>
+            </font>
+           </property>
+           <property name="text">
+            <string>Class Name</string>
+           </property>
+          </widget>
+         </item>
+         <item row="0" column="1">
+          <widget class="QLabel" name="segmentation_control_color">
+           <property name="sizePolicy">
+            <sizepolicy hsizetype="Preferred" vsizetype="Maximum">
+             <horstretch>0</horstretch>
+             <verstretch>0</verstretch>
+            </sizepolicy>
+           </property>
+           <property name="font">
+            <font>
+             <family>DejaVu Sans,Arial</family>
+             <pointsize>12</pointsize>
+             <weight>50</weight>
+             <bold>true</bold>
+            </font>
+           </property>
+           <property name="text">
+            <string>Color</string>
+           </property>
+          </widget>
+         </item>
+         <item row="0" column="2">
+          <widget class="QLabel" name="segmentation_control_checked">
+           <property name="sizePolicy">
+            <sizepolicy hsizetype="Preferred" vsizetype="Maximum">
+             <horstretch>0</horstretch>
+             <verstretch>0</verstretch>
+            </sizepolicy>
+           </property>
+           <property name="font">
+            <font>
+             <family>DejaVu Sans,Arial</family>
+             <pointsize>12</pointsize>
+             <weight>50</weight>
+             <bold>true</bold>
+            </font>
+           </property>
+           <property name="text">
+            <string>Visible</string>
+           </property>
+          </widget>
+         </item>
+
+        </layout>
+       </widget>
+
+      </item>
+
       <item>
        <widget class="QListWidget" name="label_list">
         <property name="sizePolicy">

labelCloud/view/color_button.py

@@ -0,0 +1,59 @@
+from PyQt5 import QtGui, QtWidgets
+from PyQt5.QtCore import Qt, pyqtSignal
+
+
+class ColorButton(QtWidgets.QPushButton):
+    """
+    Custom Qt Widget to show a chosen color.
+
+    Left-clicking the button shows the color-chooser, while
+    right-clicking resets the color to None (no-color).
+
+    Source: https://www.pythonguis.com/widgets/qcolorbutton-a-color-selector-tool-for-pyqt/
+    """
+
+    colorChanged = pyqtSignal(object)
+
+    def __init__(self, *args, color="#FF0000", changeable: bool = True, **kwargs):
+        super(ColorButton, self).__init__(*args, **kwargs)
+
+        self._color = None
+        self._default = color
+        if changeable:
+            self.pressed.connect(self.onColorPicker)
+
+        # Set the initial/default state.
+        self.setColor(self._default)
+
+    def setColor(self, color):
+        if color != self._color:
+            self._color = color
+            self.colorChanged.emit(color)
+
+        if self._color:
+            self.setStyleSheet("background-color: %s;" % self._color)
+        else:
+            self.setStyleSheet("")
+
+    def color(self):
+        return self._color
+
+    def onColorPicker(self):
+        """
+        Show color-picker dialog to select color.
+
+        Qt will use the native dialog by default.
+
+        """
+        dlg = QtWidgets.QColorDialog(self)
+        if self._color:
+            dlg.setCurrentColor(QtGui.QColor(self._color))
+
+        if dlg.exec_():
+            self.setColor(dlg.currentColor().name())
+
+    def mousePressEvent(self, e):
+        if e.button() == Qt.RightButton:
+            self.setColor(self._default)
+
+        return super(ColorButton, self).mousePressEvent(e)

labelCloud/view/gui.py

@@ -20,6 +20,8 @@
     QMessageBox,
 )
 
+from labelCloud.view.startup.dialog import StartupDialog
+
 from ..control.config_manager import config
 from ..definitions import Color3f, LabelingMode
 from ..io.labels.config import LabelConfig
@@ -194,6 +196,8 @@ def __init__(self, control: "Controller") -> None:
         self.button_save_label: QtWidgets.QPushButton
 
         # RIGHT PANEL
+        self.segmentation_list_group: QtWidgets.QGroupBox
+        self.segmentation_list: QtWidgets.QGridLayout
         self.label_list: QtWidgets.QListWidget
         self.current_class_dropdown: QtWidgets.QComboBox
         self.button_deselect_label: QtWidgets.QPushButton
@@ -257,6 +261,7 @@ def __init__(self, control: "Controller") -> None:
         if LabelConfig().type == LabelingMode.OBJECT_DETECTION:
             self.button_assign_label.setVisible(False)
             self.act_color_with_label.setVisible(False)
+            self.segmentation_list_group.setVisible(False)
 
         # Connect with controller
         self.controller.startup(self)
@@ -521,9 +526,6 @@ def init_progress(self, min_value, max_value):
     def update_progress(self, value) -> None:
         self.progressbar_pcds.setValue(value)
 
-    def update_current_class_dropdown(self) -> None:
-        self.controller.pcd_manager.populate_class_dropdown()
-
     def update_bbox_stats(self, bbox) -> None:
         viewing_precision = config.getint("USER_INTERFACE", "viewing_precision")
         if bbox and not self.line_edited_activated():

labelCloud/view/startup/color_button.py

@@ -14,12 +14,13 @@ class ColorButton(QtWidgets.QPushButton):
 
     colorChanged = pyqtSignal(object)
 
-    def __init__(self, *args, color="#FF0000", **kwargs):
+    def __init__(self, *args, color="#FF0000", changeable: bool = True, **kwargs):
         super(ColorButton, self).__init__(*args, **kwargs)
 
         self._color = None
         self._default = color
-        self.pressed.connect(self.onColorPicker)
+        if changeable:
+            self.pressed.connect(self.onColorPicker)
 
         # Set the initial/default state.
         self.setColor(self._default)