From 029882146724480196fc4be0df5a680477c42cc4 Mon Sep 17 00:00:00 2001
From: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Date: Sun, 25 Dec 2022 19:06:41 +0530
Subject: [PATCH] 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 | 19 +++++++++++++++++++
 ultralytics/yolo/utils/configs/default.yaml |  2 +-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/ultralytics/yolo/utils/callbacks/clearml.py b/ultralytics/yolo/utils/callbacks/clearml.py
index 946ca0b..400944b 100644
--- a/ultralytics/yolo/utils/callbacks/clearml.py
+++ b/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 {}
diff --git a/ultralytics/yolo/utils/configs/default.yaml b/ultralytics/yolo/utils/configs/default.yaml
index e40cb83..8750e7f 100644
--- a/ultralytics/yolo/utils/configs/default.yaml
+++ b/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/"