refactor: Train model for both cameras and save artifacts

main.py

@@ -59,9 +59,13 @@ def select_areas(label: str, category: str):
 @app.command()
 def train_model():
     selected_areas = load_all_selected_areas()
-    X, y = convert_selected_areas_to_train_data(selected_areas)
-    encoder, model = train_xgboost_model(X, y)
-    save_model(encoder, model)
+    matx = load_transformation_matrix()
+    X_cam1, y_cam1, X_cam2, y_cam2 = convert_selected_areas_to_train_data(
+        selected_areas, matx
+    )
+    encoder_cam1, model_cam1 = train_xgboost_model(X_cam1, y_cam1)
+    encoder_cam2, model_cam2 = train_xgboost_model(X_cam2, y_cam2)
+    save_model(encoder_cam1, model_cam1, encoder_cam2, model_cam2)
 
 
 if __name__ == "__main__":

source/helpers/artifacts.py

@@ -12,8 +12,10 @@
 _TRANSFORMATION_MATX_FILENAME = settings.artifacts_dir / "transform/matx.pkl"
 _SELECTED_AREAS_DIR = settings.artifacts_dir / "areas"
 _MODEL_DIR = settings.artifacts_dir / "model"
-_MODEL_CLF_FILENAME = _MODEL_DIR / "model.json"
-_MODEL_ENCODER_FILENAME = _MODEL_DIR / "encoder.pkl"
+_MODEL_CLF_FILENAMEM_CAM1 = _MODEL_DIR / "model_cam1.json"
+_MODEL_ENCODER_FILENAME_CAM1 = _MODEL_DIR / "encoder_cam1.pkl"
+_MODEL_CLF_FILENAMEM_CAM2 = _MODEL_DIR / "model_cam2.json"
+_MODEL_ENCODER_FILENAME_CAM2 = _MODEL_DIR / "encoder_cam2.pkl"
 
 
 def save_transformation_matrix(matx: np.ndarray):
@@ -74,25 +76,53 @@ def load_all_selected_areas() -> dict[str, dict[str, list[Pixels]]]:
     }
 
 
-def save_model(encoder: LabelEncoder, model: XGBClassifier):
+def save_model(
+    encoder_cam1: LabelEncoder,
+    model_cam1: XGBClassifier,
+    encoder_cam2: LabelEncoder,
+    model_cam2: XGBClassifier,
+):
     _MODEL_DIR.mkdir(parents=True, exist_ok=True)
-    model.save_model(_MODEL_CLF_FILENAME)
-    with open(_MODEL_ENCODER_FILENAME, "wb") as f:
-        pickle.dump(encoder, f)
-    logger.info("Model saved.")
+    model_cam1.save_model(_MODEL_CLF_FILENAMEM_CAM1)
+    with open(_MODEL_ENCODER_FILENAME_CAM1, "wb") as f:
+        pickle.dump(encoder_cam1, f)
+    logger.info("Model saved for camera 1.")
+
+    model_cam2.save_model(_MODEL_CLF_FILENAMEM_CAM2)
+    with open(_MODEL_ENCODER_FILENAME_CAM2, "wb") as f:
+        pickle.dump(encoder_cam2, f)
+    logger.info("Model saved for camera 2.")
 
     # load_model()
 
 
-def load_model() -> tuple[LabelEncoder, XGBClassifier]:
-    if not _MODEL_CLF_FILENAME.exists() or not _MODEL_ENCODER_FILENAME.exists():
-        raise FileNotFoundError("Model files not found.")
+def load_model() -> tuple[LabelEncoder, XGBClassifier, LabelEncoder, XGBClassifier]:
+    if (
+        not _MODEL_CLF_FILENAMEM_CAM1.exists()
+        or not _MODEL_ENCODER_FILENAME_CAM1.exists()
+    ):
+        raise FileNotFoundError("Model files not found for camera 1.")
+
+    model_cam1 = XGBClassifier()
+    model_cam1.load_model(_MODEL_CLF_FILENAMEM_CAM1)
+
+    with open(_MODEL_ENCODER_FILENAME_CAM1, "rb") as f:
+        encoder_cam1 = pickle.load(f)
+
+    logger.info("Model loaded for camera 1.")
+
+    if (
+        not _MODEL_CLF_FILENAMEM_CAM2.exists()
+        or not _MODEL_ENCODER_FILENAME_CAM2.exists()
+    ):
+        raise FileNotFoundError("Model files not found for camera 2.")
+
+    model_cam2 = XGBClassifier()
+    model_cam2.load_model(_MODEL_CLF_FILENAMEM_CAM2)
 
-    model = XGBClassifier()
-    model.load_model(_MODEL_CLF_FILENAME)
+    with open(_MODEL_ENCODER_FILENAME_CAM2, "rb") as f:
+        encoder_cam2 = pickle.load(f)
 
-    with open(_MODEL_ENCODER_FILENAME, "rb") as f:
-        encoder = pickle.load(f)
+    logger.info("Model loaded for camera 2.")
 
-    logger.info("Model loaded.")
-    return encoder, model
+    return encoder_cam1, model_cam1, encoder_cam2, model_cam2

source/processing/model.py

@@ -1,5 +1,6 @@
 import numpy as np
 from siapy.entities import Pixels
+from siapy.transformations import corregistrator
 from sklearn.preprocessing import LabelEncoder
 from xgboost import XGBClassifier
 
@@ -13,24 +14,36 @@
 
 def convert_selected_areas_to_train_data(
     selected_areas: dict[str, dict[str, list[Pixels]]],
-) -> tuple[list[np.ndarray], list[str]]:
+    transformation_matx: np.ndarray,
+) -> tuple[list[np.ndarray], list[str], list[np.ndarray], list[str]]:
     image_set_cam1, image_set_cam2 = read_spectral_images()
     labels_cam1, labels_cam2 = extract_labels_from_spectral_images(
         image_set_cam1, image_set_cam2
     )
-    X = []
-    y = []
+    X_cam1 = []
+    y_cam1 = []
+    X_cam2 = []
+    y_cam2 = []
     for category, data in selected_areas.items():
         for label, pixels_list in data.items():
-            image_cam1, _ = 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 in pixels_list:
-                siagnatures = image_cam1.to_signatures(pixels)
-                signal_mean = siagnatures.signals.mean()
-                X.append(signal_mean)
-                y.append(category)
-    return X, y
+            for pixels_cam1 in pixels_list:
+                # Extract data for camera 1
+                signatures = image_cam1.to_signatures(pixels_cam1)
+                signal_mean = signatures.signals.mean()
+                X_cam1.append(signal_mean)
+                y_cam1.append(category)
+
+                # Extract data for camera 2
+                pixels_cam2 = corregistrator.transform(pixels_cam1, transformation_matx)
+                signatures = image_cam2.to_signatures(pixels_cam2)
+                signal_mean = signatures.signals.mean()
+                X_cam2.append(signal_mean)
+                y_cam2.append(category)
+
+    return X_cam1, y_cam1, X_cam2, y_cam2
 
 
 def train_xgboost_model(