Suppress wandb images size mismatch warning (ultralytics#3611)

* supress wandb images size mismatch warning * supress wandb images size mismatch warning * PEP8 reformat and optimize imports Co-authored-by: Glenn Jocher <[email protected]>
fhkiel-mlaip · Jun 14, 2021 · 6c08786 · 6c08786
1 parent cbdf2b4
commit 6c08786
Showing 1 changed file with 25 additions and 7 deletions.
diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py
@@ -1,16 +1,16 @@
 """Utilities and tools for tracking runs with Weights & Biases."""
-import json
+import logging
 import sys
+from contextlib import contextmanager
 from pathlib import Path
 
-import torch
 import yaml
 from tqdm import tqdm
 
 sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
 from utils.datasets import LoadImagesAndLabels
 from utils.datasets import img2label_paths
-from utils.general import colorstr, xywh2xyxy, check_dataset, check_file
+from utils.general import colorstr, check_dataset, check_file
 
 try:
     import wandb
@@ -92,6 +92,7 @@ class WandbLogger():
     For more on how this logger is used, see the Weights & Biases documentation:
     https://docs.wandb.com/guides/integrations/yolov5
     """
+
     def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
         # Pre-training routine --
         self.job_type = job_type
@@ -272,7 +273,7 @@ def create_dataset_table(self, dataset, class_to_id, name='dataset'):
                                  "box_caption": "%s" % (class_to_id[cls])})
                 img_classes[cls] = class_to_id[cls]
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
-            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes),
+            table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
                            Path(paths).name)
         artifact.add(table, name)
         return artifact
@@ -306,8 +307,9 @@ def log(self, log_dict):
 
     def end_epoch(self, best_result=False):
         if self.wandb_run:
-            wandb.log(self.log_dict)
-            self.log_dict = {}
+            with all_logging_disabled():
+                wandb.log(self.log_dict)
+                self.log_dict = {}
             if self.result_artifact:
                 train_results = wandb.JoinedTable(self.val_table, self.result_table, "id")
                 self.result_artifact.add(train_results, 'result')
@@ -319,5 +321,21 @@ def end_epoch(self, best_result=False):
     def finish_run(self):
         if self.wandb_run:
             if self.log_dict:
-                wandb.log(self.log_dict)
+                with all_logging_disabled():
+                    wandb.log(self.log_dict)
             wandb.run.finish()
+
+
+@contextmanager
+def all_logging_disabled(highest_level=logging.CRITICAL):
+    """ source - https://gist.github.com/simon-weber/7853144
+    A context manager that will prevent any logging messages triggered during the body from being processed.
+    :param highest_level: the maximum logging level in use.
+      This would only need to be changed if a custom level greater than CRITICAL is defined.
+    """
+    previous_level = logging.root.manager.disable
+    logging.disable(highest_level)
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)