bert
bert copied to clipboard
fix Pose estimation (keypoint.ipynb) same error, output_to_keypoint()对bounding box解析有问题
使用yolov7-w6-pose.pt和keypoint.ipynb时,我想把bounding box也绘制到图像上,但是发现文件utils.plots.py中函数output_to_keypoint()对box解析有误。其中,文件utils.plots.py来自于(https://github.com/WongKinYiu/yolov7/blob/a4724d55924708b4cee0dacbc875545b4b3abe3f/utils/plots.py)
原始的函数代码段:
# utils/plots.py #L436-L444
def output_to_keypoint(output):
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
targets = []
for i, o in enumerate(output):
kpts = o[:,6:]
o = o[:,:6]
for index, (*box, conf, cls) in enumerate(o.cpu().numpy()):
targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf, *list(kpts.cpu().numpy()[index])])
return np.array(targets)
因为该函数代码段L443中,box格式为xyxy,不是cxywh,不需要执行转换函数xyxy2xywh(),于是将该行L443代码修改为
# targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf, *list(kpts.cpu().numpy()[index])])
targets.append([i, cls, *list(*box), conf, *list(kpts.cpu().numpy()[index])])
绘制bounding box到图像上,需要在文件keypoint.ipynb中添加如下代码
# add box plot_one_box()
from utils.plots import plot_one_box
output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)
output = output_to_keypoint(output)
nimg = image[0].permute(1, 2, 0) * 255
nimg = nimg.cpu().numpy().astype(np.uint8)
nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
for idx in range(output.shape[0]):
plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)
# add bbox
xyxy = output[idx, 2:6]
label = "cls%s-%.4f" % (int(output[idx, 1]), float(output[idx, 6]))
plot_one_box(xyxy, nimg, label=label, color=None, line_thickness=3)