一、前言

由于工程项目中需要对视频中的person进行关键点检测,我测试各个算法后,并没有采用比较应用化成熟的Openpose,决定采用检测精度更高的HRnet系列。但是由于官方给的算法只能测试数据集,需要自己根据算法模型编写实例化代码。
本文根据SimDR工程实现视频关键点检测。SimDR根据HRnet改进而来,整个工程既包括HRnet又包括改进后的算法,使用起来较为方便,而且本文仅在cpu上就可以跑通整个工程。

二、环境配置

python的环境主要就是按照工程中SimDR与yolov5的requirement.txt安装即可。总之缺啥装啥。

三、工程准备

1、克隆工程

git clone https://github.com/leeyegy/SimDR.git  #克隆姿态估计工程cd SimDRgit clone -b v5.0 https://github.com/ultralytics/yolov5.git #在姿态估计工程中添加yolov5算法

2、目标检测

①添加权重文件

添加yolov5x.pt(见评论区网盘)到‘ SimDR/yolov5/weights/ ’文件夹下。

②获取边界框

在yolov5文件夹下新建YOLOv5.py,复制以下内容到文件中。注意:根据大家的反馈,不同的电脑,导入yolov5相关包时会不同的方式,代码中我是from yolov5.xxx import xxx,但是有些可以不用前面的yolov5,大家自行尝试哈。一般出现No module xxx 都是有关yolov5 的包导入出错哈。

import argparseimport timefrom pathlib import Pathimport numpy as npimport cv2import torchimport torch.backends.cudnn as cudnnfrom numpy import randomimport sysimport osfrom yolov5.models.experimental import attempt_loadfrom yolov5.utils.datasets import LoadStreams, LoadImagesfrom yolov5.utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_pathfrom yolov5.utils.plots import plot_one_boxfrom yolov5.utils.torch_utils import select_device, load_classifier, time_synchronizedfrom  yolov5.utils.datasets import letterboxclass Yolov5():    def __init__(self, weights=None, opt=None, device=None):        """        @param weights:        @param save_txt:        @param opt:        @param device:        """        self.weights = weights        self.device = device        # save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run        # save_dir.mkdir(parents=True, exist_ok=True)  # make dir        self.img_size = 640        self.model = attempt_load(weights, map_location=self.device)        self.stride = int(self.model.stride.max())        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]        self.opt = opt    def detect(self,img0):        """        @param img0: 输入图片  shape=[h,w,3]        @return:        """        person_boxes = np.ones((6))        img = letterbox(img0, self.img_size, stride=self.stride)[0]        # Convert        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416        img = np.ascontiguousarray(img)        img = torch.from_numpy(img).to(self.device)        img = img.float()  # uint8 to fp16/32        img /= 255.0  # 0 - 255 to 0.0 - 1.0        if img.ndimension() == 3:            img = img.unsqueeze(0)        pred = self.model(img, augment=self.opt.augment)[0]        # Apply NMS        pred = non_max_suppression(pred, self.opt.conf_thres, self.opt.iou_thres, classes=self.opt.classes, agnostic=self.opt.agnostic_nms)        for i, det in enumerate(pred):            if len(det):                # Rescale boxes from img_size to im0 size                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()                boxes = reversed(det)                boxes = boxes.cpu().numpy() #2022.04.06修改,在GPU上跑boxes无法直接转numpy数据                #for i , box in enumerate(np.array(boxes)):                for i , box in enumerate(boxes):                    if int(box[-1]) == 0 and box[-2]>=0.7:                        person_boxes=np.vstack((person_boxes , box))        #                 label = f'{self.names[int(box[-1])]} {box[-2]:.2f}'        #                 print(label)        #                 plot_one_box(box, img0, label=label, color=self.colors[int(box[-1])], line_thickness=3)        # cv2.imwrite('result1.jpg',img0)        # print(s)        # print(person_boxes,np.ndim(person_boxes))        if np.ndim(person_boxes)>=2 :            person_boxes_result = person_boxes[1:]            boxes_result = person_boxes[1:,:4]        else:            person_boxes_result = []            boxes_result = []        return person_boxes_result,boxes_resultdef yolov5test(opt,path = ''):    detector = Yolov5(weights='weights/yolov5x.pt',opt=opt,device=torch.device('cpu'))    img0 = cv2.imread(path)    personboxes ,boxes= detector.detect(img0)    for i,(x1,y1,x2,y2) in enumerate(boxes):        print(x1,y1,x2,y2)    print(personboxes,'\n',boxes)if __name__ == '__main__':    parser = argparse.ArgumentParser()    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')    parser.add_argument('--augment', action='store_true', help='augmented inference')    parser.add_argument('--update', action='store_true', help='update all model')    parser.add_argument('--project', default='runs/detect', help='save results to project/name')    parser.add_argument('--name', default='exp', help='save results to project/name')    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')    opt = parser.parse_args()    print(opt)    # check_requirements(exclude=('pycocotools', 'thop'))    with torch.no_grad():        yolov5test(opt,'data/images/zidane.jpg')

③路径问题

