From 622591eaa1354998008e32827c5c18b2c16b1fa1 Mon Sep 17 00:00:00 2001
From: Graham Knapp <32717635+dancergraham@users.noreply.github.com>
Date: Sun, 22 Sep 2024 16:52:42 +0200
Subject: [PATCH] Feature/e57 (#21)

* feature: read e57 files

* feature: read e57 files

* style: fix linting failures

---------

Co-authored-by: Graham Knapp <graham.knapp@acernis.fr>
---
 src/lidar_visualizer/datasets/__init__.py |  1 +
 src/lidar_visualizer/datasets/generic.py  | 55 ++++++++++++++++++++---
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/src/lidar_visualizer/datasets/__init__.py b/src/lidar_visualizer/datasets/__init__.py
index 22cae4e..3bc18d8 100644
--- a/src/lidar_visualizer/datasets/__init__.py
+++ b/src/lidar_visualizer/datasets/__init__.py
@@ -34,6 +34,7 @@ def supported_file_extensions():
         "ctm",
         "off",
         "stl",
+        "e57",
     ]
 
 
diff --git a/src/lidar_visualizer/datasets/generic.py b/src/lidar_visualizer/datasets/generic.py
index bcf0b16..e9fd727 100644
--- a/src/lidar_visualizer/datasets/generic.py
+++ b/src/lidar_visualizer/datasets/generic.py
@@ -77,7 +77,10 @@ def __getitem__(self, idx):
     def _get_point_cloud_reader(self):
         """Attempt to guess with try/catch blocks which is the best point cloud reader to use for
         the given dataset folder. Supported readers so far are:
+
+        File readers are functions which take a filename as an input and return a tuple of points and colors.
             - np.fromfile
+            - pye57
             - open3d
             - trimesh.load
             - PyntCloud
@@ -102,7 +105,49 @@ def read_kitti_scan(file):
         tried_libraries = []
         missing_libraries = []
 
-        # 2. Try with Open3D
+        # 2 Try with pye57
+        if self.file_extension == "e57":
+            try:
+                import pye57
+
+                def read_e57_scan(file):
+                    e57 = pye57.E57(file)
+                    point_data = None
+                    color_data = None
+                    # One e57 file can contain multiple scans, scanned from different positions
+                    for i in range(e57.scan_count):
+                        i = e57.read_scan(i, colors=True, ignore_missing_fields=True)
+                        scan_data = np.stack(
+                            [i["cartesianX"], i["cartesianY"], i["cartesianZ"]], axis=1
+                        )
+                        point_data = (
+                            np.concat([point_data, scan_data])
+                            if point_data is not None
+                            else scan_data
+                        )
+                        try:
+                            scan_color_data = np.stack(
+                                [i["colorRed"], i["colorGreen"], i["colorBlue"]], axis=1
+                            )
+                            color_data = (
+                                np.concat([color_data, scan_color_data])
+                                if color_data is not None
+                                else scan_color_data
+                            )
+                        except KeyError:
+                            pass
+                    # e57 file colors are in 0-255 range
+                    color_data = color_data / 255.0 if color_data is not None else None
+                    return point_data, color_data
+
+                return read_e57_scan
+            except ModuleNotFoundError:
+                missing_libraries.append("pye57")
+                print("[WARNING] pye57 not installed")
+            except:
+                tried_libraries.append("pye57")
+
+        # 3. Try with Open3D
         try:
             self.o3d.t.io.read_point_cloud(first_scan_file)
 
@@ -129,7 +174,7 @@ def read_scan_with_intensities(file):
         except:
             tried_libraries.append("open3d")
 
-        # 3. Try with trimesh
+        # 4. Try with trimesh
         try:
             import trimesh
 
@@ -140,7 +185,7 @@ def read_scan_with_intensities(file):
         except:
             tried_libraries.append("trimesh")
 
-        # 4. Try with PyntCloud
+        # 5. Try with PyntCloud
         try:
             from pyntcloud import PyntCloud
 
@@ -151,10 +196,10 @@ def read_scan_with_intensities(file):
         except:
             tried_libraries.append("pyntcloud")
 
-        # If reach this point means that none of the librares exist/could read the file
+        # If reach this point means that none of the libraries exist/could read the file
         if not tried_libraries:
             print(
-                "No 3D library is insallted in your system. Install one of the following "
+                "No 3D library is installed in your system. Install one of the following "
                 "to read the pointclouds"
             )
             print("\n".join(missing_libraries))