Add Keep visualizations checkbox to training GUI

[hajin-park]

Description

This PR clarifies the use of the Visualize Predictions During Training checkbox by adding a new Keep Prediction Visualization Images After Training checkbox to the training config GUI. If training visualization images are enabled, the images are deleted after training by default unless the new checkbox is enabled (checked).

Types of changes

• Bugfix
• New feature
• Refactor / Code style update (no logical changes)
• Build / CI changes
• Documentation Update
• Other (explain)

Does this address any currently open issues?

• Possibility to save training visualization images #1016

Outside contributors checklist

• Review the guidelines for contributing to this repository
• Read and sign the CLA and add yourself to the authors list
• Make sure you are making a pull request against the develop branch (not main). Also you should start your branch off develop
• Add tests that prove your fix is effective or that your feature works
• Add necessary documentation (if appropriate)

Thank you for contributing to SLEAP!

❤️

Summary by CodeRabbit

• New Features

  • Introduced --keep_viz argument to retain prediction visualization images after training.
  • Renamed --save_viz argument to --view_viz.
  • Added --delete_viz argument for deleting visualization images after training.

• Tests

  • Added test test_keep_viz_cli to check the functionality of the --keep_viz option.
  • Updated tests to check visualization paths and save directories.

[codecov]

Codecov Report

Attention: Patch coverage is 42.85714% with 4 lines in your changes missing coverage. Please review.

Project coverage is 74.13%. Comparing base (7ed1229) to head (33c9d62).
Report is 20 commits behind head on develop.

Files	Patch %	Lines
sleap/gui/learning/runners.py	0.00%	3 Missing ⚠️
sleap/nn/training.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #1824      +/-   ##
===========================================
+ Coverage    73.30%   74.13%   +0.82%     
===========================================
  Files          134      135       +1     
  Lines        24087    24577     +490     
===========================================
+ Hits         17658    18221     +563     
+ Misses        6429     6356      -73     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[gitttt-1234]

From subgroup:

• Keep it named save_visualizations
• Add checkbox to GUI labeled "Keep visualization images after training" with a default value of off (unchecked)

[hajin-park]

@coderabbitai review

[coderabbitai]

Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 5

Outside diff range and nitpick comments (17)

sleap/config/pipeline_form.yaml (3)

Line range hint 55-55: Remove trailing spaces.

Trailing spaces detected. Remove the trailing spaces to adhere to YAML formatting standards.

-  default: "multi-animal bottom-up "
+  default: "multi-animal bottom-up"
-  default: 9000 
+  default: 9000
-  default: 0.95 
+  default: 0.95

Also applies to: 220-220, 391-391

---

Line range hint 422-422: Fix wrong indentation.

Incorrect indentation detected. Adjust the indentation to match the expected level.

-    default: false
+  default: false
-    default: 10
+  default: 10

Also applies to: 506-506

---

Line range hint 569-569: Reduce blank lines.

Too many blank lines detected. Reduce the number of blank lines to adhere to YAML formatting standards.

-

docs/guides/cli.md (2)

Line range hint 253-253: Correct the verb form.

The verb "merged" should be in its base form "merge".

- `labels_with_predictions.slp` can then merged into the base labels project from the SLEAP GUI via **File** --> **Merge into project...**.
+ `labels_with_predictions.slp` can then merge into the base labels project from the SLEAP GUI via **File** --> **Merge into project...**.

---

Line range hint 333-333: Correct article usage.

The plural noun "predictions" should not be used with the article "a".

- For example, to convert a predictions SLP file to an analysis HDF5 file:
+ For example, to convert predictions SLP file to an analysis HDF5 file:

sleap/gui/learning/runners.py (9)

Line range hint 50-50: Replace deprecated type Text with str.

The type Text is deprecated and should be replaced with str.

- def path(self) -> Text:
+ def path(self) -> str:

---

Line range hint 55-55: Replace deprecated type Text with str.

The type Text is deprecated and should be replaced with str.

- def cli_args(self) -> List[Text]:
+ def cli_args(self) -> List[str]:

---

Line range hint 200-200: Replace deprecated type Text with str.

The type Text is deprecated and should be replaced with str.

- def path(self) -> Text:
+ def path(self) -> str:

---

Line range hint 254-255: Simplify nested if statements.