本文代码是在pycharm中运行,yolov5工程的加入导致有些文件夹名称相同,pycharm会搞混,可能会出现某些包找不到。这里需要先运行一下YOLOv5.py脚本,根据报错改一下import的内容。举个例子,./SimDR/yolov5/models/experimental.py 文件中会出现图片中的问题

改成如下即可,其他的文件改法相同。

④添加SPPF模块

yolov5 v5.0工程中没有SPPF模块,此时我们需要在./SimDR/yolov5/models/common.py文件末尾加入以下代码。

import warningsclass SPPF(nn.Module):    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))        super().__init__()        c_ = c1 // 2  # hidden channels        self.cv1 = Conv(c1, c_, 1, 1)        self.cv2 = Conv(c_ * 4, c2, 1, 1)        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)    def forward(self, x):        x = self.cv1(x)        with warnings.catch_warnings():            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning            y1 = self.m(x)            y2 = self.m(y1)            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))

3、姿态估计

①添加权重

在SimDR文件夹下新建weight/hrnet文件夹,添加pose_hrnet_w48_384x288.pth等文件(见评论区网盘)

②修改yaml文件

SimDR/experiments/文件夹下是coco与mpii数据集的配置文件,本文以coco为例。

接下来,修改./SimDR/experiments/coco/hrnet/heatmap/w48_384x288_adam_lr1e-3.yaml文件中的TEST部分的MODEL_FILE路径,如图所示。(SimDR算法的配置文件同理改动。)

③获取关键点

在’ SimDR/ ‘文件夹下新建Point_detect.py ,复制以下内容到文件中。

注意:代码第12行的路径要改成自己yolov5工程的路径,有这条代码才能正常运行。

【2022.04.16更新:根据评论区的建议,为关键点增加置信度值,这个值我是根据模型输出经过softmax后取最大值(关键点坐标就是这个最大值的索引),仅供参考。根据这个置信度可以解决半身照也会绘制全部点的问题。】

