ClearML: Mosaic Logging (#91)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

ultralytics/yolo/utils/callbacks/clearml.py

@@ -1,3 +1,6 @@
+import os
+from pathlib import Path
+
 from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
 
 try:
@@ -9,6 +12,13 @@ except (ImportError, AssertionError):
     clearml = None
 
 
+def _log_images(imgs_dict, group="", step=0):
+    task = Task.current_task()
+    if task:
+        for k, v in imgs_dict.items():
+            task.get_logger().report_image(group, k, step, v)
+
+
 def on_train_start(trainer):
     # TODO: reuse existing task
     task = Task.init(project_name=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv8',
@@ -20,6 +30,14 @@ def on_train_start(trainer):
     task.connect(dict(trainer.args), name='General')
 
 
+def on_epoch_start(trainer):
+    if trainer.epoch == 1:
+        plots = [filename for filename in os.listdir(trainer.save_dir) if filename.startswith("train_batch")]
+        imgs_dict = {f"train_batch_{i}": Path(trainer.save_dir) / img for i, img in enumerate(plots)}
+        if imgs_dict:
+            _log_images(imgs_dict, "Mosaic", trainer.epoch)
+
+
 def on_val_end(trainer):
     if trainer.epoch == 0:
         model_info = {
@@ -37,5 +55,6 @@ def on_train_end(trainer):
 
 callbacks = {
     "on_train_start": on_train_start,
+    "on_epoch_start": on_epoch_start,
     "on_val_end": on_val_end,
     "on_train_end": on_train_end} if clearml else {}

ultralytics/yolo/utils/configs/default.yaml

@@ -46,7 +46,7 @@ iou_thres: 0.6
 max_det: 300
 half: True
 dnn: False # use OpenCV DNN for ONNX inference
-plots: False
+plots: True
 
 # Prediction settings:
 source: "ultralytics/assets/"