refactor: Update get_images_by_label function to return indices of th…

…e selected images
siapy · Sep 2, 2024 · 6818ded · 6818ded
1 parent 1a81437
commit 6818ded
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 8 deletions.
diff --git a/source/helpers/helpers.py b/source/helpers/helpers.py
@@ -59,7 +59,7 @@ def get_images_by_label(
     image_set_cam2: list[SpectralImage],
     labels_cam1: list[list[str]],
     labels_cam2: list[list[str]],
-) -> tuple[SpectralImage, SpectralImage]:
+) -> tuple[SpectralImage, SpectralImage, int, int]:
     index_cam1 = -1
     index_cam2 = -1
     for idx, labels_list in enumerate(labels_cam1):
@@ -82,4 +82,4 @@ def get_images_by_label(
     logger.info(f"Extracted images for label: '{label}'")
     logger.info(f"Camera1 image path: '{image_cam1.filepath}'")
     logger.info(f"Camera2 image path: '{image_cam2.filepath}'")
-    return image_cam1, image_cam2
+    return image_cam1, image_cam2, index_cam1, index_cam2
diff --git a/source/misc/display_image.py b/source/misc/display_image.py
@@ -30,7 +30,7 @@ def display_spectral_image(label: Optional[str] = None):
         labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
             image_set_cam1, image_set_cam2
         )
-        image_cam1, image_cam2 = get_images_by_label(
+        image_cam1, image_cam2, _, _ = get_images_by_label(
             label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
         )
 
@@ -49,7 +49,7 @@ def display_spectral_images_with_areas(
     labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
         image_set_cam1, image_set_cam2
     )
-    image_cam1, image_cam2 = get_images_by_label(
+    image_cam1, image_cam2, _, _ = get_images_by_label(
         label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
     )
 

diff --git a/source/processing/model.py b/source/processing/model.py
@@ -26,7 +26,7 @@ def convert_selected_areas_to_train_data(
     y_cam2 = []
     for category, data in selected_areas.items():
         for label, pixels_list in data.items():
-            image_cam1, image_cam2 = get_images_by_label(
+            image_cam1, image_cam2, _, _ = get_images_by_label(
                 label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
             )
             for pixels_cam1 in pixels_list:

diff --git a/source/processing/selector.py b/source/processing/selector.py
@@ -16,7 +16,7 @@ def select_areas_on_images(label: str) -> list[Pixels]:
     labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
         image_set_cam1, image_set_cam2
     )
-    image_cam1, _ = get_images_by_label(
+    image_cam1, _, _, _ = get_images_by_label(
         label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
     )
     selected_areas = pixels_select_lasso(image_cam1)

diff --git a/source/processing/transformator.py b/source/processing/transformator.py
@@ -21,7 +21,7 @@ def find_transformation_between_images(label: str) -> np.ndarray:
     labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
         image_set_cam1, image_set_cam2
     )
-    image_cam1, image_cam2 = get_images_by_label(
+    image_cam1, image_cam2, _, _ = get_images_by_label(
         label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
     )
 
@@ -32,4 +32,3 @@ def find_transformation_between_images(label: str) -> np.ndarray:
     logger.info("Transformation matrix calculated")
     # pixels_transformed = corregistrator.transform(pixels_cam1, matx)
     return matx
-