From dc13820c9d561bf112e773795cd75d7c40dbbff7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 30 Aug 2021 13:01:41 +0200 Subject: [PATCH] TensorBoard `on_train_end()` speed improvements (#4605) --- utils/loggers/__init__.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index 3d67e93..16d0348 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -131,10 +131,9 @@ class Loggers(): 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 + import cv2 for f in files: - self.tb.add_image(f.stem, np.asarray(Image.open(f)), epoch, dataformats='HWC') + self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC') if self.wandb: self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})