import cv2import numpy as npimport torchfrom torchvision.transforms import transformsimport torch.nn.functional as Ffrom lib.config import cfgfrom yolov5.YOLOv5 import Yolov5from lib.utils.transforms import  flip_back_simdr,transform_preds,get_affine_transformfrom lib import modelsimport argparseimport syssys.path.insert(0, 'D:\\Study\\Pose Estimation\\SimDR\\yolov5')class Points():    def __init__(self,                 model_name='sa-simdr',                 resolution=(384,288),                 opt=None,                 yolo_weights_path="./yolov5/weights/yolov5x.pt",                ):        """        Initializes a new SimpleHRNet object.        HRNet (and YOLOv3) are initialized on the torch.device("device") and        its (their) pre-trained weights will be loaded from disk.        Args:            c (int): number of channels (when using HRNet model) or resnet size (when using PoseResNet model).            nof_joints (int): number of joints.            checkpoint_path (str): path to an official hrnet checkpoint or a checkpoint obtained with `train_coco.py`.            model_name (str): model name (HRNet or PoseResNet).                Valid names for HRNet are: `HRNet`, `hrnet`                Valid names for PoseResNet are: `PoseResNet`, `poseresnet`, `ResNet`, `resnet`                Default: "HRNet"            resolution (tuple): hrnet input resolution - format: (height, width).                Default: (384, 288)            interpolation (int): opencv interpolation algorithm.                Default: cv2.INTER_CUBIC            multiperson (bool): if True, multiperson detection will be enabled.                This requires the use of a people detector (like YOLOv3).                Default: True            return_heatmaps (bool): if True, heatmaps will be returned along with poses by self.predict.                Default: False            return_bounding_boxes (bool): if True, bounding boxes will be returned along with poses by self.predict.                Default: False            max_batch_size (int): maximum batch size used in hrnet inference.                Useless without multiperson=True.                Default: 16            yolo_model_def (str): path to yolo model definition file.                Default: "./model/detectors/yolo/config/yolov3.cfg"            yolo_class_path (str): path to yolo class definition file.                Default: "./model/detectors/yolo/data/coco.names"            yolo_weights_path (str): path to yolo pretrained weights file.                Default: "./model/detectors/yolo/weights/yolov3.weights.cfg"            device (:class:`torch.device`): the hrnet (and yolo) inference will be run on this device.                Default: torch.device("cpu")        """        self.model_name = model_name        self.resolution = resolution  # in the form (height, width) as in the original implementation        self.aspect_ratio = resolution[1]/resolution[0]        self.yolo_weights_path = yolo_weights_path        self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],                           [9, 10], [11, 12], [13, 14], [15, 16]]        self.device = torch.device(opt.device)        cfg.defrost()        if model_name in ('sa-simdr','sasimdr','sa_simdr'):            if resolution ==(384,288):                cfg.merge_from_file('./experiments/coco/hrnet/sa_simdr/w48_384x288_adam_lr1e-3_split1_5_sigma4.yaml')            elif resolution == (256,192):                cfg.merge_from_file('./experiments/coco/hrnet/sa_simdr/w48_256x192_adam_lr1e-3_split2_sigma4.yaml')            else:                raise ValueError('Wrong cfg file')        elif model_name in ('simdr'):                if resolution == (256, 192):                    cfg.merge_from_file('./experiments/coco/hrnet/simdr/nmt_w48_256x192_adam_lr1e-3.yaml')                else:                    raise ValueError('Wrong cfg file')        elif model_name in ('hrnet','HRnet','Hrnet'):            if resolution == (384,288):                cfg.merge_from_file('./experiments/coco/hrnet/heatmap/w48_384x288_adam_lr1e-3.yaml')            elif resolution == (256,192):                cfg.merge_from_file('./experiments/coco/hrnet/heatmap/w48_256x192_adam_lr1e-3.yaml')            else:                raise ValueError('Wrong cfg file')        else:            raise ValueError('Wrong model name.')        cfg.freeze()        self.model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(            cfg, is_train=False)        print('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))        checkpoint_path = cfg.TEST.MODEL_FILE        checkpoint = torch.load(checkpoint_path, map_location=self.device)        if 'model' in checkpoint:            self.model.load_state_dict(checkpoint['model'])        else:            self.model.load_state_dict(checkpoint)        if 'cuda' in str(self.device):            print("device: 'cuda' - ", end="")            if 'cuda' == str(self.device):                # if device is set to 'cuda', all available GPUs will be used                print("%d GPU(s) will be used" % torch.cuda.device_count())                device_ids = None            else:                # if device is set to 'cuda:IDS', only that/those device(s) will be used                print("GPU(s) '%s' will be used" % str(self.device))                device_ids = [int(x) for x in str(self.device)[5:].split(',')]        elif 'cpu' == str(self.device):            print("device: 'cpu'")        else:            raise ValueError('Wrong device name.')        self.model = self.model.to(self.device)        self.model.eval()        self.detector = Yolov5(                               weights=yolo_weights_path,                               opt=opt ,                               device=self.device)        self.transform = transforms.Compose([            transforms.ToPILImage(),            transforms.Resize((self.resolution[0], self.resolution[1])),  # (height, width)            transforms.ToTensor(),            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),        ])    def _box2cs(self, box):        x, y, w, h = box[:4]        return self._xywh2cs(x, y, w, h)    def _xywh2cs(self, x, y, w, h):        center = np.zeros((2), dtype=np.float32)        center[0] = x + w * 0.5        center[1] = y + h * 0.5        if w > self.aspect_ratio * h:            h = w * 1.0 / self.aspect_ratio        elif w  0:                images = images.to(self.device)                with torch.no_grad():                    output_x, output_y = self.model(images)                    if cfg.TEST.FLIP_TEST:                        input_flipped = images.flip(3)                        output_x_flipped_, output_y_flipped_ = self.model(input_flipped)                        output_x_flipped = flip_back_simdr(output_x_flipped_.cpu().numpy(),                                                           self.flip_pairs, type='x')                        output_y_flipped = flip_back_simdr(output_y_flipped_.cpu().numpy(),                                                           self.flip_pairs, type='y')                        output_x_flipped = torch.from_numpy(output_x_flipped.copy()).to(self.device)                        output_y_flipped = torch.from_numpy(output_y_flipped.copy()).to(self.device)                        # feature is not aligned, shift flipped heatmap for higher accuracy                        if cfg.TEST.SHIFT_HEATMAP:                            output_x_flipped[:, :, 0:-1] = \                                output_x_flipped.clone()[:, :, 1:]                        output_x = F.softmax((output_x + output_x_flipped) * 0.5, dim=2)                        output_y = F.softmax((output_y + output_y_flipped) * 0.5, dim=2)                    else:                        output_x = F.softmax(output_x, dim=2)                        output_y = F.softmax(output_y, dim=2)                    max_val_x, preds_x = output_x.max(2, keepdim=True)                    max_val_y, preds_y = output_y.max(2, keepdim=True)                    mask = max_val_x > max_val_y                    max_val_x[mask] = max_val_y[mask]                    maxvals = max_val_x * 10.0                    output = torch.ones([images.size(0), preds_x.size(1), 3])                    output[:, :, 0] = torch.squeeze(torch.true_divide(preds_x, cfg.MODEL.SIMDR_SPLIT_RATIO))                    output[:, :, 1] = torch.squeeze(torch.true_divide(preds_y, cfg.MODEL.SIMDR_SPLIT_RATIO))                    # output[:, :, 2] = maxvals.squeeze(2)                    output = output.cpu().numpy()                    preds = output.copy()                    for i in range(output.shape[0]):                        preds[i] = transform_preds(                            output[i], c[i], s[i], [cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1]]                        )                    preds[:, :, 2] = maxvals.squeeze(2)            else:                preds = np.empty((0, 0, 3), dtype=np.float32)        return preds    def simdr_pts(self,img,detection,images,boxes):        c, s = [], []        if detection is not None:            for i, (x1, y1, x2, y2) in enumerate(detection):                x1 = int(round(x1.item()))                x2 = int(round(x2.item()))                y1 = int(round(y1.item()))                y2 = int(round(y2.item()))                boxes[i] = [x1, y1, x2, y2]                w, h = x2 - x1, y2 - y1                xx1 = np.max((0, x1))                yy1 = np.max((0, y1))                xx2 = np.min((img.shape[1] - 1, x1 + np.max((0, w - 1))))                yy2 = np.min((img.shape[0] - 1, y1 + np.max((0, h - 1))))                box = [xx1, yy1, xx2 - xx1, yy2 - yy1]                center, scale = self._box2cs(box)                c.append(center)                s.append(scale)                trans = get_affine_transform(center, scale, 0, np.array(cfg.MODEL.IMAGE_SIZE))                input = cv2.warpAffine(                    img,                    trans,                    (int(self.resolution[1]), int(self.resolution[0])),                    flags=cv2.INTER_LINEAR)                images[i] = self.transform(input)            if images.shape[0] > 0:                images = images.to(self.device)                with torch.no_grad():                    output_x, output_y = self.model(images)                    if cfg.TEST.FLIP_TEST:                        input_flipped = images.flip(3)                        output_x_flipped_, output_y_flipped_ = self.model(input_flipped)                        output_x_flipped = flip_back_simdr(output_x_flipped_.cpu().numpy(),                                                           self.flip_pairs, type='x')                        output_y_flipped = flip_back_simdr(output_y_flipped_.cpu().numpy(),                                                           self.flip_pairs, type='y')                        output_x_flipped = torch.from_numpy(output_x_flipped.copy()).to(self.device)                        output_y_flipped = torch.from_numpy(output_y_flipped.copy()).to(self.device)                        # feature is not aligned, shift flipped heatmap for higher accuracy                        if cfg.TEST.SHIFT_HEATMAP:                            output_x_flipped[:, :, 0:-1] = \                                output_x_flipped.clone()[:, :, 1:]                        output_x = (F.softmax(output_x, dim=2) + F.softmax(output_x_flipped, dim=2)) * 0.5                        output_y = (F.softmax(output_y, dim=2) + F.softmax(output_y_flipped, dim=2)) * 0.5                    else:                        output_x = F.softmax(output_x, dim=2)                        output_y = F.softmax(output_y, dim=2)                    max_val_x, preds_x = output_x.max(2, keepdim=True)                    max_val_y, preds_y = output_y.max(2, keepdim=True)                    mask = max_val_x > max_val_y                    max_val_x[mask] = max_val_y[mask]                    maxvals = max_val_x * 10.0                    output = torch.ones([images.size(0), preds_x.size(1), 3])                    output[:, :, 0] = torch.squeeze(torch.true_divide(preds_x, cfg.MODEL.SIMDR_SPLIT_RATIO))                    output[:, :, 1] = torch.squeeze(torch.true_divide(preds_y, cfg.MODEL.SIMDR_SPLIT_RATIO))                    output = output.cpu().numpy()                    preds = output.copy()                    for i in range(output.shape[0]):                        preds[i] = transform_preds(                            output[i], c[i], s[i], [cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1]]                        )                    preds[:, :, 2] = maxvals.squeeze(2)            else:                preds = np.empty((0, 0, 3), dtype=np.float32)        return preds    def hrnet_pts(self,img,detection,images,boxes):        if detection is not None:            for i, (x1, y1, x2, y2) in enumerate(detection):                x1 = int(round(x1.item()))                x2 = int(round(x2.item()))                y1 = int(round(y1.item()))                y2 = int(round(y2.item()))                # Adapt detections to match HRNet input aspect ratio (as suggested by xtyDoge in issue #14)                correction_factor = self.resolution[0] / self.resolution[1] * (x2 - x1) / (y2 - y1)                if correction_factor > 1:                    # increase y side                    center = y1 + (y2 - y1) // 2                    length = int(round((y2 - y1) * correction_factor))                    y1 = max(0, center - length // 2)                    y2 = min(img.shape[0], center + length // 2)                elif correction_factor  0:            images = images.to(self.device)            with torch.no_grad():                out = self.model(images)                out = out.detach().cpu().numpy()                pts = np.empty((out.shape[0], out.shape[1], 3), dtype=np.float32)                # For each human, for each joint: y, x, confidence                for i, human in enumerate(out):                    for j, joint in enumerate(human):                        pt = np.unravel_index(np.argmax(joint), (self.resolution[0] // 4, self.resolution[1] // 4))                        # 0: pt_x / (height // 4) * (bb_y2 - bb_y1) + bb_y1                        # 1: pt_y / (width // 4) * (bb_x2 - bb_x1) + bb_x1                        # 2: confidences                        pts[i, j, 0] = pt[1] * 1. / (self.resolution[1] // 4) * (boxes[i][2] - boxes[i][0]) + boxes[i][0]                        pts[i, j, 1] = pt[0] * 1. / (self.resolution[0] // 4) * (boxes[i][3] - boxes[i][1]) + boxes[i][1]                        pts[i, j, 2] = joint[pt]        else:            pts = np.empty((0, 0, 3), dtype=np.float32)        return pts    def _predict_single(self, image):        _,detections = self.detector.detect(image)        nof_people = len(detections) if detections is not None else 0        boxes = np.empty((nof_people, 4), dtype=np.int32)        images = torch.empty((nof_people, 3, self.resolution[0], self.resolution[1]))  # (height, width)        if self.model_name in ('sa-simdr','sasimdr'):            pts=self.sa_simdr_pts(image,detections,images,boxes)        elif self.model_name in ('hrnet','HRnet','hrnet'):            pts = self.hrnet_pts(image, detections, images, boxes)        elif self.model_name in ('simdr'):            pts = self.simdr_pts(image, detections, images, boxes)        return pts        # c,s=[],[]        # if detections is not None:        #     for i, (x1, y1, x2, y2) in enumerate(detections):        #         x1 = int(round(x1.item()))        #         x2 = int(round(x2.item()))        #         y1 = int(round(y1.item()))        #         y2 = int(round(y2.item()))        #         boxes[i] = [x1,y1,x2,y2]        #         w ,h= x2-x1,y2-y1        #         xx1 = np.max((0, x1))        #         yy1 = np.max((0, y1))        #         xx2 = np.min((image.shape[1] - 1, x1 + np.max((0, w - 1))))        #         yy2 = np.min((image.shape[0] - 1, y1 + np.max((0, h - 1))))        #         box = [xx1, yy1, xx2-xx1, yy2-yy1]        #         center,scale = self._box2cs(box)        #         c.append(center)        #         s.append(scale)        #        #         trans = get_affine_transform(center, scale, 0, np.array(cfg.MODEL.IMAGE_SIZE))        #         input = cv2.warpAffine(        #             image,        #             trans,        #             (int(self.resolution[1]), int(self.resolution[0])),        #             flags=cv2.INTER_LINEAR)        #         images[i] = self.transform(input)        # if images.shape[0] > 0:        #     images = images.to(self.device)        #     with torch.no_grad():        #         output_x,output_y = self.model(images)        #         if cfg.TEST.FLIP_TEST:        #             input_flipped = images.flip(3)        #             output_x_flipped_, output_y_flipped_ = self.model(input_flipped)        #             output_x_flipped = flip_back_simdr(output_x_flipped_.cpu().numpy(),        #                                                self.flip_pairs, type='x')        #             output_y_flipped = flip_back_simdr(output_y_flipped_.cpu().numpy(),        #                                                self.flip_pairs, type='y')        #             output_x_flipped = torch.from_numpy(output_x_flipped.copy()).to(self.device)        #             output_y_flipped = torch.from_numpy(output_y_flipped.copy()).to(self.device)        #        #             # feature is not aligned, shift flipped heatmap for higher accuracy        #             if cfg.TEST.SHIFT_HEATMAP:        #                 output_x_flipped[:, :, 0:-1] = \        #                     output_x_flipped.clone()[:, :, 1:]        #             output_x = F.softmax((output_x + output_x_flipped) * 0.5, dim=2)        #             output_y = F.softmax((output_y + output_y_flipped) * 0.5, dim=2)        #         else:        #             output_x = F.softmax(output_x, dim=2)        #             output_y = F.softmax(output_y, dim=2)        #         max_val_x, preds_x = output_x.max(2, keepdim=True)        #         max_val_y, preds_y = output_y.max(2, keepdim=True)        #        #         mask = max_val_x > max_val_y        #         max_val_x[mask] = max_val_y[mask]        #         maxvals = max_val_x.cpu().numpy()        #        #         output = torch.ones([images.size(0), preds_x.size(1), 2])        #         output[:, :, 0] = torch.squeeze(torch.true_divide(preds_x, cfg.MODEL.SIMDR_SPLIT_RATIO))        #         output[:, :, 1] = torch.squeeze(torch.true_divide(preds_y, cfg.MODEL.SIMDR_SPLIT_RATIO))        #        #         output = output.cpu().numpy()        #         preds = output.copy()        #         for i in range(output.shape[0]):        #             preds[i] = transform_preds(        #                 output[i], c[i], s[i], [cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1]]        #             )        # else:        #     preds = np.empty((0, 0, 2), dtype=np.float32)        # return preds# parser = argparse.ArgumentParser()# parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')# parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')# parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')# parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')# parser.add_argument('--augment', action='store_true', help='augmented inference')# parser.add_argument('--update', action='store_true', help='update all model')# parser.add_argument('--project', default='runs/detect', help='save results to project/name')# parser.add_argument('--name', default='exp', help='save results to project/name')# parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')# opt = parser.parse_args()# model = Points(model_name='hrnet',opt=opt)# img0 = cv2.imread('./data/test1.jpg')# pts = model.predict(img0)# print(pts.shape)# for point in pts[0]:#     image = cv2.circle(img0, (int(point[0]), int(point[1])), 3, [255,0,255], -1 , lineType= cv2.LINE_AA)#     cv2.imwrite('./data/test11_result.jpg',image)

