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 <glenn.jocher@ultralytics.com>
This commit is contained in:
1 file changed
+19
-13
@@ -3,20 +3,22 @@ import glob
|
||||
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 @@ class ClearmlLogger:
|
||||
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 @@ class ClearmlLogger:
|
||||
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,
|
||||
|
||||
Reference in new issue
Block a user