yolov7 icon indicating copy to clipboard operation
yolov7 copied to clipboard

fix Pose estimation (keypoint.ipynb) some error, output_to_keypoint()对bounding box解析有问题

Open meiqicheng opened this issue 2 years ago • 0 comments

使用yolov7-w6-pose.ptkeypoint.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)

meiqicheng avatar Aug 05 '22 08:08 meiqicheng