④绘制骨骼关键点

根据以上步骤,我们已经得到了关键点的坐标值,接下来需要在图片中描绘出来,以便展示检测结果。

首先在’ ./SimDR/lib/utils/ ‘文件夹下新建visualization.py文件,将以下内容复制到文件中。骨架绘制代码结合了simple-hrnet与Openpose工程。

【2022.04.16更新:由于之前的绘制代码被我魔改过,现在恢复成所有点与骨骼都绘制的模样,但是总觉得好丑,没有openpose那种美观,如果有人绘制出比较美观的骨架,希望能分享一下哈,共同进步!】

import cv2import matplotlib.pyplot as pltimport numpy as npimport torchimport torchvisionimport ffmpegimport randomimport mathimport copydef plot_one_box(x, img, color=None, label=None, line_thickness=3):    # Plots one bounding box on image img    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness    color = color or [random.randint(0, 255) for _ in range(3)]    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)    if label:        tf = max(tl - 1, 1)  # font thickness        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)    return imgdef joints_dict():    joints = {        "coco": {            "keypoints": {                0: "nose",                1: "left_eye",                2: "right_eye",                3: "left_ear",                4: "right_ear",                5: "left_shoulder",                6: "right_shoulder",                7: "left_elbow",                8: "right_elbow",                9: "left_wrist",                10: "right_wrist",                11: "left_hip",                12: "right_hip",                13: "left_knee",                14: "right_knee",                15: "left_ankle",                16: "right_ankle"            },            "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]                # [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7],                # [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]                [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7],                [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],  # [3, 5], [4, 6]                [0, 5], [0, 6]                # [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7],                # [6, 8], [7, 9], [8, 10],  [0, 3], [0, 4], [1, 3], [2, 4],  # [3, 5], [4, 6]                # [0, 5], [0, 6]            ]        },        "mpii": {            "keypoints": {                0: "right_ankle",                1: "right_knee",                2: "right_hip",                3: "left_hip",                4: "left_knee",                5: "left_ankle",                6: "pelvis",                7: "thorax",                8: "upper_neck",                9: "head top",                10: "right_wrist",                11: "right_elbow",                12: "right_shoulder",                13: "left_shoulder",                14: "left_elbow",                15: "left_wrist"            },            "skeleton": [                # [5, 4], [4, 3], [0, 1], [1, 2], [3, 2], [13, 3], [12, 2], [13, 12], [13, 14],                # [12, 11], [14, 15], [11, 10], # [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]                [5, 4], [4, 3], [0, 1], [1, 2], [3, 2], [3, 6], [2, 6], [6, 7], [7, 8], [8, 9],                [13, 7], [12, 7], [13, 14], [12, 11], [14, 15], [11, 10],            ]        },    }    return jointsdef draw_points(image, points, color_palette='tab20', palette_samples=16, confidence_threshold=0.1,color=None):    """    Draws `points` on `image`.    Args:        image: image in opencv format        points: list of points to be drawn.            Shape: (nof_points, 3)            Format: each point should contain (y, x, confidence)        color_palette: name of a matplotlib color palette            Default: 'tab20'        palette_samples: number of different colors sampled from the `color_palette`            Default: 16        confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]            Default: 0.1    Returns:        A new image with overlaid points    """    circle_size = max(2, int(np.sqrt(np.max(np.max(points, axis=0) - np.min(points, axis=0)) // 16)))    for i, pt in enumerate(points):        if pt[2] >= confidence_threshold:            image = cv2.circle(image, (int(pt[0]), int(pt[1])), circle_size, color[i] ,-1, lineType= cv2.LINE_AA)    return imagedef draw_skeleton(image, points, skeleton, color_palette='Set2', palette_samples=8, person_index=0,                  confidence_threshold=0.1,sk_color=None):    """    Draws a `skeleton` on `image`.    Args:        image: image in opencv format        points: list of points to be drawn.            Shape: (nof_points, 3)            Format: each point should contain (y, x, confidence)        skeleton: list of joints to be drawn            Shape: (nof_joints, 2)            Format: each joint should contain (point_a, point_b) where `point_a` and `point_b` are an index in `points`        color_palette: name of a matplotlib color palette            Default: 'Set2'        palette_samples: number of different colors sampled from the `color_palette`            Default: 8        person_index: index of the person in `image`            Default: 0        confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]            Default: 0.1    Returns:        A new image with overlaid joints    """    canvas = copy.deepcopy(image)    cur_canvas = canvas.copy()    for i, joint in enumerate(skeleton):        pt1, pt2 = points[joint]        if pt1[2] >= confidence_threshold and pt2[2]>= confidence_threshold :            length = ((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2) ** 0.5            angle = math.degrees(math.atan2(pt1[1] - pt2[1],pt1[0] - pt2[0]))            polygon = cv2.ellipse2Poly((int(np.mean((pt1[0],pt2[0]))), int(np.mean((pt1[1],pt2[1])))), (int(length / 2), 2), int(angle), 0, 360, 1)            cv2.fillConvexPoly(cur_canvas, polygon, sk_color[i],lineType=cv2.LINE_AA)            # cv2.fillConvexPoly(cur_canvas, polygon, sk_color,lineType=cv2.LINE_AA)            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)    return canvasdef draw_points_and_skeleton(image, points, skeleton, points_color_palette='tab20', points_palette_samples=16,                             skeleton_color_palette='Set2', skeleton_palette_samples=8, person_index=0,                             confidence_threshold=0.1,color=None,sk_color=None):    """    Draws `points` and `skeleton` on `image`.    Args:        image: image in opencv format        points: list of points to be drawn.            Shape: (nof_points, 3)            Format: each point should contain (y, x, confidence)        skeleton: list of joints to be drawn            Shape: (nof_joints, 2)            Format: each joint should contain (point_a, point_b) where `point_a` and `point_b` are an index in `points`        points_color_palette: name of a matplotlib color palette            Default: 'tab20'        points_palette_samples: number of different colors sampled from the `color_palette`            Default: 16        skeleton_color_palette: name of a matplotlib color palette            Default: 'Set2'        skeleton_palette_samples: number of different colors sampled from the `color_palette`            Default: 8        person_index: index of the person in `image`            Default: 0        confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]            Default: 0.1    Returns:        A new image with overlaid joints    """    colors1 = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],               [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],               [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [255, 0, 85]]    image = draw_skeleton(image, points, skeleton, color_palette=skeleton_color_palette,                          palette_samples=skeleton_palette_samples, person_index=person_index,                          confidence_threshold=confidence_threshold,sk_color=colors1)    image = draw_points(image, points, color_palette=points_color_palette, palette_samples=points_palette_samples,                        confidence_threshold=confidence_threshold,color=colors1)    return imagedef save_images(images, target, joint_target, output, joint_output, joint_visibility, summary_writer=None, step=0,                prefix=''):    """    Creates a grid of images with gt joints and a grid with predicted joints.    This is a basic function for debugging purposes only.    If summary_writer is not None, the grid will be written in that SummaryWriter with name "{prefix}_images" and    "{prefix}_predictions".    Args:        images (torch.Tensor): a tensor of images with shape (batch x channels x height x width).        target (torch.Tensor): a tensor of gt heatmaps with shape (batch x channels x height x width).        joint_target (torch.Tensor): a tensor of gt joints with shape (batch x joints x 2).        output (torch.Tensor): a tensor of predicted heatmaps with shape (batch x channels x height x width).        joint_output (torch.Tensor): a tensor of predicted joints with shape (batch x joints x 2).        joint_visibility (torch.Tensor): a tensor of joint visibility with shape (batch x joints).        summary_writer (tb.SummaryWriter): a SummaryWriter where write the grids.            Default: None        step (int): summary_writer step.            Default: 0        prefix (str): summary_writer name prefix.            Default: ""    Returns:        A pair of images which are built from torchvision.utils.make_grid    """    # Input images with gt    images_ok = images.detach().clone()    images_ok[:, 0].mul_(0.229).add_(0.485)    images_ok[:, 1].mul_(0.224).add_(0.456)    images_ok[:, 2].mul_(0.225).add_(0.406)    for i in range(images.shape[0]):        joints = joint_target[i] * 4.        joints_vis = joint_visibility[i]        for joint, joint_vis in zip(joints, joints_vis):            if joint_vis[0]:                a = int(joint[1].item())                b = int(joint[0].item())                # images_ok[i][:, a-1:a+1, b-1:b+1] = torch.tensor([1, 0, 0])                images_ok[i][0, a - 1:a + 1, b - 1:b + 1] = 1                images_ok[i][1:, a - 1:a + 1, b - 1:b + 1] = 0    grid_gt = torchvision.utils.make_grid(images_ok, nrow=int(images_ok.shape[0] ** 0.5), padding=2, normalize=False)    if summary_writer is not None:        summary_writer.add_image(prefix + 'images', grid_gt, global_step=step)    # Input images with prediction    images_ok = images.detach().clone()    images_ok[:, 0].mul_(0.229).add_(0.485)    images_ok[:, 1].mul_(0.224).add_(0.456)    images_ok[:, 2].mul_(0.225).add_(0.406)    for i in range(images.shape[0]):        joints = joint_output[i] * 4.        joints_vis = joint_visibility[i]        for joint, joint_vis in zip(joints, joints_vis):            if joint_vis[0]:                a = int(joint[1].item())                b = int(joint[0].item())                # images_ok[i][:, a-1:a+1, b-1:b+1] = torch.tensor([1, 0, 0])                images_ok[i][0, a - 1:a + 1, b - 1:b + 1] = 1                images_ok[i][1:, a - 1:a + 1, b - 1:b + 1] = 0    grid_pred = torchvision.utils.make_grid(images_ok, nrow=int(images_ok.shape[0] ** 0.5), padding=2, normalize=False)    if summary_writer is not None:        summary_writer.add_image(prefix + 'predictions', grid_pred, global_step=step)    # Heatmaps    # ToDo    # for h in range(0,17):    #     heatmap = torchvision.utils.make_grid(output[h].detach(), nrow=int(np.sqrt(output.shape[0])),    #                                            padding=2, normalize=True, range=(0, 1))    #     summary_writer.add_image('train_heatmap_%d' % h, heatmap, global_step=step + epoch*len_dl_train)    return grid_gt, grid_preddef check_video_rotation(filename):    # thanks to    # https://stackoverflow.com/questions/53097092/frame-from-video-is-upside-down-after-extracting/55747773#55747773    # this returns meta-data of the video file in form of a dictionary    meta_dict = ffmpeg.probe(filename)    # from the dictionary, meta_dict['streams'][0]['tags']['rotate'] is the key    # we are looking for    rotation_code = None    try:        if int(meta_dict['streams'][0]['tags']['rotate']) == 90:            rotation_code = cv2.ROTATE_90_CLOCKWISE        elif int(meta_dict['streams'][0]['tags']['rotate']) == 180:            rotation_code = cv2.ROTATE_180        elif int(meta_dict['streams'][0]['tags']['rotate']) == 270:            rotation_code = cv2.ROTATE_90_COUNTERCLOCKWISE        else:            raise ValueError    except KeyError:        pass    return rotation_code

