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
|
import re
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
import yaml
|
import yaml
|
||||||
from torchvision.transforms import ToPILImage
|
|
||||||
|
from utils.plots import Annotator, colors
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import clearml
|
import clearml
|
||||||
from clearml import Dataset, Task
|
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
|
assert hasattr(clearml, '__version__') # verify package import not local dir
|
||||||
except (ImportError, AssertionError):
|
except (ImportError, AssertionError):
|
||||||
clearml = None
|
clearml = None
|
||||||
|
|
||||||
|
|
||||||
def construct_dataset(clearml_info_string):
|
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_id = clearml_info_string.replace('clearml://', '')
|
||||||
dataset = Dataset.get(dataset_id=dataset_id)
|
dataset = Dataset.get(dataset_id=dataset_id)
|
||||||
dataset_root_path = Path(dataset.get_local_copy())
|
dataset_root_path = Path(dataset.get_local_copy())
|
||||||
@@ -120,9 +122,9 @@ class ClearmlLogger:
|
|||||||
local_path=str(f),
|
local_path=str(f),
|
||||||
iteration=iteration)
|
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:
|
arguments:
|
||||||
image_path (PosixPath) the path the original image file
|
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:
|
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
|
# 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:
|
if self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images:
|
||||||
converter = ToPILImage()
|
im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2))
|
||||||
labels = []
|
annotator = Annotator(im=im, pil=True)
|
||||||
for conf, class_nr in zip(boxes[:, 4], boxes[:, 5]):
|
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)]
|
class_name = class_names[int(class_nr)]
|
||||||
confidence = round(float(conf) * 100, 2)
|
confidence = round(float(conf) * 100, 2)
|
||||||
labels.append(f"{class_name}: {confidence}%")
|
label = f"{class_name}: {confidence}%"
|
||||||
annotated_image = converter(
|
|
||||||
draw_bounding_boxes(image=image.mul(255).clamp(0, 255).byte().cpu(),
|
if confidence > conf_threshold:
|
||||||
boxes=boxes[:, :4],
|
annotator.rectangle(box.cpu().numpy(), outline=color)
|
||||||
labels=labels))
|
annotator.box_label(box.cpu().numpy(), label=label, color=color)
|
||||||
|
|
||||||
|
annotated_image = annotator.result()
|
||||||
self.task.get_logger().report_image(title='Bounding Boxes',
|
self.task.get_logger().report_image(title='Bounding Boxes',
|
||||||
series=image_path.name,
|
series=image_path.name,
|
||||||
iteration=self.current_epoch,
|
iteration=self.current_epoch,
|
||||||
|
|||||||
Reference in new issue
Block a user