Clean and bump dvc callback: settings, stacked images (ultralytics#4343)

Co-authored-by: Glenn Jocher <[email protected]>
0iui0 · Aug 13, 2023 · ee606bf · ee606bf
1 parent 42aadb4
commit ee606bf
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 16 deletions.
diff --git a/docs/quickstart.md b/docs/quickstart.md
@@ -259,7 +259,7 @@ The table below provides an overview of the settings available for adjustment wi
 | `api_key`          | `''`                  | `str`     | Ultralytics HUB [API Key](https://hub.ultralytics.com/settings?tab=api+keys)                                     |
 | `clearml`          | `True`                | `bool`    | Whether to use ClearML logging                                                                                   |
 | `comet`            | `True`                | `bool`    | Whether to use [Comet ML](https://bit.ly/yolov8-readme-comet) for experiment tracking and visualization          |
-| `dvc`              | `True`                | `bool`    | Whether to use DVC for version control                                                                           |
+| `dvc`              | `True`                | `bool`    | Whether to use [DVC for experiment tracking](https://dvc.org/doc/dvclive/ml-frameworks/yolo) and version control                                                                           |
 | `hub`              | `True`                | `bool`    | Whether to use [Ultralytics HUB](https://hub.ultralytics.com) integration                                        |
 | `mlflow`           | `True`                | `bool`    | Whether to use MLFlow for experiment tracking                                                                    |
 | `neptune`          | `True`                | `bool`    | Whether to use Neptune for experiment tracking                                                                   |

diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
@@ -543,7 +543,8 @@ def plot_metrics(self):
 
     def on_plot(self, name, data=None):
         """Registers plots (e.g. to be consumed in callbacks)"""
-        self.plots[name] = {'data': data, 'timestamp': time.time()}
+        path = Path(name)
+        self.plots[path] = {'data': data, 'timestamp': time.time()}
 
     def final_eval(self):
         """Performs final evaluation and validation for object detection YOLO model."""

diff --git a/ultralytics/engine/validator.py b/ultralytics/engine/validator.py
@@ -303,7 +303,8 @@ def metric_keys(self):
 
     def on_plot(self, name, data=None):
         """Registers plots (e.g. to be consumed in callbacks)"""
-        self.plots[name] = {'data': data, 'timestamp': time.time()}
+        path = Path(name)
+        self.plots[path] = {'data': data, 'timestamp': time.time()}
 
     # TODO: may need to put these following functions into callback
     def plot_val_samples(self, batch, ni):

diff --git a/ultralytics/utils/callbacks/dvc.py b/ultralytics/utils/callbacks/dvc.py
@@ -1,6 +1,8 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
 import os
+import re
+from pathlib import Path
 
 import pkg_resources as pkg
 
@@ -32,13 +34,17 @@
 _training_epoch = False
 
 
-def _logger_disabled():
-    return os.getenv('ULTRALYTICS_DVC_DISABLED', 'false').lower() == 'true'
+def _log_images(path, prefix=''):
+    if live:
+        name = path.name
 
+        # Group images by batch to enable sliders in UI
+        if m := re.search(r'_batch(\d+)', name):
+            ni = m.group(1)
+            new_stem = re.sub(r'_batch(\d+)', '_batch', path.stem)
+            name = (Path(new_stem) / ni).with_suffix(path.suffix)
 
-def _log_images(image_path, prefix=''):
-    if live:
-        live.log_image(os.path.join(prefix, image_path.name), image_path)
+        live.log_image(os.path.join(prefix, name), path)
 
 
 def _log_plots(plots, prefix=''):
@@ -68,14 +74,10 @@ def _log_confusion_matrix(validator):
 def on_pretrain_routine_start(trainer):
     try:
         global live
-        if not _logger_disabled():
-            live = dvclive.Live(save_dvc_exp=True, cache_images=True)
-            LOGGER.info(
-                'DVCLive is detected and auto logging is enabled (can be disabled with `ULTRALYTICS_DVC_DISABLED=true`).'
-            )
-        else:
-            LOGGER.debug('DVCLive is detected and auto logging is disabled via `ULTRALYTICS_DVC_DISABLED`.')
-            live = None
+        live = dvclive.Live(save_dvc_exp=True, cache_images=True)
+        LOGGER.info(
+            f'DVCLive is detected and auto logging is enabled (can be disabled in the {SETTINGS.file} with `dvc: false`).'
+        )
     except Exception as e:
         LOGGER.warning(f'WARNING ⚠️ DVCLive installed but not initialized correctly, not logging this run. {e}')