log more than 1 image

Author: Erotemic

yolo/utils/callbacks/tensorboard_plotter.py

@@ -273,7 +273,6 @@ def _dump_measures(train_dpath, title='?name?', smoothing='auto', ignore_outlier
             plot_preferences.update(user_plot_preferences)
         else:
             plot_preferences = default_plot_preferences
-        print(f'plot_preferences = {ub.urepr(plot_preferences, nl=3)}')
 
         for item in plot_preferences['attributes']:
             item['pattern_'] = kwutil.util_pattern.MultiPattern.coerce(item['pattern'])

yolo/utils/logging_utils.py

@@ -213,18 +213,19 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b
         if batch_idx != 0:
             return
         batch_size, images, targets, rev_tensor, img_paths = batch
-        gt_boxes = targets[0] if targets.ndim == 3 else targets
-        pred_boxes = outputs[0] if isinstance(outputs, list) else outputs
-        images = [images[0]]
+        # gt_boxes = targets[0] if targets.ndim == 3 else targets
+        # pred_boxes = outputs[0] if isinstance(outputs, list) else outputs
+        # images = images[:]
         step = trainer.current_epoch
 
-        for logger in trainer.loggers:
-            if isinstance(logger, WandbLogger):
+        for _logger in trainer.loggers:
+            if isinstance(_logger, WandbLogger):
                 # FIXME: not robust to configured image sizes, need to know
                 # that info.
-                logger.log_image("Input Image", images, step=step)
-                logger.log_image("Ground Truth", images, step=step, boxes=[log_bbox(gt_boxes)])
-                logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)])
+                for image, gt_boxes, pred_boxes in zip(images, targets, outputs):
+                    _logger.log_image("Input Image", [image], step=step)
+                    _logger.log_image("Ground Truth", [image], step=step, boxes=[log_bbox(gt_boxes)])
+                    _logger.log_image("Prediction", [image], step=step, boxes=[log_bbox(pred_boxes)])
 
         # TODO: better config
         import os
@@ -243,6 +244,8 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b
 
             for bx in range(len(images)):
                 image_chw = images[bx].data.cpu().numpy()
+                gt_boxes = targets[bx]
+                pred_boxes = outputs[bx]
                 image_hwc = einops.rearrange(image_chw, 'c h w -> h w c')
                 image_hwc = kwimage.ensure_uint255(image_hwc)