Use a single if statement instead of nested if statements.

- if self.frame_filter == "user":
-     args_list.append("--only-labeled-frames")
- elif self.frame_filter == "suggested":
-     args_list.append("--only-suggested-frames")
+ if self.frame_filter in ["user", "suggested"]:
+     args_list.append(f"--only-{self.frame_filter}-frames")

---

Line range hint 288-288: Replace deprecated type Text with str.

The type Text is deprecated and should be replaced with str.

- def make_predict_cli_call(
+ def make_predict_cli_call(
- ) -> List[Text]:
+ ) -> List[str]:

---

Line range hint 312-312: Avoid bare except statements.

Specify the exception type to avoid catching unintended exceptions.

- except:
+ except Exception:

---

Line range hint 669-669: Remove extraneous f prefix from string.

The f prefix is used for strings that need formatting. Remove it where not necessary.

- win.set_message(f"Preparing to run training...")
+ win.set_message("Preparing to run training...")

---

Line range hint 799-799: Function definition does not bind loop variable i.

Ensure that the loop variable i is bound within the function definition.

- for i, item_for_inference in enumerate(items_for_inference.items):
+ for i, item_for_inference in enumerate(items_for_inference.items):
+     def waiting_item(**kwargs):
+         kwargs["current_item"] = i
+         kwargs["total_items"] = len(items_for_inference.items)
+         return waiting(**kwargs)

---

Line range hint 841-841: Remove assignment to unused variable success.

The assignment to the variable success is not used and should be removed.

- success = False

sleap/nn/training.py (3)

Line range hint 16-16: Remove unused import.

The import cattr is not used in the code.

- import cattr

---

Line range hint 79-79: Remove unused import.

The import plot_pafs from sleap.nn.viz is not used in the code.

- from sleap.nn.viz import plot_confmaps, plot_img, plot_pafs, plot_peaks
+ from sleap.nn.viz import plot_confmaps, plot_img, plot_peaks

---

Line range hint 106-109: Replace deprecated Text annotation with str.

The Text annotation from typing is deprecated. Replace it with str.

- from typing import Callable, List, Optional, Text, TypeVar, Union
+ from typing import Callable, List, Optional, TypeVar, Union

- def get_timestamp() -> Text:
+ def get_timestamp() -> str:

- def setup_new_run_folder(config: OutputsConfig, base_run_name: Optional[Text] = None) -> Text:
+ def setup_new_run_folder(config: OutputsConfig, base_run_name: Optional[str] = None) -> str:

- def setup_checkpointing(config: CheckpointingConfig, run_path: Text) -> List[tf.keras.callbacks.Callback]:
+ def setup_checkpointing(config: CheckpointingConfig, run_path: str) -> List[tf.keras.callbacks.Callback]:

- def setup_tensorboard(config: TensorBoardConfig, run_path: Text) -> List[tf.keras.callbacks.Callback]:
+ def setup_tensorboard(config: TensorBoardConfig, run_path: str) -> List[tf.keras.callbacks.Callback]:

- def setup_output_callbacks(config: OutputsConfig, run_path: Optional[Text] = None) -> List[tf.keras.callbacks.Callback]:
+ def setup_output_callbacks(config: OutputsConfig, run_path: Optional[str] = None) -> List[tf.keras.callbacks.Callback]:

- def setup_visualization(config: OutputsConfig, run_path: Text, viz_fn: Callable[[], matplotlib.figure.Figure], name: Text) -> List[tf.keras.callbacks.Callback]:
+ def setup_visualization(config: OutputsConfig, run_path: str, viz_fn: Callable[[], matplotlib.figure.Figure], name: str) -> List[tf.keras.callbacks.Callback]:

- def sanitize_scope_name(name: Text) -> Text:
+ def sanitize_scope_name(name: str) -> str:

- run_path: Optional[Text] = attr.ib(default=None, init=False)
+ run_path: Optional[str] = attr.ib(default=None, init=False)

- def cleanup(self):
+ def cleanup(self) -> None:

- def package(self):
+ def package(self) -> None:

- def visualize_example(example):
+ def visualize_example(example) -> matplotlib.figure.Figure:

- def visualize_confmaps_example(example):
+ def visualize_confmaps_example(example) -> matplotlib.figure.Figure:

