Merge pull request #11 from ch-sa/file-menu

File menu

labelCloud/control/alignmode.py

@@ -11,8 +11,8 @@
 
 
 class AlignMode:
-    def __init__(self, pcd_controller: PointCloudManger):
-        self.pcd_controller = pcd_controller
+    def __init__(self, pcd_manager: PointCloudManger):
+        self.pcd_manager = pcd_manager
         self.view: Union[GUI, None] = None
         self.activated = False
         self.point_color = (1, 1, 0, 1)
@@ -105,7 +105,7 @@ def calculate_angles(self):
         print("Alignment rotation: %s around %s" % (round(rotation_angle, 2), np.round(rotation_axis, 2)))
 
         # Initiate point cloud rotation
-        self.pcd_controller.rotate_pointcloud(rotation_axis, rotation_angle, self.plane1)
+        self.pcd_manager.rotate_pointcloud(rotation_axis, rotation_angle, self.plane1)
 
         self.view.update_status("Aligned point cloud with the selected floor.", "navigation")
         self.change_activation(force=False)

labelCloud/control/controller.py

@@ -17,12 +17,12 @@ class Controller:
     def __init__(self):
         """Initializes all controllers and managers."""
         self.view: Union['GUI', None] = None
-        self.pcd_controller = PointCloudManger()
+        self.pcd_manager = PointCloudManger()
         self.bbox_controller = BoundingBoxController()
 
         # Drawing states
         self.drawing_mode = DrawingManager(self.bbox_controller)
-        self.align_mode = AlignMode(self.pcd_controller)
+        self.align_mode = AlignMode(self.pcd_manager)
 
         # Control states
         self.curr_cursor_pos = None  # updated by mouse movement
@@ -34,16 +34,19 @@ def __init__(self):
         self.side_mode = False
         self.selected_side = None
 
-    def set_view(self, view: 'GUI'):
+    def startup(self, view: 'GUI'):
         """Sets the view in all controllers and dependent modules; Loads labels from file."""
         self.view = view
         self.bbox_controller.set_view(self.view)
-        self.pcd_controller.set_view(self.view)
+        self.pcd_manager.set_view(self.view)
         self.drawing_mode.set_view(self.view)
         self.align_mode.set_view(self.view)
         self.view.glWidget.set_bbox_controller(self.bbox_controller)
-        self.bbox_controller.pcdc = self.pcd_controller
-        self.bbox_controller.set_bboxes(self.pcd_controller.get_labels_from_file())  # Load labels for first pcd
+        self.bbox_controller.pcdc = self.pcd_manager
+
+        # Read labels from folders
+        self.pcd_manager.read_pointcloud_folder()
+        self.next_pcd(save=False)
 
     def loop_gui(self):
         """Function collection called during each event loop iteration."""
@@ -52,27 +55,28 @@ def loop_gui(self):
         self.view.glWidget.updateGL()
 
     # POINT CLOUD METHODS
-    def next_pcd(self):
-        self.save()
-        if self.pcd_controller.pcds_left():
-            self.pcd_controller.get_next_pcd()
+    def next_pcd(self, save: bool = True):
+        if save:
+            self.save()
+        if self.pcd_manager.pcds_left():
+            self.pcd_manager.get_next_pcd()
             self.reset()
-            self.bbox_controller.set_bboxes(self.pcd_controller.get_labels_from_file())
+            self.bbox_controller.set_bboxes(self.pcd_manager.get_labels_from_file())
         else:
-            self.view.update_progress(self.pcd_controller.no_of_pcds)
+            self.view.update_progress(len(self.pcd_manager.pcds))
             self.view.button_next_pcd.setEnabled(False)
 
     def prev_pcd(self):
         self.save()
-        if self.pcd_controller.current_id > 0:
-            self.pcd_controller.get_prev_pcd()
+        if self.pcd_manager.current_id > 0:
+            self.pcd_manager.get_prev_pcd()
             self.reset()
-            self.bbox_controller.set_bboxes(self.pcd_controller.get_labels_from_file())
+            self.bbox_controller.set_bboxes(self.pcd_manager.get_labels_from_file())
 
     # CONTROL METHODS
     def save(self):
         """Saves all bounding boxes in the label file."""
