Merge pull request #63 from ch-sa/ch-sa/pcd-loader

Add point cloud handlers

config.ini

@@ -60,4 +60,3 @@ far_plane = 300
 keep_perspective = False
 ; show button to visualize related images in a separate window
 show_2d_image = False
-

labelCloud/__init__.py

@@ -1 +1 @@
-__version__ = "0.7.4"
+__version__ = "0.7.5"

labelCloud/control/label_manager.py

@@ -2,8 +2,8 @@
 from pathlib import Path
 from typing import List
 
-from ..label_formats import BaseLabelFormat, CentroidFormat, KittiFormat, VerticesFormat
-from ..model.bbox import BBox
+from ..io.labels import BaseLabelFormat, CentroidFormat, KittiFormat, VerticesFormat
+from ..model import BBox
 from .config_manager import config
 
 

labelCloud/control/pcd_manager.py

@@ -3,64 +3,44 @@
 Sets the point cloud and original point cloud path. Initiate the writing to the virtual object buffer.
 """
 import logging
-from dataclasses import dataclass
 from pathlib import Path
 from shutil import copyfile
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
+import pkg_resources
+
 import numpy as np
 import open3d as o3d
 
-from ..model import BBox, PointCloud
-from ..utils.logger import end_section, green, print_column, start_section
+from ..io.pointclouds import BasePointCloudHandler, Open3DHandler
+from ..model import BBox, Perspective, PointCloud
+from ..utils.logger import blue, green, print_column
 from .config_manager import config
 from .label_manager import LabelManager
 
 if TYPE_CHECKING:
     from ..view.gui import GUI
 
-import pkg_resources
-
-
-@dataclass
-class Perspective(object):
-    zoom: float
-    rotation: Tuple[float, float, float]
-
-
-def color_pointcloud(points, z_min, z_max) -> np.ndarray:
-    palette = np.loadtxt(
-        pkg_resources.resource_filename("labelCloud.resources", "rocket-palette.txt")
-    )
-    palette_len = len(palette) - 1
-
-    colors = np.zeros(points.shape)
-    for ind, height in enumerate(points[:, 2]):
-        colors[ind] = palette[round((height - z_min) / (z_max - z_min) * palette_len)]
-    return colors
-
 
 class PointCloudManger(object):
-    PCD_EXTENSIONS = [".pcd", ".ply", ".pts", ".xyz", ".xyzn", ".xyzrgb", ".bin"]
+    PCD_EXTENSIONS = BasePointCloudHandler.get_supported_extensions()
     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) -> None:
         # Point cloud management
         self.pcd_folder = config.getpath("FILE", "pointcloud_folder")
         self.pcds: List[Path] = []
         self.current_id = -1
 
-        self.current_o3d_pcd = None
         self.view: Optional[GUI] = None
         self.label_manager = LabelManager()
 
         # Point cloud control
         self.pointcloud = None
         self.collected_object_classes = set()
-        self.saved_perspective: Perspective = None
+        self.saved_perspective: Optional[Perspective] = None
 
     @property
     def pcd_path(self) -> Path:
@@ -95,7 +75,7 @@ def read_pointcloud_folder(self) -> None:
             self.view.update_status(
                 "Please set the point cloud folder to a location that contains point cloud files."
             )
-            self.pointcloud = self.load_pointcloud(
+            self.pointcloud = PointCloud.from_file(
                 Path(
                     pkg_resources.resource_filename(
                         "labelCloud.resources", "labelCloud_icon.pcd"
@@ -115,7 +95,10 @@ def get_next_pcd(self) -> None:
         logging.info("Loading next point cloud...")
         if self.pcds_left():
             self.current_id += 1
-            self.pointcloud = self.load_pointcloud(self.pcd_path)
+            self.save_current_perspective()
+            self.pointcloud = PointCloud.from_file(
+                self.pcd_path, self.saved_perspective
+            )
             self.update_pcd_infos()
         else:
             logging.warning("No point clouds left!")
@@ -124,7 +107,10 @@ def get_prev_pcd(self) -> None:
         logging.info("Loading previous point cloud...")
         if self.current_id > 0:
             self.current_id -= 1
-            self.pointcloud = self.load_pointcloud(self.pcd_path)
+            self.save_current_perspective()
+            self.pointcloud = PointCloud.from_file(
+                self.pcd_path, self.saved_perspective
+            )
             self.update_pcd_infos()
         else:
             raise Exception("No point cloud left for loading!")
@@ -150,86 +136,12 @@ def save_labels_into_file(self, bboxes: List[BBox]) -> None:
         else:
             logging.warning("No point clouds to save labels for!")
 
-    def save_current_perspective(self, active: bool = True) -> None:
-        if active and self.pointcloud:
-            self.saved_perspective = Perspective(
-                zoom=self.pointcloud.trans_z,
-                rotation=tuple(self.pointcloud.get_rotations()),
-            )
+    def save_current_perspective(self) -> None:
+        if config.getboolean("USER_INTERFACE", "KEEP_PERSPECTIVE") and self.pointcloud:
+            self.saved_perspective = Perspective.from_point_cloud(self.pointcloud)
             logging.info(f"Saved current perspective ({self.saved_perspective}).")
-        else:
-            self.saved_perspective = None
-            logging.info("Reset saved perspective.")
 
     # MANIPULATOR
-    def load_pointcloud(self, path_to_pointcloud: Path) -> PointCloud:
-        start_section(f"Loading {path_to_pointcloud.name}")
-
-        if config.getboolean("USER_INTERFACE", "keep_perspective"):
-            self.save_current_perspective()
-
-        if path_to_pointcloud.suffix == ".bin":  # Loading binary pcds with numpy
-            bin_pcd = np.fromfile(path_to_pointcloud, dtype=np.float32)
-            points = bin_pcd.reshape((-1, 4))[
-                :, 0:3
-            ]  # Reshape and drop reflection values
-            points = points[~np.isnan(points).any(axis=1)]  # drop rows with nan
-            self.current_o3d_pcd = o3d.geometry.PointCloud(
-                o3d.utility.Vector3dVector(points)
-            )
-        else:  # Load point cloud with open3d
-            self.current_o3d_pcd = o3d.io.read_point_cloud(
-                str(path_to_pointcloud), remove_nan_points=True
-            )
-
-        tmp_pcd = PointCloud(path_to_pointcloud)
-        tmp_pcd.points = np.asarray(self.current_o3d_pcd.points).astype(
-            "float32"
-        )  # Unpack point cloud
-        tmp_pcd.colors = np.asarray(self.current_o3d_pcd.colors).astype("float32")
-
-        tmp_pcd.colorless = len(tmp_pcd.colors) == 0
-
-        print_column(["Number of Points:", f"{len(tmp_pcd.points):n}"])
-        # Calculate and set initial translation to view full pointcloud
-        tmp_pcd.center = self.current_o3d_pcd.get_center()
-        tmp_pcd.set_mins_maxs()
-
-        if PointCloudManger.COLORIZE and tmp_pcd.colorless:
-            logging.info("Generating colors for colorless point cloud!")
-            min_height, max_height = tmp_pcd.get_min_max_height()
-            tmp_pcd.colors = color_pointcloud(tmp_pcd.points, min_height, max_height)
-            tmp_pcd.colorless = False
-
-        max_dims = np.subtract(tmp_pcd.pcd_maxs, tmp_pcd.pcd_mins)
-        diagonal = min(
-            np.linalg.norm(max_dims),
-            config.getfloat("USER_INTERFACE", "far_plane") * 0.9,
-        )
-
-        tmp_pcd.init_translation = -self.current_o3d_pcd.get_center() - [0, 0, diagonal]
-
-        if self.saved_perspective != None:
-            tmp_pcd.init_translation = tuple(
-                list(tmp_pcd.init_translation[:2]) + [self.saved_perspective.zoom]
-            )
-            tmp_pcd.set_rotations(*self.saved_perspective.rotation)
-
-        tmp_pcd.reset_translation()
-        tmp_pcd.print_details()
-        if self.pointcloud is not None:  # Skip first pcd to intialize OpenGL first
-            tmp_pcd.write_vbo()
-
-        logging.info(
-            green(f"Successfully loaded point cloud from {path_to_pointcloud}!")
-        )
-        if tmp_pcd.colorless:
-            logging.warning(
-                "Did not find colors for the loaded point cloud, drawing in colorless mode!"
-            )
-        end_section()
-        return tmp_pcd
-
     def rotate_around_x(self, dangle) -> None:
         self.pointcloud.set_rot_x(self.pointcloud.rot_x - dangle)
 
@@ -259,7 +171,7 @@ def zoom_into(self, distance) -> None:
         self.pointcloud.set_trans_z(self.pointcloud.trans_z + zoom_distance)
 
     def reset_translation(self) -> None:
-        self.pointcloud.reset_translation()
+        self.pointcloud.reset_perspective()
 
     def reset_rotation(self) -> None:
         self.pointcloud.rot_x, self.pointcloud.rot_y, self.pointcloud.rot_z = (0, 0, 0)
@@ -278,31 +190,32 @@ def rotate_pointcloud(
             str(self.pcd_path),
             str(originals_path.joinpath(self.pcd_name)),
         )
+        logging.info("Copyied the original point cloud to %s.", blue(originals_path))
 
         # Rotate and translate point cloud
         rotation_matrix = o3d.geometry.get_rotation_matrix_from_axis_angle(
             np.multiply(axis, angle)
         )
-        self.current_o3d_pcd.rotate(rotation_matrix, center=tuple(rotation_point))
-        self.current_o3d_pcd.translate([0, 0, -rotation_point[2]])
+        o3d_pointcloud = Open3DHandler().to_open3d_point_cloud(self.pointcloud)
+        o3d_pointcloud.rotate(rotation_matrix, center=tuple(rotation_point))
+        o3d_pointcloud.translate([0, 0, -rotation_point[2]])
+        logging.info("Rotating point cloud...")
+        print_column(["Angle:", np.round(angle, 3)])
+        print_column(["Axis:", np.round(axis, 3)])
+        print_column(["Point:", np.round(rotation_point, 3)], last=True)
 
         # Check if pointcloud is upside-down
-        pcd_mins = np.amin(self.current_o3d_pcd.points, axis=0)
-        pcd_maxs = np.amax(self.current_o3d_pcd.points, axis=0)
-
-        if abs(pcd_mins[2]) > pcd_maxs[2]:
+        if abs(self.pointcloud.pcd_mins[2]) > self.pointcloud.pcd_maxs[2]:
             logging.warning("Point cloud is upside down, rotating ...")
-            self.current_o3d_pcd.rotate(
+            o3d_pointcloud.rotate(
                 o3d.geometry.get_rotation_matrix_from_xyz([np.pi, 0, 0]),
                 center=(0, 0, 0),
             )
 
-        save_path = self.pcd_path
-        if save_path.suffix == ".bin":  # save .bin point clouds as .pcd
-            save_path = save_path.parent.joinpath(save_path.stem + ".pcd")
-
-        o3d.io.write_point_cloud(str(save_path), self.current_o3d_pcd)
-        self.pointcloud = self.load_pointcloud(save_path)
+        self.pointcloud = PointCloud(
+            self.pcd_path, *Open3DHandler().to_point_cloud(o3d_pointcloud)
+        )
+        self.pointcloud.to_file()
 
     # HELPER
 

labelCloud/io/labels/__init__.py

@@ -0,0 +1,8 @@
+from .base import (
+    BaseLabelFormat,
+    abs2rel_rotation,
+    rel2abs_rotation,
+)
+from .centroid import CentroidFormat
+from .kitti import KittiFormat
+from .vertices import VerticesFormat

labelCloud/label_formats/base.py → labelCloud/io/labels/base.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from ..model import BBox
+from ...model import BBox
 
 
 class BaseLabelFormat(ABC):

labelCloud/label_formats/centroid_format.py → labelCloud/io/labels/centroid.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import List
 
-from ..model import BBox
+from ...model import BBox
 from . import BaseLabelFormat, abs2rel_rotation, rel2abs_rotation
 
 

labelCloud/label_formats/kitti_format.py → labelCloud/io/labels/kitti.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import List
 
-from ..model import BBox
+from ...model import BBox
 from . import BaseLabelFormat, abs2rel_rotation, rel2abs_rotation
 
 

labelCloud/label_formats/vertices_format.py → labelCloud/io/labels/vertices.py

@@ -5,8 +5,8 @@
 
 import numpy as np
 
-from ..model import BBox
-from ..utils import math3d
+from ...model import BBox
+from ...utils import math3d
 from . import BaseLabelFormat
 
 

labelCloud/io/pointclouds/__init__.py

@@ -0,0 +1,3 @@
+from .base import BasePointCloudHandler
+from .numpy import NumpyHandler
+from .open3d import Open3DHandler

labelCloud/io/pointclouds/base.py

@@ -0,0 +1,46 @@
+import logging
+from abc import abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Set, Tuple
+
+import numpy as np
+
+from ...utils.logger import blue
+from ...utils.singleton import SingletonABCMeta
+
+if TYPE_CHECKING:
+    from ...model import PointCloud
+
+
+class BasePointCloudHandler(object, metaclass=SingletonABCMeta):
+    EXTENSIONS = set()  # should be set in subclasses
+
+    @abstractmethod
+    def read_point_cloud(self, path: Path) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        """Read a point cloud file and return only the points and colors as array."""
+        logging.info(
+            blue("Loading point cloud from %s using %s."), path, self.__class__.__name__
+        )
+        pass
+
+    @abstractmethod
+    def write_point_cloud(self, path: Path, pointcloud: "PointCloud") -> None:
+        logging.info(
+            blue("Writing point cloud to %s using %s."), path, self.__class__.__name__
+        )
+        pass
+
+    @classmethod
+    def get_supported_extensions(cls) -> Set[str]:
+        return set().union(*[handler.EXTENSIONS for handler in cls.__subclasses__()])
+
+    @classmethod
+    def get_handler(cls, file_extension: str) -> "BasePointCloudHandler":
+        """Return a point cloud handler for the given file extension."""
+        for subclass in cls.__subclasses__():
+            if file_extension in subclass.EXTENSIONS:
+                return subclass()
+
+        logging.error(
+            "No point cloud handler found for file extension %s.", file_extension
+        )

labelCloud/io/pointclouds/numpy.py

@@ -0,0 +1,32 @@
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Tuple
+
+import numpy as np
+
+from . import BasePointCloudHandler
+
+if TYPE_CHECKING:
+    from ...model import PointCloud
+
+
+class NumpyHandler(BasePointCloudHandler):
+    EXTENSIONS = {".bin"}
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def read_point_cloud(self, path: Path) -> Tuple[np.ndarray, None]:
+        """Read point cloud file as array and drop reflection and nan values."""
+        super().read_point_cloud(path)
+        points = np.fromfile(path, dtype=np.float32)
+        points = points.reshape((-1, 4 if len(points) % 4 == 0 else 3))[:, 0:3]
+        return (points[~np.isnan(points).any(axis=1)], None)
+
+    def write_point_cloud(self, path: Path, pointcloud: "PointCloud") -> None:
+        """Write point cloud points into binary file."""
+        super().write_point_cloud(path, pointcloud)
+        logging.warning(
+            "Only writing point coordinates, any previous reflection values will be dropped."
+        )
+        pointcloud.points.tofile(path)