- def visualize_pafs_example(example):
+ def visualize_pafs_example(example) -> matplotlib.figure.Figure:

- def visualize_class_maps_example(example):
+ def visualize_class_maps_example(example) -> matplotlib.figure.Figure:

- def main(args: Optional[List] = None):
+ def main(args: Optional[List] = None) -> None:

Also applies to: 132-132, 134-134, 138-138, 156-159, 308-308, 358-358, 364-365, 416-416, 464-464, 485-485, 503-503, 505-505, 539-539, 613-613, 619-622, 904-904, 907-907, 909-909, 911-911, 913-913, 915-915, 917-917, 926-926, 932-932, 1056-1056, 1061-1061, 1112-1112, 1120-1120, 1169-1169, 1174-1174, 1216-1216, 1224-1224, 1289-1289, 1294-1294, 1405-1405, 1410-1410, 1470-1470, 1478-1478, 1488-1488, 1496-1496, 1546-1546, 1551-1551, 1608-1608, 1616-1616, 1625-1625, 1635-1635, 1699-1699, 1704-1704

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 22733f1 and c75a0ef.

Files selected for processing (29)

• docs/guides/cli.md (2 hunks)
• docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb (2 hunks)
• sleap/config/pipeline_form.yaml (1 hunks)
• sleap/gui/learning/runners.py (9 hunks)
• sleap/nn/config/outputs.py (2 hunks)
• sleap/nn/training.py (4 hunks)
• sleap/training_profiles/baseline.centroid.json (1 hunks)
• sleap/training_profiles/baseline_large_rf.bottomup.json (1 hunks)
• sleap/training_profiles/baseline_large_rf.single.json (1 hunks)
• sleap/training_profiles/baseline_large_rf.topdown.json (1 hunks)
• sleap/training_profiles/baseline_medium_rf.bottomup.json (1 hunks)
• sleap/training_profiles/baseline_medium_rf.single.json (1 hunks)
• sleap/training_profiles/baseline_medium_rf.topdown.json (1 hunks)
• sleap/training_profiles/pretrained.bottomup.json (1 hunks)
• sleap/training_profiles/pretrained.centroid.json (1 hunks)
• sleap/training_profiles/pretrained.single.json (1 hunks)
• sleap/training_profiles/pretrained.topdown.json (1 hunks)
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json (1 hunks)
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json (1 hunks)
• tests/data/models/min_tracks_2node.UNet.topdown_multiclass/initial_config.json (1 hunks)
• tests/data/models/min_tracks_2node.UNet.topdown_multiclass/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.bottomup/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.bottomup/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centered_instance/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centroid/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centroid/training_config.json (1 hunks)
• tests/data/models/minimal_robot.UNet.single_instance/initial_config.json (1 hunks)
• tests/data/models/minimal_robot.UNet.single_instance/training_config.json (1 hunks)

Files skipped from review due to trivial changes (9)

• sleap/training_profiles/baseline_large_rf.bottomup.json
• sleap/training_profiles/baseline_large_rf.topdown.json
• sleap/training_profiles/baseline_medium_rf.bottomup.json
• sleap/training_profiles/pretrained.topdown.json
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json
• tests/data/models/min_tracks_2node.UNet.topdown_multiclass/initial_config.json
• tests/data/models/minimal_instance.UNet.centroid/initial_config.json
• tests/data/models/minimal_robot.UNet.single_instance/training_config.json

Additional context used

Ruff

sleap/nn/config/outputs.py

66-66: typing.Text is deprecated, use str

Replace with str

(UP019)

---

106-106: typing.Text is deprecated, use str

Replace with str

(UP019)

---

109-109: typing.Text is deprecated, use str

Replace with str

(UP019)

---

167-167: typing.Text is deprecated, use str

Replace with str

(UP019)

---

168-168: typing.Text is deprecated, use str

Replace with str

(UP019)

---

169-169: typing.Text is deprecated, use str

Replace with str

(UP019)

---

170-170: typing.Text is deprecated, use str

Replace with str

(UP019)

---

171-171: typing.Text is deprecated, use str

Replace with str

(UP019)

---

181-181: typing.Text is deprecated, use str

Replace with str

(UP019)

sleap/gui/learning/runners.py

22-22: sleap.io.video.SingleImageVideo imported but unused