-        self.pcd_controller.save_labels_into_file(self.bbox_controller.get_bboxes())
+        self.pcd_manager.save_labels_into_file(self.bbox_controller.get_bboxes())
 
     def reset(self):
         """Resets the controllers and bounding boxes from the current screen."""
@@ -147,11 +151,11 @@ def mouse_move_event(self, a0: QtGui.QMouseEvent):
                     self.bbox_controller.set_center(*new_center)  # absolute positioning
             else:
                 if a0.buttons() & QtCore.Qt.LeftButton:  # pcd rotation
-                    self.pcd_controller.rotate_around_x(dy)
-                    self.pcd_controller.rotate_around_z(dx)
+                    self.pcd_manager.rotate_around_x(dy)
+                    self.pcd_manager.rotate_around_z(dx)
                 elif a0.buttons() & QtCore.Qt.RightButton:  # pcd translation
-                    self.pcd_controller.translate_along_x(dx)
-                    self.pcd_controller.translate_along_y(dy)
+                    self.pcd_manager.translate_along_x(dx)
+                    self.pcd_manager.translate_along_y(dy)
 
             if dx > 0.1 or dy > 0.1:  # Reset scroll locks if significant cursor movements
                 if self.side_mode:
@@ -172,7 +176,7 @@ def mouse_scroll_event(self, a0: QtGui.QWheelEvent):
             self.bbox_controller.get_active_bbox().change_side(self.selected_side,
                                                                -a0.angleDelta().y() / 4000)  # ToDo implement method
         else:
-            self.pcd_controller.zoom_into(a0.angleDelta().y())
+            self.pcd_manager.zoom_into(a0.angleDelta().y())
             self.scroll_mode = True
 
     def key_press_event(self, a0: QtGui.QKeyEvent):
@@ -184,7 +188,7 @@ def key_press_event(self, a0: QtGui.QKeyEvent):
 
         # Reset point cloud pose to intial rotation and translation
         elif (a0.key() == QtCore.Qt.Key_R) or (a0.key() == QtCore.Qt.Key_Home):
-            self.pcd_controller.reset_transformations()
+            self.pcd_manager.reset_transformations()
             print("Reseted position to default.")
 
         elif a0.key() == QtCore.Qt.Key_Delete:  # Delete active bbox

labelCloud/control/drawing_manager.py

@@ -288,7 +288,7 @@ def register_tmp_point(self, new_tmp_point: List[float]):
             self.tmp_p2 = new_tmp_point
 
     def get_bbox(self):
-        # self.view.controller.pcd_controller.
+        # self.view.controller.pcd_manager.
         pass
 
     def draw_preview(self):

labelCloud/control/label_manager.py

@@ -11,27 +11,29 @@
 from model.bbox import BBox
 
 
+def get_label_strategy(export_format: str, label_folder: str) -> 'IFormattingInterface':
+    if export_format == "vertices":
+        return VerticesFormat(label_folder)
+    elif export_format == "centroid_rel":
+        return CentroidFormat(label_folder, relative_rotation=True)
+    elif export_format == "kitti":
+        return KittiFormat(label_folder, relative_rotation=True)
+    elif export_format != "centroid_abs":
+        print(f"Unknown export strategy '{export_format}'. Proceeding with default (centroid_abs)!")
+    return CentroidFormat(label_folder, relative_rotation=False)
+
+
 class LabelManager:
     LABEL_FORMATS = ["vertices", "centroid_rel", "centroid_abs", "kitti"]  # supported export formats
     STD_LABEL_FORMAT = config.get("LABEL", "label_format")
     EXPORT_PRECISION = config.getint("LABEL", "export_precision")  # Number of decimal places
-    STD_LABEL_FOLDER = config.get("FILE", "label_folder")
 
-    def __init__(self, strategy: str = STD_LABEL_FORMAT, path_to_label_folder: str = STD_LABEL_FOLDER):
-        self.label_folder = path_to_label_folder
+    def __init__(self, strategy: str = STD_LABEL_FORMAT, path_to_label_folder: str = None):
+        self.label_folder = path_to_label_folder or config.get("FILE", "label_folder")
         if not os.path.isdir(self.label_folder):
             os.mkdir(self.label_folder)
-        if strategy == "vertices":
-            self.label_strategy = VerticesFormat(self.label_folder)
-        elif strategy == "centroid_abs":
-            self.label_strategy = CentroidFormat(self.label_folder, relative_rotation=False)
-        elif strategy == "centroid_rel":
-            self.label_strategy = CentroidFormat(self.label_folder, relative_rotation=True)
-        elif strategy == "kitti":
-            self.label_strategy = KittiFormat(self.label_folder, relative_rotation=True)  # KITTI is always relative
-        else:
-            self.label_strategy = CentroidFormat(self.label_folder, relative_rotation=False)
-            print("Unknown export strategy '%s'. Proceeding with default (corners)!" % strategy)
+
+        self.label_strategy = get_label_strategy(strategy, self.label_folder)
 
     def import_labels(self, pcd_name: str) -> List[BBox]:
         try:
@@ -116,6 +118,9 @@ def import_labels(self, pcd_name_stripped):
     def export_labels(self, bboxes, pcd_name, pcd_folder, pcd_path):
         raise NotImplementedError
 
+    def update_label_folder(self, new_label_folder):
+        self.label_folder = new_label_folder
+
 
 class VerticesFormat(IFormattingInterface, ABC):
 

labelCloud/control/pcd_manager.py

@@ -4,12 +4,10 @@
 """
 import ntpath
 import os
-import sys
-from typing import List, Tuple, TYPE_CHECKING
+from typing import List, Tuple, TYPE_CHECKING, Optional
 
 import numpy as np
 import open3d as o3d
-from PyQt5.QtWidgets import QMessageBox
 from shutil import copyfile
 
 from control.config_manager import config
@@ -21,30 +19,6 @@
     from view.gui import GUI
 
 
-def find_pcd_files(path: str) -> List[str]:
-    if not os.path.isdir(path):  # Check if point cloud folder exists
-        show_no_pcd_dialog()
-        sys.exit()
-
-    pcd_files = []
-    for file in os.listdir(path):
-        if file.endswith(PointCloudManger.PCD_EXTENSIONS):
-            pcd_files.append(file)
-
-    return sorted(pcd_files)
-
-
-def show_no_pcd_dialog():
-    msg = QMessageBox()
-    msg.setIcon(QMessageBox.Warning)
-    msg.setText("Did not find any point cloud.")
-    msg.setInformativeText("Please copy all your point clouds into the <i>%s</i> folder. "
-                           "If you already have done that check the supported formats %s."
-                           % (PointCloudManger.PCD_FOLDER, str(PointCloudManger.PCD_EXTENSIONS)))
-    msg.setWindowTitle("No Point Clouds Found")
-    msg.exec_()
-
-
 def color_pointcloud(points, z_min, z_max):
     palette = np.loadtxt("labelCloud/ressources/rocket-palette.txt")
     palette_len = len(palette) - 1
@@ -56,29 +30,51 @@ def color_pointcloud(points, z_min, z_max):
 
 
 class PointCloudManger:
-    PCD_EXTENSIONS = (".pcd", ".ply", ".pts", ".xyz", ".xyzn", ".xyzrgb", ".bin")
-    PCD_FOLDER = config.get("FILE", "POINTCLOUD_FOLDER")
+    PCD_EXTENSIONS = [".pcd", ".ply", ".pts", ".xyz", ".xyzn", ".xyzrgb", ".bin"]
     ORIGINALS_FOLDER = "original_pointclouds"
     TRANSLATION_FACTOR = config.getfloat("POINTCLOUD", "STD_TRANSLATION")
     ZOOM_FACTOR = config.getfloat("POINTCLOUD", "STD_ZOOM")
     COLORIZE = config.getboolean("POINTCLOUD", "COLORLESS_COLORIZE")
 
     def __init__(self):
         # Point cloud management
-        self.pcd_folder = PointCloudManger.PCD_FOLDER
-        self.pcds = find_pcd_files(self.pcd_folder)
-        self.no_of_pcds = len(self.pcds)
+        self.pcd_folder = config.get("FILE", "pointcloud_folder")
+        self.pcds = []
         self.current_id = -1
+
         self.current_o3d_pcd = None
-        self.view = None
+        self.view: Optional[GUI] = None
         self.label_manager = LabelManager()
+
         # Point cloud control
         self.pointcloud = None
         self.collected_object_classes = set()
 
+    def read_pointcloud_folder(self):
+        """Checks point cloud folder and sets self.pcds to all valid point cloud file names."""
+        if os.path.isdir(self.pcd_folder):
+            self.pcds = []
+            for file in sorted(os.listdir(self.pcd_folder), key=str.casefold):
+                if file.endswith(tuple(PointCloudManger.PCD_EXTENSIONS)):
+                    self.pcds.append(file)
+        else:
+            print(f"Point cloud path {self.pcd_folder} is not a valid directory.")
+
+        if self.pcds:
+            self.view.update_status(f"Found {len(self.pcds)} point clouds in the point cloud folder.")
+            self.update_pcd_infos()
+        else:
+            self.view.show_no_pointcloud_dialog(self.pcd_folder, PointCloudManger.PCD_EXTENSIONS)
+            self.view.update_status("Please set the point cloud folder to a location that contains point cloud files.")
+            self.pointcloud = self.load_pointcloud("labelCloud/ressources/labelCloud_icon.pcd")
+            self.update_pcd_infos(pointcloud_label=" – (select folder!)")
+
+        self.view.init_progress(min_value=0, max_value=len(self.pcds))
+        self.current_id = -1
+
     # GETTER
     def pcds_left(self) -> bool:
-        return self.current_id + 1 < self.no_of_pcds
+        return self.current_id + 1 < len(self.pcds)
 
     def get_next_pcd(self):
         print("Loading next point cloud...")
@@ -87,10 +83,7 @@ def get_next_pcd(self):
             self.pointcloud = self.load_pointcloud(self.get_current_path())
             self.update_pcd_infos()
         else:
-            if self.current_id == -1:
-                show_no_pcd_dialog()
-                sys.exit()
-            raise Exception("No point cloud left for loading!")
+            print("No point clouds left!")
 
     def get_prev_pcd(self):
         print("Loading previous point cloud...")
@@ -110,7 +103,7 @@ def get_current_name(self) -> str:
 
     def get_current_details(self) -> Tuple[str, int, int]:
         if self.current_id >= 0:
-            return self.get_current_name(), self.current_id + 1, self.no_of_pcds
+            return self.get_current_name(), self.current_id + 1, len(self.pcds)
 
     def get_current_path(self) -> str:
         return os.path.join(self.pcd_folder, self.pcds[self.current_id])
@@ -123,13 +116,16 @@ def get_labels_from_file(self):
     # SETTER
     def set_view(self, view: 'GUI') -> None:
         self.view = view
-        self.view.init_progress(min_value=0, max_value=self.no_of_pcds)
         self.view.glWidget.set_pointcloud_controller(self)
-        self.get_next_pcd()
 
     def save_labels_into_file(self, bboxes: List[BBox]):
-        self.label_manager.export_labels(self.get_current_path(), bboxes)
-        self.view.update_label_completer({bbox.get_classname() for bbox in bboxes})
+        if self.pcds:
+            self.label_manager.export_labels(self.get_current_path(), bboxes)
+            self.collected_object_classes.update({bbox.get_classname() for bbox in bboxes})
+            self.view.update_label_completer(self.collected_object_classes)
+            self.view.update_default_object_class_menu(self.collected_object_classes)
+        else:
+            print("No point clouds to save labels for!")
 
     # MANIPULATOR
     def load_pointcloud(self, path_to_pointcloud: str) -> PointCloud:
@@ -260,12 +256,14 @@ def get_perspective(self):
 
     # UPDATE GUI
 
-    def update_pcd_infos(self):
-        self.view.set_pcd_label(self.get_current_name())
+    def update_pcd_infos(self, pointcloud_label: str = None):
+        self.view.set_pcd_label(pointcloud_label or self.get_current_name())
         self.view.update_progress(self.current_id)
 
         if self.current_id <= 0:
             self.view.button_prev_pcd.setEnabled(False)
+            if self.pcds:
+                self.view.button_next_pcd.setEnabled(True)
         else:
             self.view.button_next_pcd.setEnabled(True)
             self.view.button_prev_pcd.setEnabled(True)