ultralytics 8.2.77 new color_mode=instance plot arg (#15034)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-08-14 00:45:53 +08:00 committed by GitHub
parent b7c5db94b4
commit 265166ae4d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 32 additions and 9 deletions

View file

@ -369,7 +369,7 @@ class Annotator:
# Convert im back to PIL and update draw
self.fromarray(self.im)
def kpts(self, kpts, shape=(640, 640), radius=5, kpt_line=True, conf_thres=0.25):
def kpts(self, kpts, shape=(640, 640), radius=5, kpt_line=True, conf_thres=0.25, kpt_color=None):
"""
Plot keypoints on the image.
@ -379,6 +379,7 @@ class Annotator:
radius (int, optional): Radius of the drawn keypoints. Default is 5.
kpt_line (bool, optional): If True, the function will draw lines connecting keypoints
for human pose. Default is True.
kpt_color (tuple, optional): The color of the keypoints (B, G, R).
Note:
`kpt_line=True` currently only supports human pose plotting.
@ -391,7 +392,7 @@ class Annotator:
is_pose = nkpt == 17 and ndim in {2, 3}
kpt_line &= is_pose # `kpt_line=True` for now only supports human pose plotting
for i, k in enumerate(kpts):
color_k = [int(x) for x in self.kpt_color[i]] if is_pose else colors(i)
color_k = kpt_color or (self.kpt_color[i].tolist() if is_pose else colors(i))
x_coord, y_coord = k[0], k[1]
if x_coord % shape[1] != 0 and y_coord % shape[0] != 0:
if len(k) == 3:
@ -414,7 +415,14 @@ class Annotator:
continue
if pos2[0] % shape[1] == 0 or pos2[1] % shape[0] == 0 or pos2[0] < 0 or pos2[1] < 0:
continue
cv2.line(self.im, pos1, pos2, [int(x) for x in self.limb_color[i]], thickness=2, lineType=cv2.LINE_AA)
cv2.line(
self.im,
pos1,
pos2,
kpt_color or self.limb_color[i].tolist(),
thickness=2,
lineType=cv2.LINE_AA,
)
if self.pil:
# Convert im back to PIL and update draw
self.fromarray(self.im)