Remove unused import: sleap.io.video.SingleImageVideo

(F401)

---

50-50: typing.Text is deprecated, use str

Replace with str

(UP019)

---

55-55: typing.Text is deprecated, use str

Replace with str

(UP019)

---

200-200: typing.Text is deprecated, use str

Replace with str

(UP019)

---

254-255: Use a single if statement instead of nested if statements

(SIM102)

---

288-288: typing.Text is deprecated, use str

Replace with str

(UP019)

---

312-312: Do not use bare except

(E722)

---

594-594: typing.Text is deprecated, use str

Replace with str

(UP019)

---

594-594: typing.Text is deprecated, use str

Replace with str

(UP019)

---

669-669: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

799-799: Function definition does not bind loop variable i

(B023)

---

833-833: typing.Text is deprecated, use str

Replace with str

(UP019)

---

841-841: Local variable success is assigned to but never used

Remove assignment to unused variable success

(F841)

---

888-891: Use ternary operator ret = "success" if proc.returncode == 0 else proc.returncode instead of if-else-block

Replace if-else-block with ret = "success" if proc.returncode == 0 else proc.returncode

(SIM108)

sleap/nn/training.py

16-16: cattr imported but unused

Remove unused import: cattr

(F401)

---

79-79: sleap.nn.viz.plot_pafs imported but unused

Remove unused import: sleap.nn.viz.plot_pafs

(F401)

---

106-106: typing.Text is deprecated, use str

Replace with str

(UP019)

---

107-107: typing.Text is deprecated, use str

Replace with str

(UP019)

---

108-108: typing.Text is deprecated, use str

Replace with str

(UP019)

---

109-109: typing.Text is deprecated, use str

Replace with str

(UP019)

---

132-132: typing.Text is deprecated, use str

Replace with str

(UP019)

---

134-134: typing.Text is deprecated, use str

Replace with str

(UP019)

---

138-138: typing.Text is deprecated, use str

Replace with str

(UP019)

---

156-156: typing.Text is deprecated, use str

Replace with str

(UP019)

---

157-157: typing.Text is deprecated, use str

Replace with str

(UP019)

---

158-158: typing.Text is deprecated, use str

Replace with str

(UP019)

---

159-159: typing.Text is deprecated, use str

Replace with str

(UP019)

---

308-308: typing.Text is deprecated, use str

Replace with str

(UP019)

---

358-358: typing.Text is deprecated, use str

Replace with str

(UP019)

---

364-364: typing.Text is deprecated, use str

Replace with str

(UP019)

---

365-365: typing.Text is deprecated, use str

Replace with str

(UP019)

---

416-416: typing.Text is deprecated, use str

Replace with str

(UP019)

---

464-464: typing.Text is deprecated, use str

Replace with str

(UP019)

---

485-485: typing.Text is deprecated, use str

Replace with str

(UP019)

---

503-503: typing.Text is deprecated, use str

Replace with str

(UP019)

---

505-505: typing.Text is deprecated, use str

Replace with str

(UP019)

---

539-539: typing.Text is deprecated, use str

Replace with str

(UP019)

---

539-539: typing.Text is deprecated, use str

Replace with str

(UP019)

---

613-613: typing.Text is deprecated, use str

Replace with str

(UP019)

---

619-619: typing.Text is deprecated, use str

Replace with str

(UP019)

---

620-620: typing.Text is deprecated, use str

Replace with str

(UP019)

---

621-621: typing.Text is deprecated, use str

Replace with str

(UP019)

---

622-622: typing.Text is deprecated, use str

Replace with str

(UP019)

---

904-904: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

907-907: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

909-909: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

911-911: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

913-913: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

915-915: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

917-917: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

926-926: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

932-932: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1056-1056: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1061-1061: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1112-1112: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1120-1120: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1169-1169: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1174-1174: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1216-1216: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1224-1224: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1289-1289: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1294-1294: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1405-1405: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1410-1410: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1470-1470: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1478-1478: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1488-1488: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1496-1496: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1546-1546: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1551-1551: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1608-1608: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1616-1616: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1625-1625: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1635-1635: f-string without any placeholders

Remove extraneous f prefix

(F541)

---

1699-1699: typing.Text is deprecated, use str

Replace with str

(UP019)

---

1704-1704: typing.Text is deprecated, use str

