You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
443 lines
20 KiB
443 lines
20 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
import contextlib
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from PIL import __version__ as pil_version
|
|
|
|
from ultralytics.yolo.utils import LOGGER, TryExcept, plt_settings, threaded
|
|
|
|
from .checks import check_font, check_version, is_ascii
|
|
from .files import increment_path
|
|
from .ops import clip_boxes, scale_image, xywh2xyxy, xyxy2xywh
|
|
|
|
|
|
class Colors:
|
|
# Ultralytics color palette https://ultralytics.com/
|
|
def __init__(self):
|
|
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
|
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
|
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
|
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
|
|
self.n = len(self.palette)
|
|
self.pose_palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
|
|
[153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
|
|
[255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
|
|
[51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]],
|
|
dtype=np.uint8)
|
|
|
|
def __call__(self, i, bgr=False):
|
|
c = self.palette[int(i) % self.n]
|
|
return (c[2], c[1], c[0]) if bgr else c
|
|
|
|
@staticmethod
|
|
def hex2rgb(h): # rgb order (PIL)
|
|
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
|
|
|
|
|
|
colors = Colors() # create instance for 'from utils.plots import colors'
|
|
|
|
|
|
class Annotator:
|
|
# YOLOv8 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
|
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
|
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
|
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
|
self.pil = pil or non_ascii
|
|
if self.pil: # use PIL
|
|
self.pil_9_2_0_check = check_version(pil_version, '9.2.0') # deprecation check
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
self.draw = ImageDraw.Draw(self.im)
|
|
try:
|
|
font = check_font('Arial.Unicode.ttf' if non_ascii else font)
|
|
size = font_size or max(round(sum(self.im.size) / 2 * 0.035), 12)
|
|
self.font = ImageFont.truetype(str(font), size)
|
|
except Exception:
|
|
self.font = ImageFont.load_default()
|
|
else: # use cv2
|
|
self.im = im
|
|
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
|
# pose
|
|
self.skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9],
|
|
[8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
|
|
|
|
self.limb_color = colors.pose_palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]
|
|
self.kpt_color = colors.pose_palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]
|
|
|
|
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
|
# Add one xyxy box to image with label
|
|
if isinstance(box, torch.Tensor):
|
|
box = box.tolist()
|
|
if self.pil or not is_ascii(label):
|
|
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
|
if label:
|
|
if self.pil_9_2_0_check:
|
|
_, _, w, h = self.font.getbbox(label) # text width, height (New)
|
|
else:
|
|
w, h = self.font.getsize(label) # text width, height (Old, deprecated in 9.2.0)
|
|
outside = box[1] - h >= 0 # label fits outside box
|
|
self.draw.rectangle(
|
|
(box[0], box[1] - h if outside else box[1], box[0] + w + 1,
|
|
box[1] + 1 if outside else box[1] + h + 1),
|
|
fill=color,
|
|
)
|
|
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
|
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
|
else: # cv2
|
|
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
|
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
|
|
if label:
|
|
tf = max(self.lw - 1, 1) # font thickness
|
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
|
outside = p1[1] - h >= 3
|
|
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
|
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
|
cv2.putText(self.im,
|
|
label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
|
|
0,
|
|
self.lw / 3,
|
|
txt_color,
|
|
thickness=tf,
|
|
lineType=cv2.LINE_AA)
|
|
|
|
def masks(self, masks, colors, im_gpu, alpha=0.5, retina_masks=False):
|
|
"""Plot masks at once.
|
|
Args:
|
|
masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
|
colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
|
im_gpu (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
|
alpha (float): mask transparency: 0.0 fully transparent, 1.0 opaque
|
|
"""
|
|
if self.pil:
|
|
# convert to numpy first
|
|
self.im = np.asarray(self.im).copy()
|
|
if len(masks) == 0:
|
|
self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
|
if im_gpu.device != masks.device:
|
|
im_gpu = im_gpu.to(masks.device)
|
|
colors = torch.tensor(colors, device=masks.device, dtype=torch.float32) / 255.0 # shape(n,3)
|
|
colors = colors[:, None, None] # shape(n,1,1,3)
|
|
masks = masks.unsqueeze(3) # shape(n,h,w,1)
|
|
masks_color = masks * (colors * alpha) # shape(n,h,w,3)
|
|
|
|
inv_alph_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
|
|
mcs = (masks_color * inv_alph_masks).sum(0) * 2 # mask color summand shape(n,h,w,3)
|
|
|
|
im_gpu = im_gpu.flip(dims=[0]) # flip channel
|
|
im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
|
|
im_gpu = im_gpu * inv_alph_masks[-1] + mcs
|
|
im_mask = (im_gpu * 255)
|
|
im_mask_np = im_mask.byte().cpu().numpy()
|
|
self.im[:] = im_mask_np if retina_masks else scale_image(im_mask_np, self.im.shape)
|
|
if self.pil:
|
|
# convert im back to PIL and update draw
|
|
self.fromarray(self.im)
|
|
|
|
def kpts(self, kpts, shape=(640, 640), radius=5, kpt_line=True):
|
|
"""Plot keypoints.
|
|
Args:
|
|
kpts (tensor): predicted kpts, shape: [17, 3]
|
|
shape (tuple): image shape, (h, w)
|
|
steps (int): keypoints step
|
|
radius (int): size of drawing points
|
|
"""
|
|
if self.pil:
|
|
# convert to numpy first
|
|
self.im = np.asarray(self.im).copy()
|
|
nkpt, ndim = kpts.shape
|
|
is_pose = nkpt == 17 and ndim == 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)
|
|
x_coord, y_coord = k[0], k[1]
|
|
if x_coord % shape[1] != 0 and y_coord % shape[0] != 0:
|
|
if len(k) == 3:
|
|
conf = k[2]
|
|
if conf < 0.5:
|
|
continue
|
|
cv2.circle(self.im, (int(x_coord), int(y_coord)), radius, color_k, -1, lineType=cv2.LINE_AA)
|
|
|
|
if kpt_line:
|
|
ndim = kpts.shape[-1]
|
|
for i, sk in enumerate(self.skeleton):
|
|
pos1 = (int(kpts[(sk[0] - 1), 0]), int(kpts[(sk[0] - 1), 1]))
|
|
pos2 = (int(kpts[(sk[1] - 1), 0]), int(kpts[(sk[1] - 1), 1]))
|
|
if ndim == 3:
|
|
conf1 = kpts[(sk[0] - 1), 2]
|
|
conf2 = kpts[(sk[1] - 1), 2]
|
|
if conf1 < 0.5 or conf2 < 0.5:
|
|
continue
|
|
if pos1[0] % shape[1] == 0 or pos1[1] % shape[0] == 0 or pos1[0] < 0 or pos1[1] < 0:
|
|
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)
|
|
if self.pil:
|
|
# convert im back to PIL and update draw
|
|
self.fromarray(self.im)
|
|
|
|
def rectangle(self, xy, fill=None, outline=None, width=1):
|
|
# Add rectangle to image (PIL-only)
|
|
self.draw.rectangle(xy, fill, outline, width)
|
|
|
|
def text(self, xy, text, txt_color=(255, 255, 255), anchor='top'):
|
|
# Add text to image (PIL-only)
|
|
if anchor == 'bottom': # start y from font bottom
|
|
w, h = self.font.getsize(text) # text width, height
|
|
xy[1] += 1 - h
|
|
if self.pil:
|
|
self.draw.text(xy, text, fill=txt_color, font=self.font)
|
|
else:
|
|
tf = max(self.lw - 1, 1) # font thickness
|
|
cv2.putText(self.im, text, xy, 0, self.lw / 3, txt_color, thickness=tf, lineType=cv2.LINE_AA)
|
|
|
|
def fromarray(self, im):
|
|
# Update self.im from a numpy array
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
self.draw = ImageDraw.Draw(self.im)
|
|
|
|
def result(self):
|
|
# Return annotated image as array
|
|
return np.asarray(self.im)
|
|
|
|
|
|
@TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395
|
|
@plt_settings()
|
|
def plot_labels(boxes, cls, names=(), save_dir=Path('')):
|
|
import pandas as pd
|
|
import seaborn as sn
|
|
|
|
# plot dataset labels
|
|
LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ")
|
|
b = boxes.transpose() # classes, boxes
|
|
nc = int(cls.max() + 1) # number of classes
|
|
x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
|
|
|
|
# seaborn correlogram
|
|
sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
|
|
plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
|
|
plt.close()
|
|
|
|
# matplotlib labels
|
|
ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
|
|
y = ax[0].hist(cls, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
|
|
with contextlib.suppress(Exception): # color histogram bars by class
|
|
[y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # known issue #3195
|
|
ax[0].set_ylabel('instances')
|
|
if 0 < len(names) < 30:
|
|
ax[0].set_xticks(range(len(names)))
|
|
ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10)
|
|
else:
|
|
ax[0].set_xlabel('classes')
|
|
sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
|
|
sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
|
|
|
|
# rectangles
|
|
boxes[:, 0:2] = 0.5 # center
|
|
boxes = xywh2xyxy(boxes) * 1000
|
|
img = Image.fromarray(np.ones((1000, 1000, 3), dtype=np.uint8) * 255)
|
|
for cls, box in zip(cls[:500], boxes[:500]):
|
|
ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot
|
|
ax[1].imshow(img)
|
|
ax[1].axis('off')
|
|
|
|
for a in [0, 1, 2, 3]:
|
|
for s in ['top', 'right', 'left', 'bottom']:
|
|
ax[a].spines[s].set_visible(False)
|
|
|
|
plt.savefig(save_dir / 'labels.jpg', dpi=200)
|
|
plt.close()
|
|
|
|
|
|
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
|
|
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
|
b = xyxy2xywh(xyxy.view(-1, 4)) # boxes
|
|
if square:
|
|
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
|
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
|
xyxy = xywh2xyxy(b).long()
|
|
clip_boxes(xyxy, im.shape)
|
|
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
|
if save:
|
|
file.parent.mkdir(parents=True, exist_ok=True) # make directory
|
|
f = str(increment_path(file).with_suffix('.jpg'))
|
|
# cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
|
|
Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0) # save RGB
|
|
return crop
|
|
|
|
|
|
@threaded
|
|
def plot_images(images,
|
|
batch_idx,
|
|
cls,
|
|
bboxes,
|
|
masks=np.zeros(0, dtype=np.uint8),
|
|
kpts=np.zeros((0, 51), dtype=np.float32),
|
|
paths=None,
|
|
fname='images.jpg',
|
|
names=None):
|
|
# Plot image grid with labels
|
|
if isinstance(images, torch.Tensor):
|
|
images = images.cpu().float().numpy()
|
|
if isinstance(cls, torch.Tensor):
|
|
cls = cls.cpu().numpy()
|
|
if isinstance(bboxes, torch.Tensor):
|
|
bboxes = bboxes.cpu().numpy()
|
|
if isinstance(masks, torch.Tensor):
|
|
masks = masks.cpu().numpy().astype(int)
|
|
if isinstance(kpts, torch.Tensor):
|
|
kpts = kpts.cpu().numpy()
|
|
if isinstance(batch_idx, torch.Tensor):
|
|
batch_idx = batch_idx.cpu().numpy()
|
|
|
|
max_size = 1920 # max image size
|
|
max_subplots = 16 # max image subplots, i.e. 4x4
|
|
bs, _, h, w = images.shape # batch size, _, height, width
|
|
bs = min(bs, max_subplots) # limit plot images
|
|
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
|
if np.max(images[0]) <= 1:
|
|
images *= 255 # de-normalise (optional)
|
|
|
|
# Build Image
|
|
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
|
for i, im in enumerate(images):
|
|
if i == max_subplots: # if last batch has fewer images than we expect
|
|
break
|
|
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
|
im = im.transpose(1, 2, 0)
|
|
mosaic[y:y + h, x:x + w, :] = im
|
|
|
|
# Resize (optional)
|
|
scale = max_size / ns / max(h, w)
|
|
if scale < 1:
|
|
h = math.ceil(scale * h)
|
|
w = math.ceil(scale * w)
|
|
mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
|
|
|
|
# Annotate
|
|
fs = int((h + w) * ns * 0.01) # font size
|
|
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
|
|
for i in range(i + 1):
|
|
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
|
annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
|
|
if paths:
|
|
annotator.text((x + 5, y + 5), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
|
|
if len(cls) > 0:
|
|
idx = batch_idx == i
|
|
|
|
boxes = xywh2xyxy(bboxes[idx, :4]).T
|
|
classes = cls[idx].astype('int')
|
|
labels = bboxes.shape[1] == 4 # labels if no conf column
|
|
conf = None if labels else bboxes[idx, 4] # check for confidence presence (label vs pred)
|
|
|
|
if boxes.shape[1]:
|
|
if boxes.max() <= 1.01: # if normalized with tolerance 0.01
|
|
boxes[[0, 2]] *= w # scale to pixels
|
|
boxes[[1, 3]] *= h
|
|
elif scale < 1: # absolute coords need scale if image scales
|
|
boxes *= scale
|
|
boxes[[0, 2]] += x
|
|
boxes[[1, 3]] += y
|
|
for j, box in enumerate(boxes.T.tolist()):
|
|
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
|
|
label = f'{c}' if labels else f'{c} {conf[j]:.1f}'
|
|
annotator.box_label(box, label, color=color)
|
|
|
|
# Plot keypoints
|
|
if len(kpts):
|
|
kpts_ = kpts[idx].copy()
|
|
if len(kpts_):
|
|
if kpts_[..., 0].max() <= 1.01 or kpts_[..., 1].max() <= 1.01: # if normalized with tolerance .01
|
|
kpts_[..., 0] *= w # scale to pixels
|
|
kpts_[..., 1] *= h
|
|
elif scale < 1: # absolute coords need scale if image scales
|
|
kpts_ *= scale
|
|
kpts_[..., 0] += x
|
|
kpts_[..., 1] += y
|
|
for j in range(len(kpts_)):
|
|
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
|
annotator.kpts(kpts_[j])
|
|
|
|
# Plot masks
|
|
if len(masks):
|
|
if idx.shape[0] == masks.shape[0]: # overlap_masks=False
|
|
image_masks = masks[idx]
|
|
else: # overlap_masks=True
|
|
image_masks = masks[[i]] # (1, 640, 640)
|
|
nl = idx.sum()
|
|
index = np.arange(nl).reshape((nl, 1, 1)) + 1
|
|
image_masks = np.repeat(image_masks, nl, axis=0)
|
|
image_masks = np.where(image_masks == index, 1.0, 0.0)
|
|
|
|
im = np.asarray(annotator.im).copy()
|
|
for j, box in enumerate(boxes.T.tolist()):
|
|
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
|
color = colors(classes[j])
|
|
mh, mw = image_masks[j].shape
|
|
if mh != h or mw != w:
|
|
mask = image_masks[j].astype(np.uint8)
|
|
mask = cv2.resize(mask, (w, h))
|
|
mask = mask.astype(bool)
|
|
else:
|
|
mask = image_masks[j].astype(bool)
|
|
with contextlib.suppress(Exception):
|
|
im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6
|
|
annotator.fromarray(im)
|
|
annotator.im.save(fname) # save
|
|
|
|
|
|
@plt_settings()
|
|
def plot_results(file='path/to/results.csv', dir='', segment=False, pose=False):
|
|
# Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
|
|
import pandas as pd
|
|
save_dir = Path(file).parent if file else Path(dir)
|
|
if segment:
|
|
fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
|
|
index = [1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]
|
|
elif pose:
|
|
fig, ax = plt.subplots(2, 9, figsize=(21, 6), tight_layout=True)
|
|
index = [1, 2, 3, 4, 5, 6, 7, 10, 11, 14, 15, 16, 17, 18, 8, 9, 12, 13]
|
|
else:
|
|
fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
|
|
index = [1, 2, 3, 4, 5, 8, 9, 10, 6, 7]
|
|
ax = ax.ravel()
|
|
files = list(save_dir.glob('results*.csv'))
|
|
assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.'
|
|
for f in files:
|
|
try:
|
|
data = pd.read_csv(f)
|
|
s = [x.strip() for x in data.columns]
|
|
x = data.values[:, 0]
|
|
for i, j in enumerate(index):
|
|
y = data.values[:, j].astype('float')
|
|
# y[y == 0] = np.nan # don't show zero values
|
|
ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8)
|
|
ax[i].set_title(s[j], fontsize=12)
|
|
# if j in [8, 9, 10]: # share train and val loss y axes
|
|
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
|
|
except Exception as e:
|
|
LOGGER.warning(f'WARNING: Plotting error for {f}: {e}')
|
|
ax[1].legend()
|
|
fig.savefig(save_dir / 'results.png', dpi=200)
|
|
plt.close()
|
|
|
|
|
|
def output_to_target(output, max_det=300):
|
|
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting
|
|
targets = []
|
|
for i, o in enumerate(output):
|
|
box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
|
|
j = torch.full((conf.shape[0], 1), i)
|
|
targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
|
|
targets = torch.cat(targets, 0).numpy()
|
|
return targets[:, 0], targets[:, 1], targets[:, 2:]
|