Add train.py and val.py callbacks (#4220)

* added callbacks

* Update callbacks.py

* Update train.py

* Update val.py

* Fix CamlCase add staticmethod

* Refactor logger into callbacks

* Cleanup

* New callback on_val_image_end()

* Add curves and results images to TensorBoard

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

utils/loggers/__init__.py

@@ -29,10 +29,12 @@ class Loggers():
         self.hyp = hyp
         self.logger = logger  # for printing results to console
         self.include = include
+        self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
+                     'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',  # metrics
+                     'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
+                     'x/lr0', 'x/lr1', 'x/lr2']  # params
         for k in LOGGERS:
             setattr(self, k, None)  # init empty logger dictionary
-
-    def start(self):
         self.csv = True  # always log to csv
 
         # Message
@@ -57,7 +59,11 @@ class Loggers():
         else:
             self.wandb = None
 
-        return self
+    def on_pretrain_routine_end(self):
+        # Callback runs on pre-train routine end
+        paths = self.save_dir.glob('*labels*.jpg')  # training labels
+        if self.wandb:
+            self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
 
     def on_train_batch_end(self, ni, model, imgs, targets, paths, plots):
         # Callback runs on train batch end
@@ -78,8 +84,8 @@ class Loggers():
         if self.wandb:
             self.wandb.current_epoch = epoch + 1
 
-    def on_val_batch_end(self, pred, predn, path, names, im):
-        # Callback runs on train batch end
+    def on_val_image_end(self, pred, predn, path, names, im):
+        # Callback runs on val image end
         if self.wandb:
             self.wandb.val_one_image(pred, predn, path, names, im)
 
@@ -89,25 +95,20 @@ class Loggers():
             files = sorted(self.save_dir.glob('val*.jpg'))
             self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
 
-    def on_train_val_end(self, mloss, results, lr, epoch, best_fitness, fi):
-        # Callback runs on val end during training
+    def on_fit_epoch_end(self, mloss, results, lr, epoch, best_fitness, fi):
+        # Callback runs at the end of each fit (train+val) epoch
         vals = list(mloss) + list(results) + lr
-        keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
-                'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',  # metrics
-                'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
-                'x/lr0', 'x/lr1', 'x/lr2']  # params
-        x = {k: v for k, v in zip(keys, vals)}  # dict
-
+        x = {k: v for k, v in zip(self.keys, vals)}  # dict
         if self.csv:
             file = self.save_dir / 'results.csv'
             n = len(x) + 1  # number of cols
-            s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + keys)).rstrip(',') + '\n')  # add header
+            s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n')  # add header
             with open(file, 'a') as f:
                 f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
 
         if self.tb:
             for k, v in x.items():
-                self.tb.add_scalar(k, v, epoch)  # TensorBoard
+                self.tb.add_scalar(k, v, epoch)
 
         if self.wandb:
             self.wandb.log(x)
@@ -119,20 +120,22 @@ class Loggers():
             if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
                 self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
 
-    def on_train_end(self, last, best, plots):
+    def on_train_end(self, last, best, plots, epoch):
         # Callback runs on training end
         if plots:
             plot_results(dir=self.save_dir)  # save results.png
         files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
         files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()]  # filter
+
+        if self.tb:
+            from PIL import Image
+            import numpy as np
+            for f in files:
+                self.tb.add_image(f.stem, np.asarray(Image.open(f)), epoch, dataformats='HWC')
+
         if self.wandb:
             wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
             wandb.log_artifact(str(best if best.exists() else last), type='model',
                                name='run_' + self.wandb.wandb_run.id + '_model',
                                aliases=['latest', 'best', 'stripped'])
             self.wandb.finish_run()
-
-    def log_images(self, paths):
-        # Log images
-        if self.wandb:
-            self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})