Replace with str

(UP019)

yamllint

sleap/config/pipeline_form.yaml

[error] 55-55: trailing spaces

(trailing-spaces)

---

[error] 220-220: trailing spaces

(trailing-spaces)

---

[error] 391-391: trailing spaces

(trailing-spaces)

---

[warning] 422-422: wrong indentation: expected 2 but found 4

(indentation)

---

[warning] 506-506: wrong indentation: expected 2 but found 4

(indentation)

---

[warning] 569-569: too many blank lines

(4 > 2) (empty-lines)

LanguageTool

docs/guides/cli.md

[typographical] ~120-~120: Consider adding a comma.
Context: ...ook. If you specify how many identities there should be in a frame (i.e., the number ...

(IF_THERE_COMMA)

---

[misspelling] ~241-~241: This word is normally spelled as one.
Context: ...put_video.mp4" **2. Inference with multi-model pipelines (e.g., top-down):** none ...

(EN_COMPOUNDS_MULTI_MODEL)

---

[grammar] ~253-~253: The modal verb ‘can’ requires the verb’s base form.
Context: ... labels_with_predictions.slp can then merged into the base labels project from the S...

(MD_BASEFORM)

---

[style] ~291-~291: The preposition ‘Amongst’ is correct, but some people think that it is old-fashioned or literary. A more frequently used alternative is the preposition “among”.
Context: ...t between various dataset file formats. Amongst other things, it can be used to export ...

(AMONGST)

---

[grammar] ~333-~333: The plural noun “predictions” cannot be used with the article “a”. Did you mean “a prediction” or “predictions”?
Context: ...nversion). For example, to convert a predictions SLP file to an analysis HDF5 file: ...

(A_NNS)

---

[style] ~345-~345: Consider using a different verb to strengthen your wording.
Context: ... sleap-inspect {code}sleap-inspect gives you various information about a SLEAP d...

(GIVE_INFORMATION)

---

[uncategorized] ~345-~345: Possible missing comma found.
Context: ...rious information about a SLEAP dataset file such as a list of videos and a count of...

(AI_HYDRA_LEO_MISSING_COMMA)

---

[uncategorized] ~406-~406: Possible missing comma found.
Context: ...re's also a script to output diagnostic information which may help us if you need to contac...

(AI_HYDRA_LEO_MISSING_COMMA)

Markdownlint

docs/guides/cli.md

233-233: Punctuation: ':'
Trailing punctuation in heading

(MD026, no-trailing-punctuation)

---

335-335: null
Fenced code blocks should have a language specified

(MD040, fenced-code-language)

Gitleaks

docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb

99-99: Identified a pattern that may indicate AWS credentials, risking unauthorized cloud resource access and data breaches on AWS platforms.

(aws-access-token)

---

100-100: Identified a pattern that may indicate AWS credentials, risking unauthorized cloud resource access and data breaches on AWS platforms.

(aws-access-token)

Additional comments not posted (42)

sleap/training_profiles/pretrained.single.json (1)

116-116: Addition of keep_viz_images key.

The new key keep_viz_images has been added to the outputs section. Ensure that this key is consistently used across all relevant configuration files and that its default value aligns with the intended behavior.

sleap/training_profiles/pretrained.centroid.json (1)

116-116: Addition of keep_viz_images key.

The new key keep_viz_images has been added to the outputs section. Ensure that this key is consistently used across all relevant configuration files and that its default value aligns with the intended behavior.

sleap/training_profiles/baseline_medium_rf.single.json (1)

119-119: Addition of keep_viz_images key.

The new key keep_viz_images has been added to the outputs section. Ensure that this key is consistently used across all relevant configuration files and that its default value aligns with the intended behavior.

sleap/training_profiles/baseline_large_rf.single.json (1)

119-119: Addition of keep_viz_images key.

The new key keep_viz_images has been added to the outputs section. Ensure that this key is consistently used across all relevant configuration files and that its default value aligns with the intended behavior.

tests/data/models/minimal_robot.UNet.single_instance/initial_config.json (1)

123-123: Addition of keep_viz_images parameter looks good.

The keep_viz_images parameter has been correctly added to the configuration. Ensure consistency across all related configuration files.

sleap/training_profiles/baseline.centroid.json (1)