4、测试算法

①主程序

在SimDR文件夹下新建main.py ,复制以下代码到文件中,修改parser参数source的默认值,运行代码。

import argparseimport timeimport osimport cv2 as cvimport numpy as npfrom pathlib import Pathfrom Point_detect import Pointsfrom lib.utils.visualization import draw_points_and_skeleton,joints_dictdef image_detect(opt):    skeleton = joints_dict()['coco']['skeleton']    hrnet_model = Points(model_name='hrnet', opt=opt,resolution=(384,288))  #resolution = (384,288)  or (256,192)    # simdr_model = Points(model_name='simdr', opt=opt,resolution=(256,192))  #resolution = (256,192)    # sa_simdr_model = Points(model_name='sa-simdr', opt=opt,resolution=(384,288))  #resolution = (384,288)  or (256,192)    img0 = cv.imread(opt.source)    frame = img0.copy()   #predict    pred = hrnet_model.predict(img0)    # pred = simdr_model.predict(frame)    # pred = sa_simdr_model.predict(frame)   #vis    for i, pt in enumerate(pred):        frame = draw_points_and_skeleton(frame, pt, skeleton)    #save    cv.imwrite('test_result.jpg', frame)def video_detect(opt):    hrnet_model = Points(model_name='hrnet', opt=opt, resolution=(384, 288))  # resolution = (384,288)  or (256,192)    # simdr_model = Points(model_name='simdr', opt=opt,resolution=(256,192))  #resolution = (256,192)    # sa_simdr_model = Points(model_name='sa-simdr', opt=opt,resolution=(384,288))  #resolution = (384,288)  or (256,192)    skeleton = joints_dict()['coco']['skeleton']    cap = cv.VideoCapture(opt.source)    if opt.save_video:        fourcc = cv.VideoWriter_fourcc(*'MJPG')        out = cv.VideoWriter('data/runs/{}_out.avi'.format(os.path.basename(opt.source).split('.')[0]), fourcc, 24, (int(cap.get(3)), int(cap.get(4))))    while cap.isOpened():        ret, frame = cap.read()        if not ret:            break        pred = hrnet_model.predict(frame)        # pred = simdr_model.predict(frame)        # pred = sa_simdr_model.predict(frame)        for pt in pred:            frame = draw_points_and_skeleton(frame,pt,skeleton)        if opt.show:            cv.imshow('result', frame)        if opt.save_video:            out.write(frame)        if cv.waitKey(1) == 27:            break    out.release()    cap.release()    cv.destroyAllWindows()# video_detect(0)if __name__ == '__main__':    parser = argparse.ArgumentParser()    parser.add_argument('--source', type=str, default='./data/images/test1.jpg', help='source')  # file/folder, 0 for webcam    parser.add_argument('--detect_weight', type=str, default="./yolov5/weights/yolov5x.pt", help='e.g "./yolov5/weights/yolov5x.pt"')    parser.add_argument('--save_video', action='store_true', default=False,help='save results to *.avi')    parser.add_argument('--show', action='store_true', default=True, help='save results to *.avi')    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')    parser.add_argument('--augment', action='store_true', help='augmented inference')    opt = parser.parse_args()    image_detect(opt)

②结果展示

四、总结

全文较长,主要都是些代码,整个工程从跑数据集到实际检测需要对代码工程有一定的理解,整个项目不难,主要考验类的构造。如果需要整个工程可以私聊我。由于我也是刚入门的萌新,所以代码格式写法或者理论看法有很多错误,欢迎指正,共同进步,如果有帮助欢迎点赞评论,万分感谢。

五、参考内容

1、GitHub – leeyegy/SimDR: PyTorch implementation for: Is 2D Heatmap Representation Even Necessary for Human Pose Estimation? (http://arxiv.org/abs/2107.03332)

2、https://github.com/ultralytics/yolov5

3、GitHub – GreenTeaHua/simple-HRNet: Multi-person Human Pose Estimation with HRNet in Pytorch