Revert "remove unnecessary for loop, adjust tests, change default thr…

…eshold" This reverts commit 3baf219, reversing changes made to 8cc046c.
gqcpm · Jul 17, 2024 · 61fe572 · 61fe572
1 parent 3baf219
commit 61fe572
Show file tree

Hide file tree

Showing 31 changed files with 17 additions and 114 deletions.
diff --git a/docs/guides/cli.md b/docs/guides/cli.md
@@ -36,8 +36,8 @@ optional arguments:
 
 ```none
 usage: sleap-train [-h] [--video-paths VIDEO_PATHS] [--val_labels VAL_LABELS]
-                   [--test_labels TEST_LABELS] [--tensorboard] [--save_viz] 
-                   [--keep_viz] [--zmq] [--run_name RUN_NAME] [--prefix PREFIX]
+                   [--test_labels TEST_LABELS] [--tensorboard] [--save_viz]
+                   [--zmq] [--run_name RUN_NAME] [--prefix PREFIX]
                    [--suffix SUFFIX]
                    training_job_path [labels_path]
 
@@ -68,8 +68,6 @@ optional arguments:
   --save_viz            Enable saving of prediction visualizations to the run
                         folder if not already specified in the training job
                         config.
-  --keep_viz            Keep prediction visualization images in the run
-                        folder after training if --save_viz is enabled.
   --zmq                 Enable ZMQ logging (for GUI) if not already specified
                         in the training job config.
   --run_name RUN_NAME   Run name to use when saving file, overrides other run

diff --git a/docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb b/docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb
@@ -335,7 +335,7 @@
       "        \"runs_folder\": \"models\",\n",
       "        \"tags\": [],\n",
       "        \"save_visualizations\": true,\n",
-      "        \"keep_viz_images\": true,\n",
+      "        \"delete_viz_images\": true,\n",
       "        \"zip_outputs\": false,\n",
       "        \"log_to_csv\": true,\n",
       "        \"checkpointing\": {\n",
@@ -727,7 +727,7 @@
       "        \"runs_folder\": \"models\",\n",
       "        \"tags\": [],\n",
       "        \"save_visualizations\": true,\n",
-      "        \"keep_viz_images\": true,\n",
+      "        \"delete_viz_images\": true,\n",
       "        \"zip_outputs\": false,\n",
       "        \"log_to_csv\": true,\n",
       "        \"checkpointing\": {\n",

diff --git a/sleap/config/pipeline_form.yaml b/sleap/config/pipeline_form.yaml
@@ -286,11 +286,6 @@ training:
   type: bool
   default: true
 
-- name: _keep_viz
-  label: Keep Prediction Visualization Images After Training
-  type: bool
-  default: false
-
 - name: _predict_frames
   label: Predict On
   type: list

diff --git a/sleap/gui/learning/runners.py b/sleap/gui/learning/runners.py
@@ -1,5 +1,4 @@
 """Run training/inference in background process via CLI."""
-
 import abc
 import attr
 import os
@@ -501,11 +500,9 @@ def write_pipeline_files(
                 "data_path": os.path.basename(data_path),
                 "models": [Path(p).as_posix() for p in new_cfg_filenames],
                 "output_path": prediction_output_path,
-                "type": (
-                    "labels"
-                    if type(item_for_inference) == DatasetItemForInference
-                    else "video"
-                ),
+                "type": "labels"
+                if type(item_for_inference) == DatasetItemForInference
+                else "video",
                 "only_suggested_frames": only_suggested_frames,
                 "tracking": tracking_args,
             }
@@ -547,7 +544,6 @@ def run_learning_pipeline(
     """
 
     save_viz = inference_params.get("_save_viz", False)
-    keep_viz = inference_params.get("_keep_viz", False)
 
     if "movenet" in inference_params["_pipeline"]:
         trained_job_paths = [inference_params["_pipeline"]]
@@ -561,7 +557,6 @@ def run_learning_pipeline(
             inference_params=inference_params,
             gui=True,
             save_viz=save_viz,
-            keep_viz=keep_viz,
         )
 
         # Check that all the models were trained
@@ -590,7 +585,6 @@ def run_gui_training(
     inference_params: Dict[str, Any],
     gui: bool = True,
     save_viz: bool = False,
-    keep_viz: bool = False,
 ) -> Dict[Text, Text]:
     """
     Runs training for each training job.
@@ -600,7 +594,6 @@ def run_gui_training(
         config_info_list: List of ConfigFileInfo with configs for training.
         gui: Whether to show gui windows and process gui events.
         save_viz: Whether to save visualizations from training.
-        keep_viz: Whether to keep prediction visualization images after training.
 
     Returns:
         Dictionary, keys are head name, values are path to trained config.
@@ -690,7 +683,6 @@ def waiting():
                 video_paths=video_path_list,
                 waiting_callback=waiting,
                 save_viz=save_viz,
-                keep_viz=keep_viz,
             )
 
             if ret == "success":
@@ -833,7 +825,6 @@ def train_subprocess(
     video_paths: Optional[List[Text]] = None,
     waiting_callback: Optional[Callable] = None,
     save_viz: bool = False,
-    keep_viz: bool = False,
 ):
     """Runs training inside subprocess."""
     run_path = job_config.outputs.run_path
@@ -862,8 +853,6 @@ def train_subprocess(
 
         if save_viz:
             cli_args.append("--save_viz")
-        if keep_viz:
-            cli_args.append("--keep_viz")
 
         # Use cli arg since cli ignores setting in config
         if job_config.outputs.tensorboard.write_logs:

diff --git a/sleap/nn/config/outputs.py b/sleap/nn/config/outputs.py
@@ -151,8 +151,8 @@ class OutputsConfig:
         save_visualizations: If True, will render and save visualizations of the model
             predictions as PNGs to "{run_folder}/viz/{split}.{epoch:04d}.png", where the
             split is one of "train", "validation", "test".
-        keep_viz_images: If True, keep the saved visualization images after training
-            completes. This is useful unchecked to reduce the model folder size if you do not need
+        delete_viz_images: If True, delete the saved visualizations after training
+            completes. This is useful to reduce the model folder size if you do not need
             to keep the visualization images.
         zip_outputs: If True, compress the run folder to a zip file. This will be named
             "{run_folder}.zip".
@@ -170,7 +170,7 @@ class OutputsConfig:
     runs_folder: Text = "models"
     tags: List[Text] = attr.ib(factory=list)
     save_visualizations: bool = True
-    keep_viz_images: bool = False
+    delete_viz_images: bool = True
     zip_outputs: bool = False
     log_to_csv: bool = True
     checkpointing: CheckpointingConfig = attr.ib(factory=CheckpointingConfig)

diff --git a/sleap/nn/training.py b/sleap/nn/training.py
@@ -946,7 +946,7 @@ def train(self):
         if self.config.outputs.save_outputs:
             if (
                 self.config.outputs.save_visualizations
-                and not self.config.outputs.keep_viz_images
+                and self.config.outputs.delete_viz_images
             ):
                 self.cleanup()
 
@@ -997,7 +997,7 @@ def cleanup(self):
 
     def package(self):
         """Package model folder into a zip file for portability."""
-        if not self.config.outputs.keep_viz_images:
+        if self.config.outputs.delete_viz_images:
             self.cleanup()
         logger.info(f"Packaging results to: {self.run_path}.zip")
         shutil.make_archive(
@@ -1864,14 +1864,6 @@ def create_trainer_using_cli(args: Optional[List] = None):
             "already specified in the training job config."
         ),
     )
-    parser.add_argument(
-        "--keep_viz",
-        action="store_true",
-        help=(
-            "Keep prediction visualization images in the run folder after training when "
-            "--save_viz is enabled."
-        ),
-    )
     parser.add_argument(
         "--zmq",
         action="store_true",
@@ -1957,7 +1949,6 @@ def create_trainer_using_cli(args: Optional[List] = None):
     if args.suffix != "":
         job_config.outputs.run_name_suffix = args.suffix
     job_config.outputs.save_visualizations |= args.save_viz
-    job_config.outputs.keep_viz_images = args.keep_viz
     if args.labels_path == "":
         args.labels_path = None
     args.video_paths = args.video_paths.split(",")

diff --git a/sleap/training_profiles/baseline.centroid.json b/sleap/training_profiles/baseline.centroid.json
@@ -116,7 +116,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_large_rf.bottomup.json b/sleap/training_profiles/baseline_large_rf.bottomup.json
@@ -125,7 +125,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_large_rf.single.json b/sleap/training_profiles/baseline_large_rf.single.json
@@ -116,7 +116,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_large_rf.topdown.json b/sleap/training_profiles/baseline_large_rf.topdown.json
@@ -117,7 +117,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_medium_rf.bottomup.json b/sleap/training_profiles/baseline_medium_rf.bottomup.json
@@ -125,7 +125,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_medium_rf.single.json b/sleap/training_profiles/baseline_medium_rf.single.json
@@ -116,7 +116,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/baseline_medium_rf.topdown.json b/sleap/training_profiles/baseline_medium_rf.topdown.json
@@ -117,7 +117,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/pretrained.bottomup.json b/sleap/training_profiles/pretrained.bottomup.json
@@ -122,7 +122,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/pretrained.centroid.json b/sleap/training_profiles/pretrained.centroid.json
@@ -113,7 +113,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/pretrained.single.json b/sleap/training_profiles/pretrained.single.json
@@ -113,7 +113,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/sleap/training_profiles/pretrained.topdown.json b/sleap/training_profiles/pretrained.topdown.json
@@ -114,7 +114,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": true,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json b/tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json
@@ -128,7 +128,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json b/tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json
@@ -191,7 +191,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/min_tracks_2node.UNet.topdown_multiclass/initial_config.json b/tests/data/models/min_tracks_2node.UNet.topdown_multiclass/initial_config.json
@@ -141,7 +141,7 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
+        "delete_viz_images": true,
         "zip_outputs": false,
         "log_to_csv": true,
         "checkpointing": {

diff --git a/tests/data/models/min_tracks_2node.UNet.topdown_multiclass/training_config.json b/tests/data/models/min_tracks_2node.UNet.topdown_multiclass/training_config.json
@@ -208,7 +208,7 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
+        "delete_viz_images": true,
         "zip_outputs": false,
         "log_to_csv": true,
         "checkpointing": {

diff --git a/tests/data/models/minimal_instance.UNet.bottomup/initial_config.json b/tests/data/models/minimal_instance.UNet.bottomup/initial_config.json
@@ -127,7 +127,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_instance.UNet.bottomup/training_config.json b/tests/data/models/minimal_instance.UNet.bottomup/training_config.json
@@ -192,7 +192,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json b/tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json
@@ -119,7 +119,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_instance.UNet.centered_instance/training_config.json b/tests/data/models/minimal_instance.UNet.centered_instance/training_config.json
@@ -179,7 +179,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_instance.UNet.centroid/initial_config.json b/tests/data/models/minimal_instance.UNet.centroid/initial_config.json
@@ -118,7 +118,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_instance.UNet.centroid/training_config.json b/tests/data/models/minimal_instance.UNet.centroid/training_config.json
@@ -175,7 +175,6 @@
         "runs_folder": "models",
         "tags": [],
         "save_visualizations": false,
-        "keep_viz_images": false,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_robot.UNet.single_instance/initial_config.json b/tests/data/models/minimal_robot.UNet.single_instance/initial_config.json
@@ -120,7 +120,6 @@
             ""
         ],
         "save_visualizations": false,
-        "keep_viz_images": true,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/data/models/minimal_robot.UNet.single_instance/training_config.json b/tests/data/models/minimal_robot.UNet.single_instance/training_config.json
@@ -180,7 +180,6 @@
             ""
         ],
         "save_visualizations": false,
-        "keep_viz_images": true,
         "log_to_csv": true,
         "checkpointing": {
             "initial_model": false,

diff --git a/tests/gui/test_dialogs.py b/tests/gui/test_dialogs.py
@@ -1,5 +1,6 @@
 """Module to test the dialogs of the GUI (contained in sleap/gui/dialogs)."""
 
+
 import os
 from pathlib import Path