119-119: Addition of keep_viz_images parameter looks good.

The keep_viz_images parameter has been correctly added to the configuration. Ensure consistency across all related configuration files.

tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json (1)

122-122: Addition of keep_viz_images parameter looks good.

The keep_viz_images parameter has been correctly added to the configuration. Ensure consistency across all related configuration files.

sleap/training_profiles/baseline_medium_rf.topdown.json (1)

120-120: Addition of keep_viz_images parameter looks good.

The keep_viz_images parameter has been correctly added to the configuration. Ensure consistency across all related configuration files.

sleap/training_profiles/pretrained.bottomup.json (1)

125-125: Addition of keep_viz_images parameter.

The parameter keep_viz_images has been added to the "outputs" section with a default value of false. This addition aligns with the PR objective to control the retention of visualization images after training.

tests/data/models/minimal_instance.UNet.bottomup/initial_config.json (1)

130-130: Addition of keep_viz_images parameter.

The parameter keep_viz_images has been added to the "outputs" section with a default value of true. This addition aligns with the PR objective to control the retention of visualization images after training.

tests/data/models/minimal_instance.UNet.centroid/training_config.json (1)

178-178: Addition of keep_viz_images parameter.

The parameter keep_viz_images has been added to the "outputs" section with a default value of true. This addition aligns with the PR objective to control the retention of visualization images after training.

tests/data/models/minimal_instance.UNet.centered_instance/training_config.json (1)

182-182: Addition of keep_viz_images parameter.

The parameter keep_viz_images has been added to the "outputs" section with a default value of true. This addition aligns with the PR objective to control the retention of visualization images after training.

tests/data/models/minimal_instance.UNet.bottomup/training_config.json (1)

195-195: Addition of keep_viz_images parameter.

The new parameter keep_viz_images is correctly added to the outputs section. Ensure the logic to handle this parameter is implemented in the codebase.

tests/data/models/min_tracks_2node.UNet.topdown_multiclass/training_config.json (1)

211-211: Addition of keep_viz_images parameter.

The new parameter keep_viz_images is correctly added to the outputs section. Ensure the logic to handle this parameter is implemented in the codebase.

sleap/nn/config/outputs.py (1)

154-155: Addition of keep_viz_images attribute.

The new attribute keep_viz_images is correctly added to the OutputsConfig class. Ensure the logic to handle this attribute is implemented in the codebase.

sleap/config/pipeline_form.yaml (1)

289-293: Addition of _keep_viz parameter.

The new parameter _keep_viz is correctly added to the output options section. Ensure the logic to handle this parameter is implemented in the codebase.

docs/guides/cli.md (8)

Line range hint 120-120: Add a comma after "book".

A comma is needed to improve readability.

- If you specify how many identities there should be in a frame (i.e., the number of animals) with the {code}`--tracking.clean_instance_count` argument, then we will use a heuristic method to connect "breaks" in the track identities where we lose one identity and spawn another.
+ If you specify how many identities there should be in a frame (i.e., the number of animals), with the {code}`--tracking.clean_instance_count` argument, then we will use a heuristic method to connect "breaks" in the track identities where we lose one identity and spawn another.

---

Line range hint 241-241: Correct the spelling of "multi-model".

The term "multi-model" should be hyphenated.

- **2. Inference with multi-model pipelines (e.g., top-down):**
+ **2. Inference with multi-model pipelines (e.g., top-down):**

---

Line range hint 291-291: Replace "Amongst" with "Among".

"Amongst" is old-fashioned; "Among" is more frequently used.

- Amongst other things, it can be used to export data from a SLEAP dataset into an HDF5 file that can be easily used for analysis (e.g., read from MATLAB).
+ Among other things, it can be used to export data from a SLEAP dataset into an HDF5 file that can be easily used for analysis (e.g., read from MATLAB).

---

Line range hint 345-345: Improve wording.

Use a stronger verb to enhance the sentence.

- `sleap-inspect`  {code}`sleap-inspect` gives you various information about a SLEAP dataset file such as a list of videos and a count of the frames with labels.
+ `sleap-inspect`  {code}`sleap-inspect` provides various information about a SLEAP dataset file such as a list of videos and a count of the frames with labels.

---

Line range hint 345-345: Add a comma after "file".

A comma is needed to improve readability.

