Staging external contribution integration

.devcontainer/devcontainer.json

@@ -0,0 +1,16 @@
+{
+  "name": "Myria3D",
+  "build": {
+    "dockerfile": "../Dockerfile",
+    "context": ".."
+  },
+  "settings": {},
+  "extensions": [
+    "ms-python.python",
+    "ms-azuretools.vscode-docker"
+  ],
+  "runArgs": ["--gpus", "all", "--shm-size=8g"],
+  "remoteUser": "root"
+}
+
+

CHANGELOG.md

@@ -1,4 +1,5 @@
 # CHANGELOG
+- Fix lidar_hd_pre_transform to support missing RGB channels #138 (contrib from @CEZERT)
 
 - Add a github action workflow to run a trained model on the lidar-prod thresholds optimisation dataset
 (in order to automate thresholds optimization)

myria3d/pctl/points_pre_transform/lidar_hd.py

@@ -1,5 +1,5 @@
-# function to turn points loaded via pdal into a pyg Data object, with additional channels
 import numpy as np
+from numpy.lib.recfunctions import append_fields
 from torch_geometric.data import Data
 
 COLORS_NORMALIZATION_MAX_VALUE = 255.0 * 256.0
@@ -9,69 +9,79 @@
 def lidar_hd_pre_transform(points):
     """Turn pdal points into torch-geometric Data object.
 
-    Builds a composite (average) color channel on the fly.     Calculate NDVI on the fly.
+    Builds a composite (average) color channel on the fly. Calculate NDVI on the fly.
 
     Args:
-        las_filepath (str): path to the LAS file.
+        points (np.ndarray): points loaded via PDAL
 
     Returns:
-        Data: the point cloud formatted for later deep learning training.
-
+        Data: the point cloud formatted for deep learning training.
     """
-    # Positions and base features
+    # Positions
     pos = np.asarray([points["X"], points["Y"], points["Z"]], dtype=np.float32).transpose()
+
     # normalization
     occluded_points = points["ReturnNumber"] > 1
 
-    points["ReturnNumber"] = (points["ReturnNumber"]) / (RETURN_NUMBER_NORMALIZATION_MAX_VALUE)
-    points["NumberOfReturns"] = (points["NumberOfReturns"]) / (
-        RETURN_NUMBER_NORMALIZATION_MAX_VALUE
-    )
+    points["ReturnNumber"] = (points["ReturnNumber"]) / RETURN_NUMBER_NORMALIZATION_MAX_VALUE
+    points["NumberOfReturns"] = (points["NumberOfReturns"]) / RETURN_NUMBER_NORMALIZATION_MAX_VALUE
+
+    # Ensure all color fields exist, even if missing (filled with 0)
+    for color in ["Red", "Green", "Blue", "Infrared"]:
+        if color not in points.dtype.names:
+            print(f"Color channel {color} not found. Creating fake {color} filled with 0.")
+            fake_color = np.zeros(points.shape[0], dtype=np.float32)
+            points = append_fields(points, color, fake_color, dtypes=np.float32, usemask=False)
 
+    # Normalize colors if available
+    available_colors = []
     for color in ["Red", "Green", "Blue", "Infrared"]:
-        assert points[color].max() <= COLORS_NORMALIZATION_MAX_VALUE
-        points[color][:] = points[color] / COLORS_NORMALIZATION_MAX_VALUE
-        points[color][occluded_points] = 0.0
-
-    # Additional features :
-    # Average color, that will be normalized on the fly based on single-sample
-    rgb_avg = (
-        np.asarray([points["Red"], points["Green"], points["Blue"]], dtype=np.float32)
-        .transpose()
-        .mean(axis=1)
-    )
+        if color in points.dtype.names:
+            assert points[color].max() <= COLORS_NORMALIZATION_MAX_VALUE, f"{color} max too high!"
+            points[color][:] = points[color] / COLORS_NORMALIZATION_MAX_VALUE
+            points[color][occluded_points] = 0.0
+            available_colors.append(color)
+        else:
+            print(f"Warning: {color} channel not found. Skipping.")
+
+    # Additional features
+    rgb_avg = np.zeros(points.shape[0], dtype=np.float32)
+    if all(c in points.dtype.names for c in ["Red", "Green", "Blue"]):
+        rgb_avg = (
+            np.asarray([points["Red"], points["Green"], points["Blue"]], dtype=np.float32)
+            .transpose()
+            .mean(axis=1)
+        )
 
     # NDVI
-    ndvi = (points["Infrared"] - points["Red"]) / (points["Infrared"] + points["Red"] + 10**-6)
-
-    # todo
-    x = np.stack(
-        [
-            points[name]
-            for name in [
-                "Intensity",
-                "ReturnNumber",
-                "NumberOfReturns",
-                "Red",
-                "Green",
-                "Blue",
-                "Infrared",
-            ]
-        ]
-        + [rgb_avg, ndvi],
-        axis=0,
-    ).transpose()
+    ndvi = np.zeros(points.shape[0], dtype=np.float32)
+    if "Infrared" in points.dtype.names and "Red" in points.dtype.names:
+        ndvi = (points["Infrared"] - points["Red"]) / (points["Infrared"] + points["Red"] + 1e-6)
+
+    # Features list: dynamically adapt based on what exists
+    x_list = [
+        points["Intensity"],
+        points["ReturnNumber"],
+        points["NumberOfReturns"],
+    ]
+
     x_features_names = [
         "Intensity",
         "ReturnNumber",
         "NumberOfReturns",
-        "Red",
-        "Green",
-        "Blue",
-        "Infrared",
-        "rgb_avg",
-        "ndvi",
     ]
+
+    for color in ["Red", "Green", "Blue", "Infrared"]:
+        if color in points.dtype.names:
+            x_list.append(points[color])
+            x_features_names.append(color)
+
+    # Always add computed features
+    x_list += [rgb_avg, ndvi]
+    x_features_names += ["rgb_avg", "ndvi"]
+
+    x = np.stack(x_list, axis=0).transpose()
+
     y = points["Classification"]
 
     data = Data(pos=pos, x=x, y=y, x_features_names=x_features_names)