fix readme and documentation

README.md

@@ -24,7 +24,7 @@ from aces.model_trainer import ModelTrainer
 
 if __name__ == "__main__":
     config_file = "config.env"
-    config = Config()
+    config = Config(config_file)
     trainer = ModelTrainer(config)
     trainer.train_model()
 ```

aces/config.py

@@ -96,7 +96,13 @@ class Config:
     """
 
     def __init__(self, config_file, override=False) -> None:
+        """
+        ACES Configuration Class Constructor
 
+        Args:
+            config_file (str): The path to the configuration file.
+            override (bool): Flag to override the configuration settings.
+        """
         load_dotenv(config_file, override=override)
 
         self.BASEDIR = Path(os.getenv("BASEDIR"))

aces/ee_utils.py

@@ -39,10 +39,14 @@ def get_credentials_by_service_account_key(key):
     def initialize_session(use_highvolume : bool = False, key : Union[str, None] = None, project: str = None):
         """
         Initialize the Earth Engine session.
+        If use_highvolume is True, the high-volume Earth Engine API will be used.
+        If a project is provided, the session will be initialized with the project ID. Recommended to use project.
+        If a key is provided, the service account key will be used.
 
         Parameters:
         use_highvolume (bool): Whether to use the high-volume Earth Engine API.
         key (str or None): The path to the service account key JSON file. If None, the default credentials will be used.
+        project (str): The Google Cloud project ID to use for the session.
         """
         if key is None:
             if use_highvolume and project:

aces/model_trainer.py

@@ -365,13 +365,15 @@ def save_plots(self) -> None:
         Saves the model architecture plot, training history plot, and model object.
         """
         print(f"Saving plots and model visualization at {self.config.MODEL_SAVE_DIR}...")
+
+        Utils.plot_metrics([key.replace("val_", "") for key in self.history.history.keys() if key.startswith("val_")],
+                           self.history.history, len(self.history.epoch), self.config.MODEL_SAVE_DIR)
+
         if self.config.USE_AI_PLATFORM:
             keras.utils.plot_model(self._model, f"{self.config.MODEL_SAVE_DIR}/model.png", show_shapes=True, rankdir="TB")
             keras.utils.plot_model(self.model, f"{self.config.MODEL_SAVE_DIR}/wrapped_model.png", show_shapes=True, rankdir="LR") # rankdir='TB'
         else:
             keras.utils.plot_model(self.model, f"{self.config.MODEL_SAVE_DIR}/model.png", show_shapes=True, rankdir="TB") # rankdir='TB'
-        Utils.plot_metrics([key.replace("val_", "") for key in self.history.history.keys() if key.startswith("val_")],
-                           self.history.history, len(self.history.epoch), self.config.MODEL_SAVE_DIR)
 
     def save_history_object(self) -> None:
         """

notebook/count_sample_size.ipynb

@@ -31,6 +31,8 @@
       "source": [
         "# Rice mapping in Bhutan with U-Net using high resolution satellite imagery\n",
         "\n",
+        "### This notebook shows an example of counting the sample size from the `tfrecords`\n",
+        "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/SERVIR/servir-aces/blob/main/notebooks/count_sample_size.ipynb\">\n",
@@ -45,9 +47,7 @@
         "</table>\n",
         "</br>\n",
         "</br>\n",
-        "</br>\n",
-        "\n",
-        "### This notebook shows an example of counting the sample size from the `tfrecords`"
+        "</br>"
       ]
     },
     {

notebook/prediction_dnn.ipynb

@@ -31,6 +31,8 @@
       "source": [
         "# Rice mapping in Bhutan with U-Net using high resolution satellite imagery\n",
         "\n",
+        "### This notebook shows an example of using the TFRecord images for prediction using saved DNN Model.\n",
+        "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/SERVIR/servir-aces/blob/main/notebooks/prediction_dnn.ipynb\">\n",
@@ -45,9 +47,7 @@
         "</table>\n",
         "</br>\n",
         "</br>\n",
-        "</br>\n",
-        "\n",
-        "### This notebook shows an example of using the TFRecord images for prediction using saved DNN Model."
+        "</br>"
       ]
     },
     {