- `sleap-inspect`  {code}`sleap-inspect` gives you various information about a SLEAP dataset file such as a list of videos and a count of the frames with labels.
+ `sleap-inspect`  {code}`sleap-inspect` provides various information about a SLEAP dataset file, such as a list of videos and a count of the frames with labels.

---

Line range hint 406-406: Add a comma after "information".

A comma is needed to improve readability.

- There's also a script to output diagnostic information which may help us if you need to contact us about problems installing or running SLEAP.
+ There's also a script to output diagnostic information, which may help us if you need to contact us about problems installing or running SLEAP.

---

Line range hint 233-233: Remove trailing punctuation in heading.

Trailing punctuation in headings should be avoided.

- #### Examples:
+ #### Examples

---

Line range hint 335-335: Specify language in fenced code block.

Fenced code blocks should have a language specified.

- ```
+ ```none

sleap/nn/training.py (16)

Line range hint 24-103: LGTM!

The DataReaders class is well-structured and correctly implemented.

---

Line range hint 106-121: LGTM!

The setup_optimizer function is correctly implemented.

---

Line range hint 124-138: LGTM!

The setup_losses function is correctly implemented.

---

Line range hint 141-153: LGTM!

The setup_metrics function is correctly implemented.

---

Line range hint 156-177: LGTM!

The setup_optimization_callbacks function is correctly implemented.

---

Line range hint 179-181: LGTM!

The get_timestamp function is correctly implemented.

---

Line range hint 183-204: LGTM!

The setup_new_run_folder function is correctly implemented.

---

Line range hint 207-226: LGTM!

The setup_zmq_callbacks function is correctly implemented.

---

Line range hint 308-1039: LGTM!

The Trainer class is well-structured and correctly implemented.

---

Line range hint 1042-1185: LGTM!

The SingleInstanceModelTrainer class is well-structured and correctly implemented.

---

Line range hint 1188-1314: LGTM!

The CentroidConfmapsModelTrainer class is well-structured and correctly implemented.

---

Line range hint 1317-1451: LGTM!

The TopdownConfmapsModelTrainer class is well-structured and correctly implemented.

---

Line range hint 1454-1593: LGTM!

The BottomUpModelTrainer class is well-structured and correctly implemented.

---

Line range hint 1596-1730: LGTM!

The BottomUpMultiClassModelTrainer class is well-structured and correctly implemented.

---

Line range hint 1733-1970: LGTM!

The create_trainer_using_cli function is correctly implemented.

---

Line range hint 1973-1984: LGTM!

The main function is correctly implemented.

docs/notebooks/Training_and_inference_on_an_example_dataset.ipynb (2)

338-338: Add a new line before the added parameter for readability.

To improve readability, add a new line before the keep_viz_images parameter.

...
        "save_visualizations": true,
+
        "keep_viz_images": true,
        "zip_outputs": false,
...

---

730-730: Add a new line before the added parameter for readability.

To improve readability, add a new line before the keep_viz_images parameter.

...
        "save_visualizations": true,
+
        "keep_viz_images": true,
        "zip_outputs": false,
...

[coderabbitai]

Actionable comments posted: 1

Outside diff range and nitpick comments (3)

docs/guides/cli.md (2)

39-40: Update the description for --save_viz.

The description for --save_viz should be updated to clarify that it works in conjunction with --keep_viz.

-  --save_viz            Enable saving of prediction visualizations to the run
+  --save_viz            Enable saving of prediction visualizations to the run folder during training.

---

71-72: Clarify the dependency of --keep_viz on --save_viz.

The description for --keep_viz should explicitly mention that it only has an effect if --save_viz is enabled.

-  --keep_viz            Keep prediction visualization images in the run
-                        folder after training when save_viz is enabled.
+  --keep_viz            Keep prediction visualization images in the run folder after training if --save_viz is enabled.

sleap/nn/training.py (1)

1868-1874: Add a period at the end of the help text.

For consistency with other help texts in the CLI parser, add a period at the end.

-            "--save_viz is enabled."
+            "--save_viz is enabled."

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between c75a0ef and 227138d.

Files selected for processing (4)

• docs/guides/cli.md (2 hunks)
• sleap/nn/training.py (4 hunks)
• tests/gui/test_dialogs.py (1 hunks)
• tests/nn/test_training.py (2 hunks)

