Add plot_images conf_thresh parameter (#8446)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Dean Mark 2024-02-25 16:46:56 +02:00 committed by GitHub
parent 30729d7346
commit 2d75f72598
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -713,6 +713,7 @@ def plot_images(
on_plot=None,
max_subplots=16,
save=True,
conf_thres=0.25,
):
"""Plot image grid with labels."""
if isinstance(images, torch.Tensor):
@ -778,7 +779,7 @@ def plot_images(
c = classes[j]
color = colors(c)
c = names.get(c, c) if names else c
if labels or conf[j] > 0.25: # 0.25 conf thresh
if labels or conf[j] > conf_thres:
label = f"{c}" if labels else f"{c} {conf[j]:.1f}"
annotator.box_label(box, label, color=color, rotated=is_obb)
@ -800,7 +801,7 @@ def plot_images(
kpts_[..., 0] += x
kpts_[..., 1] += y
for j in range(len(kpts_)):
if labels or conf[j] > 0.25: # 0.25 conf thresh
if labels or conf[j] > conf_thres:
annotator.kpts(kpts_[j])
# Plot masks
@ -816,7 +817,7 @@ def plot_images(
im = np.asarray(annotator.im).copy()
for j in range(len(image_masks)):
if labels or conf[j] > 0.25: # 0.25 conf thresh
if labels or conf[j] > conf_thres:
color = colors(classes[j])
mh, mw = image_masks[j].shape
if mh != h or mw != w: