Fix ClearML Mosaic callback to 'on_train_epoch_end' (#92)

This commit is contained in:
Glenn Jocher
2022-12-25 17:04:24 +01:00
committed by GitHub
parent 0298821467
commit 249dfbdc05
5 changed files with 17 additions and 24 deletions

View File

@ -1,6 +1,3 @@
import os
from pathlib import Path
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
try:
@ -30,12 +27,9 @@ def on_train_start(trainer):
task.connect(dict(trainer.args), name='General')
def on_epoch_start(trainer):
def on_train_epoch_end(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)
_log_images({f.stem: str(f) for f in trainer.save_dir.glob('train_batch*.jpg')}, "Mosaic", trainer.epoch)
def on_val_end(trainer):
@ -55,6 +49,6 @@ def on_train_end(trainer):
callbacks = {
"on_train_start": on_train_start,
"on_epoch_start": on_epoch_start,
"on_train_epoch_end": on_train_epoch_end,
"on_val_end": on_val_end,
"on_train_end": on_train_end} if clearml else {}

View File

@ -343,7 +343,7 @@ def compute_ap(recall, precision):
return ap, mpre, mrec
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=""):
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=Path(), names=(), eps=1e-16, prefix=""):
""" Compute the average precision, given the recall and precision curves.
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
# Arguments
@ -398,10 +398,10 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
names = dict(enumerate(names)) # to dict
if plot:
plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names)
plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1')
plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision')
plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall')
plot_pr_curve(px, py, ap, save_dir / f'{prefix}PR_curve.png', names)
plot_mc_curve(px, f1, save_dir / f'{prefix}F1_curve.png', names, ylabel='F1')
plot_mc_curve(px, p, save_dir / f'{prefix}P_curve.png', names, ylabel='Precision')
plot_mc_curve(px, r, save_dir / f'{prefix}R_curve.png', names, ylabel='Recall')
i = smooth(f1.mean(0), 0.1).argmax() # max F1 index
p, r, f1 = p[:, i], r[:, i], f1[:, i]