Additional context used

Ruff

tests/gui/test_dialogs.py

3-3: os imported but unused

Remove unused import: os

(F401)

---

4-4: pathlib.Path imported but unused

Remove unused import: pathlib.Path

(F401)

Additional comments not posted (5)

tests/gui/test_dialogs.py (1)

3-3: Remove unused import os.

The os import is unused and can be safely removed.

Tools

Ruff

3-3: os imported but unused

Remove unused import: os

(F401)

tests/nn/test_training.py (3)

126-149: LGTM! Ensure the keep_viz_images parameter is tested correctly.

The changes correctly set and test the keep_viz_images parameter.

---

152-179: LGTM! Ensure the keep_viz_images parameter is tested correctly.

The changes correctly set and test the keep_viz_images parameter.

---

392-410: LGTM! The new test for --keep_viz CLI option is well-implemented.

The function correctly tests the --keep_viz CLI option.

sleap/nn/training.py (1)

1960-1960: LGTM!

The new CLI argument --keep_viz is correctly added to the job configuration.

[coderabbitai]

Remove unused import pathlib.Path.

The Path import from pathlib is unused and can be safely removed.

- from pathlib import Path

Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

from pathlib import Path

Tools

Ruff

4-4: pathlib.Path imported but unused

Remove unused import: pathlib.Path

(F401)

[roomrys]

Ready to merge after tests pass (post updating branch)

[coderabbitai]

Actionable comments posted: 1

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 227138d and 33c9d62.

Files selected for processing (9)

• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json (1 hunks)
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.bottomup/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.bottomup/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centered_instance/training_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centroid/initial_config.json (1 hunks)
• tests/data/models/minimal_instance.UNet.centroid/training_config.json (1 hunks)
• tests/nn/test_training.py (2 hunks)

Files skipped from review due to trivial changes (8)

• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/initial_config.json
• tests/data/models/min_tracks_2node.UNet.bottomup_multiclass/training_config.json
• tests/data/models/minimal_instance.UNet.bottomup/initial_config.json
• tests/data/models/minimal_instance.UNet.bottomup/training_config.json
• tests/data/models/minimal_instance.UNet.centered_instance/initial_config.json
• tests/data/models/minimal_instance.UNet.centered_instance/training_config.json
• tests/data/models/minimal_instance.UNet.centroid/initial_config.json
• tests/data/models/minimal_instance.UNet.centroid/training_config.json

Additional context used

Ruff

tests/nn/test_training.py

411-411: Remove unnecessary True if ... else False

Remove unnecessary True if ... else False

(SIM210)

Additional comments not posted (2)

tests/nn/test_training.py (2)

126-149: LGTM!

The changes to manage visualization images during training are consistent with the PR objectives and context.

---

152-179: LGTM!

The changes to manage visualization images during training are consistent with the PR objectives and context.

[coderabbitai]

Simplify the ternary operator.

The ternary operator can be simplified to a direct comparison.

-  assert trainer.config.outputs.keep_viz_images == (
-      True if keep_viz_cli == "--keep_viz" else False
-  )
+  assert trainer.config.outputs.keep_viz_images == (keep_viz_cli == "--keep_viz")

Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	@pytest.mark.parametrize("keep_viz_cli", ["", "--keep_viz"])
	def test_keep_viz_cli(
	keep_viz_cli,
	min_single_instance_robot_model_path: str,
	tmp_path: str,
	):
	"""Test training CLI for --keep_viz option."""
	cfg_dir = min_single_instance_robot_model_path
	cfg = TrainingJobConfig.load_json(str(Path(cfg_dir, "training_config.json")))
	# Save training config to tmp folder
	cfg_path = str(Path(tmp_path, "training_config.json"))
	cfg.save_json(cfg_path)
	cli_args = [cfg_path, keep_viz_cli]
	trainer = sleap_train(cli_args)
	# Check that --keep_viz is set correctly
	assert trainer.config.outputs.keep_viz_images == (
	True if keep_viz_cli == "--keep_viz" else False
	)
	assert trainer.config.outputs.keep_viz_images == (keep_viz_cli == "--keep_viz")

Tools

Ruff

411-411: Remove unnecessary True if ... else False

Remove unnecessary True if ... else False

(SIM210)