chore: Update display of spectral images with selected areas

main.py

@@ -7,11 +7,16 @@
 from source.core import logger, settings
 from source.helpers import (
     load_all_selected_areas,
+    load_transformation_matrix,
     save_model,
     save_selected_areas,
     save_transformation_matrix,
 )
-from source.misc import check_spectral_images, display_spectral_image
+from source.misc import (
+    check_spectral_images,
+    display_spectral_image,
+    display_spectral_images_with_areas,
+)
 from source.processing import (
     convert_selected_areas_to_train_data,
     find_transformation_between_images,
@@ -46,6 +51,8 @@ def calculate_transformation(label: str):
 @app.command()
 def select_areas(label: str, category: str):
     selected_areas = select_areas_on_images(label)
+    matx = load_transformation_matrix()
+    display_spectral_images_with_areas(label, selected_areas, matx)
     save_selected_areas(selected_areas, category, label)
 
 

source/misc/__init__.py

@@ -1,7 +1,8 @@
 from .check_images import check_spectral_images
-from .display_image import display_spectral_image
+from .display_image import display_spectral_image, display_spectral_images_with_areas
 
 __all__ = [
     "check_spectral_images",
     "display_spectral_image",
+    "display_spectral_images_with_areas",
 ]

source/misc/display_image.py

@@ -1,6 +1,12 @@
 from typing import Optional
 
+import numpy as np
 from matplotlib import pyplot as plt
+from siapy.entities import Pixels, SpectralImage
+from siapy.transformations import corregistrator
+from siapy.utils.plots import (
+    display_image_with_areas,
+)
 
 from source.core import logger, settings
 from source.helpers import (
@@ -32,3 +38,24 @@ def display_spectral_image(label: Optional[str] = None):
     ax[0].imshow(image_cam1.to_display())  # noqa
     ax[1].imshow(image_cam2.to_display())  # noqa
     plt.show()
+
+
+def display_spectral_images_with_areas(
+    label: str,
+    selected_areas_cam1: list[Pixels],
+    transformation_matx: np.ndarray,
+):
+    image_set_cam1, image_set_cam2 = read_spectral_images()
+    labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
+        image_set_cam1, image_set_cam2
+    )
+    image_cam1, image_cam2 = get_images_by_label(
+        label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
+    )
+
+    selected_areas_cam2 = [
+        corregistrator.transform(pixels_cam1, transformation_matx)
+        for pixels_cam1 in selected_areas_cam1
+    ]
+    display_image_with_areas(image_cam1, selected_areas_cam1)
+    display_image_with_areas(image_cam2, selected_areas_cam2)

source/processing/selector.py

@@ -1,6 +1,5 @@
 from siapy.entities import Pixels
 from siapy.utils.plots import (
-    display_image_with_areas,
     pixels_select_lasso,
 )
 
@@ -22,5 +21,4 @@ def select_areas_on_images(label: str) -> list[Pixels]:
     )
     selected_areas = pixels_select_lasso(image_cam1)
     logger.info(f"Selected '{len(selected_areas)}' areas")
-    display_image_with_areas(image_cam1, selected_areas)
     return selected_areas