refactor: Las and laz I/O (#100)

Author: josafatburmeister

src/pointtorch/io/_las_reader.py

@@ -95,8 +95,8 @@ def _read_points(
         for column_name in las_header.point_format.extra_dimension_names:
             if column_name.lower() in ["x", "y", "z"] or columns is not None and column_name not in columns:
                 continue
-            column_values = np.array(las_data[column_name])
-            point_cloud_df[column_name] = column_values
+
+            point_cloud_df[column_name] = np.array(las_data[column_name])
 
         return point_cloud_df
 

src/pointtorch/io/_las_writer.py

@@ -3,7 +3,7 @@
 __all__ = ["LasWriter"]
 
 import pathlib
-from typing import Optional, Union
+from typing import Optional, Tuple, Union
 
 import laspy
 from laspy.compression import LazrsBackend
@@ -44,15 +44,15 @@ def supported_file_formats(self) -> list[str]:
 
         return ["las", "laz"]
 
-    def _select_point_format(self, point_cloud: pd.DataFrame) -> int:
+    def _select_point_format(self, point_cloud: pd.DataFrame) -> laspy.point.format.PointFormat:
         """
         Determines the las file format that covers the most columns of the given point cloud.
 
         Returns:
-            ID of the chosen las file format.
+            Chosen las file format.
         """
         columns = set(point_cloud.columns).difference(["x", "y", "z"])
-        best_format = 0
+        best_format = laspy.point.format.PointFormat(0)
         covered_columns = 0
 
         for f in LasWriter._supported_las_formats:
@@ -72,6 +72,46 @@ def _select_point_format(self, point_cloud: pd.DataFrame) -> int:
 
         return best_format
 
+    @classmethod
+    def _compute_scales_and_offsets(
+        cls,
+        xyz: np.ndarray,
+        default_scales: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Computes LAS scales and offsets for storing coordinates in LAS integer fields.
+
+        The returned offsets are centered on the coordinate range to maximize the usable signed 32-bit integer span.
+        The requested scales are preserved unless they would cause the coordinate range to exceed that span, in which
+        case the affected scales are increased to the minimum supported values.
+
+        Args:
+            xyz: Point coordinates.
+            default_scales: Preferred LAS scales for the x, y, and z coordinates.
+
+        Returns:
+            Tuple containing the LAS scales and offsets for the x, y, and z coordinates.
+
+        Shape:
+            - :code:`xyz`: :math:`(num_points, 3)`
+        """
+
+        if len(xyz) == 0:
+            return default_scales, np.zeros(3, dtype=np.float64)
+
+        lower_bounds = xyz.min(axis=0)
+        upper_bounds = xyz.max(axis=0)
+
+        offsets = (lower_bounds + upper_bounds) / 2
+
+        coordinate_spans = upper_bounds - lower_bounds
+        minimum_supported_scales = coordinate_spans / (2 * np.iinfo(np.int32).max)
+        scales = np.maximum(default_scales, minimum_supported_scales)
+
+        offsets = np.round(offsets / scales) * scales
+
+        return scales, offsets
+
     def _write_data(  # pylint: disable=too-many-locals
         self,
         point_cloud: pd.DataFrame,
@@ -101,17 +141,18 @@ def _write_data(  # pylint: disable=too-many-locals
         ):
             point_cloud = point_cloud.rename({"r": "red", "g": "green", "b": "blue"}, axis=1)
 
-        las_data = laspy.create(point_format=self._select_point_format(point_cloud))
+        las_format = self._select_point_format(point_cloud)
+        las_data = laspy.create(point_format=las_format)
         point_coords = point_cloud[["x", "y", "z"]].values
-        offsets = point_coords.min(axis=0)
-        scales = [self.maximum_resolution] * 3
+        scales = np.array([self.maximum_resolution] * 3)
         if x_max_resolution is not None:
             scales[0] = x_max_resolution
         if y_max_resolution is not None:
             scales[1] = y_max_resolution
         if z_max_resolution is not None:
             scales[2] = z_max_resolution
 
+        scales, offsets = self._compute_scales_and_offsets(point_coords, scales)
         las_data.change_scaling(scales=scales, offsets=offsets)
         las_data.xyz = point_coords
 
@@ -121,12 +162,16 @@ def _write_data(  # pylint: disable=too-many-locals
             if column_name.lower() in ["x", "y", "z"]:
                 continue
 
+            column_format = las_format[column_name]
+
             if column_name.lower() in point_cloud.columns:
-                las_data[column_name] = point_cloud[column_name.lower()]
+                las_data[column_name] = point_cloud[column_name.lower()].to_numpy(dtype=column_format.dtype)
                 extra_columns.remove(column_name)
             else:
                 default_value = LasWriter._standard_field_defaults.get(column_name, 0)
-                las_data[column_name] = np.full_like(las_data[column_name], fill_value=default_value)
+                las_data[column_name] = np.full_like(
+                    las_data[column_name], fill_value=default_value, dtype=column_format.dtype
+                )
 
         extra_dims = []
         for column_name in extra_columns:
@@ -138,7 +183,7 @@ def _write_data(  # pylint: disable=too-many-locals
         las_data.update_header()
 
         for extra_dim in extra_dims:
-            las_data[extra_dim.name] = point_cloud[extra_dim.name]
+            las_data[extra_dim.name] = point_cloud[extra_dim.name].to_numpy(dtype=extra_dim.type)
 
         if crs is not None:
             las_data.header.add_crs(CRS.from_string(crs))

test/io/test_las_writer.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import laspy
 
 from pointtorch.io import LasWriter, LasReader, PointCloudIoData
 
@@ -127,7 +128,7 @@ def test_write_max_resolutions(
 
         read_point_cloud_data = las_reader.read(file_path)
 
-        assert (point_cloud_df.to_numpy() == read_point_cloud_data.data.to_numpy()).all()
+        assert np.allclose(point_cloud_df.to_numpy(), read_point_cloud_data.data.to_numpy())
 
     @pytest.mark.parametrize("file_format", ["las", "laz"])
     @pytest.mark.parametrize("use_pathlib", [True, False])
@@ -148,3 +149,63 @@ def test_write_crs(
         read_point_cloud_data = las_reader.read(file_path)
 
         assert expected_crs == read_point_cloud_data.crs
+
+    @pytest.mark.parametrize("file_format", ["las", "laz"])
+    def test_offsets_for_wide_coordinate_ranges(
+        self, las_reader: LasReader, las_writer: LasWriter, cache_dir: str, file_format: str
+    ):
+        point_cloud_df = pd.DataFrame([[0.0, 0.0, 0.0], [3000.0, 10.0, 5.0]], columns=["x", "y", "z"])
+        point_cloud_data = PointCloudIoData(point_cloud_df)
+        file_path = os.path.join(cache_dir, f"wide_range.{file_format}")
+
+        las_writer.write(point_cloud_data, file_path)
+
+        with laspy.open(file_path) as las_file:
+            assert las_file.header.offsets[0] == pytest.approx(1500.0)
+            assert las_file.header.offsets[1] == pytest.approx(5.0)
+            assert las_file.header.offsets[2] == pytest.approx(2.5)
+
+        read_point_cloud_data = las_reader.read(file_path)
+        assert np.allclose(point_cloud_df.to_numpy(), read_point_cloud_data.data.to_numpy())
+
+    @pytest.mark.parametrize("file_format", ["las", "laz"])
+    def test_relaxes_resolution_to_prevent_overflow(
+        self, las_reader: LasReader, las_writer: LasWriter, cache_dir: str, file_format: str
+    ):
+        point_cloud_df = pd.DataFrame([[0.0, 0.0, 0.0], [50000.0, 0.0, 0.0]], columns=["x", "y", "z"])
+        point_cloud_data = PointCloudIoData(point_cloud_df, x_max_resolution=1e-6)
+        file_path = os.path.join(cache_dir, f"very_wide_range.{file_format}")
+
+        las_writer.write(point_cloud_data, file_path)
+
+        read_point_cloud_data = las_reader.read(file_path)
+
+        assert read_point_cloud_data.x_max_resolution is not None and read_point_cloud_data.x_max_resolution > 1e-6
+        assert np.allclose(point_cloud_df.to_numpy(), read_point_cloud_data.data.to_numpy())
+
+    @pytest.mark.parametrize("file_format", ["las", "laz"])
+    def test_write_empty_point_cloud(
+        self, las_reader: LasReader, las_writer: LasWriter, cache_dir: str, file_format: str
+    ):
+        x_max_resolution = 0.1
+        y_max_resolution = 0.01
+        z_max_resolution = 1.0
+
+        point_cloud_df = pd.DataFrame(columns=["x", "y", "z"], dtype=np.float64)
+        point_cloud_data = PointCloudIoData(
+            point_cloud_df,
+            x_max_resolution=x_max_resolution,
+            y_max_resolution=y_max_resolution,
+            z_max_resolution=z_max_resolution,
+        )
+        file_path = os.path.join(cache_dir, f"empty_point_cloud.{file_format}")
+
+        las_writer.write(point_cloud_data, file_path)
+
+        read_point_cloud_data = las_reader.read(file_path)
+
+        assert len(read_point_cloud_data.data) == 0
+
+        assert read_point_cloud_data.x_max_resolution == x_max_resolution
+        assert read_point_cloud_data.y_max_resolution == y_max_resolution
+        assert read_point_cloud_data.z_max_resolution == z_max_resolution

test/io/test_point_cloud_writer.py

@@ -94,7 +94,7 @@ def test_write_max_resolutions(
 
         read_point_cloud_data = point_cloud_reader.read(file_path)
 
-        assert (point_cloud_df.to_numpy() == read_point_cloud_data.data.to_numpy()).all()
+        assert np.allclose(point_cloud_df.to_numpy(), read_point_cloud_data.data.to_numpy())
 
     @pytest.mark.parametrize("file_format", ["h5", "hdf", "las", "laz"])
     @pytest.mark.parametrize("use_pathlib", [True, False])