Fix torchvision dependency for ClearML logging (#8993)

* Replace torchvision annotator with native local Annotator

* Add colored bounding boxes and rewrite for efficiency

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix pep8 issue

* Update clearml_utils.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update clearml_utils.py

* Reduced conf_threshold to 0.25 to match defaults

The 0.25 default conf threshold value is used a default throughout YOLOv5 (i.e. in detect.py) and also in other tools like CoreML

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <[email protected]>

Author: 3 people

utils/loggers/clearml/clearml_utils.py

@@ -3,20 +3,22 @@
 import re
 from pathlib import Path
 
+import numpy as np
 import yaml
-from torchvision.transforms import ToPILImage
+
+from utils.plots import Annotator, colors
 
 try:
     import clearml
     from clearml import Dataset, Task
-    from torchvision.utils import draw_bounding_boxes  # WARNING: requires torchvision>=0.9.0
-
     assert hasattr(clearml, '__version__')  # verify package import not local dir
 except (ImportError, AssertionError):
     clearml = None
 
 
 def construct_dataset(clearml_info_string):
+    """Load in a clearml dataset and fill the internal data_dict with its contents.
+    """
     dataset_id = clearml_info_string.replace('clearml://', '')
     dataset = Dataset.get(dataset_id=dataset_id)
     dataset_root_path = Path(dataset.get_local_copy())
@@ -120,9 +122,9 @@ def log_debug_samples(self, files, title='Debug Samples'):
                                                     local_path=str(f),
                                                     iteration=iteration)
 
-    def log_image_with_boxes(self, image_path, boxes, class_names, image):
+    def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_threshold=0.25):
         """
-        Draw the bounding boxes on a single image and report the result as a ClearML debug sample
+        Draw the bounding boxes on a single image and report the result as a ClearML debug sample.
 
         arguments:
         image_path (PosixPath) the path the original image file
@@ -133,16 +135,20 @@ def log_image_with_boxes(self, image_path, boxes, class_names, image):
         if len(self.current_epoch_logged_images) < self.max_imgs_to_log_per_epoch and self.current_epoch >= 0:
             # Log every bbox_interval times and deduplicate for any intermittend extra eval runs
             if self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images:
-                converter = ToPILImage()
-                labels = []
-                for conf, class_nr in zip(boxes[:, 4], boxes[:, 5]):
+                im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2))
+                annotator = Annotator(im=im, pil=True)
+                for i, (conf, class_nr, box) in enumerate(zip(boxes[:, 4], boxes[:, 5], boxes[:, :4])):
+                    color = colors(i)
+
                     class_name = class_names[int(class_nr)]
                     confidence = round(float(conf) * 100, 2)
-                    labels.append(f"{class_name}: {confidence}%")
-                annotated_image = converter(
-                    draw_bounding_boxes(image=image.mul(255).clamp(0, 255).byte().cpu(),
-                                        boxes=boxes[:, :4],
-                                        labels=labels))
+                    label = f"{class_name}: {confidence}%"
+
+                    if confidence > conf_threshold:
+                        annotator.rectangle(box.cpu().numpy(), outline=color)
+                        annotator.box_label(box.cpu().numpy(), label=label, color=color)
+
+                annotated_image = annotator.result()
                 self.task.get_logger().report_image(title='Bounding Boxes',
                                                     series=image_path.name,
                                                     iteration=self.current_epoch,