feat: Add functionality to select areas on images

main.py

@@ -5,9 +5,9 @@
 from pydantic.json import pydantic_encoder
 
 from source.core import logger, settings
-from source.helpers import save_transformation_matrix
+from source.helpers import save_selected_areas, save_transformation_matrix
 from source.misc import check_spectral_images, display_spectral_image
-from source.processing import find_transformation_between_images
+from source.processing import find_transformation_between_images, select_areas_on_images
 
 app = typer.Typer()
 
@@ -33,6 +33,12 @@ def calculate_transformation(label: str):
     save_transformation_matrix(matx)
 
 
+@app.command()
+def select_areas(label: str, category: str):
+    selected_areas = select_areas_on_images(label)
+    save_selected_areas(selected_areas, category, label)
+
+
 if __name__ == "__main__":
     logger.info(f"Project name: '{settings.project_name}'")
     app()

source/helpers/__init__.py

@@ -1,4 +1,9 @@
-from .artifacts import load_transformation_matrix, save_transformation_matrix
+from .artifacts import (
+    load_selected_areas,
+    load_transformation_matrix,
+    save_selected_areas,
+    save_transformation_matrix,
+)
 from .helpers import (
     extract_labels_from_spectral_image,
     extract_labels_from_spectral_images,
@@ -13,4 +18,6 @@
     "save_transformation_matrix",
     "load_transformation_matrix",
     "get_images_by_label",
+    "save_selected_areas",
+    "load_selected_areas",
 ]

source/helpers/artifacts.py

@@ -1,10 +1,14 @@
+import os
 import pickle
+import shutil
 
 import numpy as np
+from siapy.entities import Pixels
 
 from source.core import logger, settings
 
 _TRANSFORMATION_MATX_FILENAME = settings.artifacts_dir / "transform/matx.pkl"
+_SELECTED_AREAS_DIR = settings.artifacts_dir / "areas"
 
 
 def save_transformation_matrix(matx: np.ndarray):
@@ -26,3 +30,33 @@ def load_transformation_matrix() -> np.ndarray:
         matx = pickle.load(f)
 
     return matx
+
+
+def save_selected_areas(selected_areas: list[Pixels], category: str, label: str):
+    category_areas_dir = _SELECTED_AREAS_DIR / category / label
+    if os.path.exists(category_areas_dir):
+        shutil.rmtree(category_areas_dir)
+    category_areas_dir.mkdir(parents=True, exist_ok=True)
+    logger.info("Saving selected areas to: %s" % category_areas_dir)
+    for idx, area in enumerate(selected_areas):
+        with open(category_areas_dir / f"{idx}.pkl", "wb") as f:
+            pickle.dump(area.df, f)
+
+
+def load_selected_areas(category: str) -> list[Pixels]:
+    category_areas_dir = _SELECTED_AREAS_DIR / category
+    if not category_areas_dir.exists():
+        raise FileNotFoundError(
+            f"Selected areas directory not found: '{category_areas_dir}'"
+        )
+    logger.info(f"Loading selected areas from: '{category_areas_dir}'")
+
+    selected_areas = []
+    labels = category_areas_dir.glob("*")
+    for label in labels:
+        category_areas_dir_ = category_areas_dir / label.stem
+        for area_file in sorted(category_areas_dir_.glob("*.pkl")):
+            with open(area_file, "rb") as f:
+                area_df = pickle.load(f)
+                selected_areas.append(Pixels(area_df))
+    return selected_areas

source/processing/__init__.py

@@ -1,5 +1,4 @@
+from .selector import select_areas_on_images
 from .transformator import find_transformation_between_images
 
-__all__ = [
-    "find_transformation_between_images",
-]
+__all__ = ["find_transformation_between_images", "select_areas_on_images"]

source/processing/selector.py

@@ -0,0 +1,26 @@
+from siapy.entities import Pixels
+from siapy.utils.plots import (
+    display_image_with_areas,
+    pixels_select_lasso,
+)
+
+from source.core import logger
+from source.helpers import (
+    extract_labels_from_spectral_images,
+    get_images_by_label,
+    read_spectral_images,
+)
+
+
+def select_areas_on_images(label: str) -> list[Pixels]:
+    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, _ = get_images_by_label(
+        label, image_set_cam1, image_set_cam2, labels_cam1, labels_cam2
+    )
+    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