From 0c29e7b8e0ab69365dc5091b8406a2bac39a26f0 Mon Sep 17 00:00:00 2001
From: janezlapajne <lapajne.janez@gmail.com>
Date: Sun, 8 Sep 2024 23:47:51 +0200
Subject: [PATCH] refactor: Update read_spectral_images function to accept
 custom images directory

---
 source/helpers/helpers.py        |  8 ++++++--
 source/processing/segmentator.py | 15 +++++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/source/helpers/helpers.py b/source/helpers/helpers.py
index 033557a..7754d9e 100644
--- a/source/helpers/helpers.py
+++ b/source/helpers/helpers.py
@@ -5,9 +5,13 @@
 from source.core import logger, settings
 
 
-def read_spectral_images() -> tuple[list[SpectralImage], list[SpectralImage]]:
+def read_spectral_images(
+    images_dir: Path | str | None = None,
+) -> tuple[list[SpectralImage], list[SpectralImage]]:
+    if images_dir is None:
+        images_dir = settings.images_dir
     file_paths = sorted(
-        [Path(file) for file in Path(settings.images_dir).glob("*") if file.is_file()]
+        [Path(file) for file in Path(images_dir).glob("*") if file.is_file()]
     )
 
     header_paths = [
diff --git a/source/processing/segmentator.py b/source/processing/segmentator.py
index 398e3a2..bb4224a 100644
--- a/source/processing/segmentator.py
+++ b/source/processing/segmentator.py
@@ -90,8 +90,18 @@ def save_segmented_images(
 ):
     filename_cam1 = image_cam1.filepath.stem
     filename_cam2 = image_cam2.filepath.stem
+    meta_keys_to_extract = [
+        "description",
+        "data type",
+        "data ignore value",
+        "interleave",
+        "default bands",
+        "byte order",
+        "wavelength",
+    ]
 
     logger.info("Saving segmented images for Camera 1 ...")
+    metadata_cam1 = {key: image_cam1.metadata[key] for key in meta_keys_to_extract}
     for idx, area in enumerate(selected_areas_cam1):
         filename_widx = (
             f"{index}{settings.labels_between_deliminator}"
@@ -99,9 +109,10 @@ def save_segmented_images(
             f"{filename_cam1}.hdr"
         )
         subarray = image_cam1.to_subarray(area)
-        save_spectral_image(subarray, filename_widx)
+        save_spectral_image(subarray, filename_widx, metadata_cam1)
 
     logger.info("Saving segmented images for Camera 2 ...")
+    metadata_cam2 = {key: image_cam2.metadata[key] for key in meta_keys_to_extract}
     for idx, area in enumerate(selected_areas_cam2):
         filename_widx = (
             f"{index}{settings.labels_between_deliminator}"
@@ -109,7 +120,7 @@ def save_segmented_images(
             f"{filename_cam2}.hdr"
         )
         subarray = image_cam2.to_subarray(area)
-        save_spectral_image(subarray, filename_widx)
+        save_spectral_image(subarray, filename_widx, metadata_cam2)
 
 
 def perform_segmentation(