diff --git a/airborne_lidar/airborne_lidar_inference.py b/airborne_lidar/airborne_lidar_inference.py
index 399b3eb..5309d68 100644
--- a/airborne_lidar/airborne_lidar_inference.py
+++ b/airborne_lidar/airborne_lidar_inference.py
@@ -13,7 +13,7 @@
 from collections import OrderedDict
 from airborne_lidar_seg import get_model, nearest_correspondance, count_parameters, class_mode
 import laspy
-from airborne_lidar_utils import write_features, read_parameters, gpu_stats
+from airborne_lidar_utils import write_features, read_parameters, gpu_stats, write_to_las
 from airborne_lidar_datasets import PartDatasetTest
 
 
@@ -70,25 +70,6 @@ def write_las_to_h5(filename):
         return filename.parent / filename.stem
 
 
-def write_to_las(filename, xyz, pred, header, info_class):
-    """Write xyz and ASPRS predictions to las file format. """
-    # TODO: Write CRS info with file.
-    with laspy.file.File(filename, mode='w', header=header) as out_file:
-        out_file.x = xyz[:, 0]
-        out_file.y = xyz[:, 1]
-        out_file.z = xyz[:, 2]
-        pred = pred_to_asprs(pred, info_class)
-        out_file.classification = pred
-
-
-def pred_to_asprs(pred, info_class):
-    """Converts predicted values (0->n) to the corresponding ASPRS class."""
-    labels2 = np.full(shape=pred.shape, fill_value=0, dtype=int)
-    for key, value in info_class.items():
-        labels2[pred == value['mode']] = int(key)
-    return labels2
-
-
 def test(args, filename, info_class):
     model_folder = Path(args['global']['model_dir'])
     data_dir = Path(args['global']['data_dir'])
diff --git a/airborne_lidar/airborne_lidar_utils.py b/airborne_lidar/airborne_lidar_utils.py
index 7de1427..8525d16 100644
--- a/airborne_lidar/airborne_lidar_utils.py
+++ b/airborne_lidar/airborne_lidar_utils.py
@@ -5,6 +5,7 @@
 import shutil
 import warnings
 import numpy as np
+import laspy
 
 try:
     from pynvml import *
@@ -99,3 +100,22 @@ def wblue(str):
 
 def wgreen(str):
     return bcolors.OKGREEN + str + bcolors.ENDC
+
+
+def write_to_las(filename, xyz, pred, header, info_class):
+    """Write xyz and ASPRS predictions to las file format. """
+    # TODO: Write CRS info with file.
+    with laspy.file.File(filename, mode='w', header=header) as out_file:
+        out_file.x = xyz[:, 0]
+        out_file.y = xyz[:, 1]
+        out_file.z = xyz[:, 2]
+        pred = pred_to_asprs(pred, info_class)
+        out_file.classification = pred
+
+
+def pred_to_asprs(pred, info_class):
+    """Converts predicted values (0->n) to the corresponding ASPRS class."""
+    labels2 = np.full(shape=pred.shape, fill_value=0, dtype=int)
+    for key, value in info_class.items():
+        labels2[pred == value['mode']] = int(key)
+    return labels2
diff --git a/airborne_lidar/prepare_airborne_lidar_label.py b/airborne_lidar/prepare_airborne_lidar_label.py
index 23420e2..c599a30 100644
--- a/airborne_lidar/prepare_airborne_lidar_label.py
+++ b/airborne_lidar/prepare_airborne_lidar_label.py
@@ -88,7 +88,7 @@ def main():
             path_prepare_label = Path(args.dest, dst)
             path_prepare_label.mkdir(exist_ok=True)
 
-            xyzni, label, nb_pts = read_las_format(f"{elem}.las")
+            xyzni, label, nb_pts, header = read_las_format(f"{elem}.las")
             write_features(Path(f"{path_prepare_label / Path(elem).stem}_prepared.hdfs"), xyzni=xyzni, labels=label)
             print(f"File {dst}/{Path(elem).stem} prepared. {nb_pts:,} points written.")