关键点检测——HRNet源码解析篇

🍊作者简介:秃头小苏,致力于用最通俗的语言描述问题

🍊专栏推荐:深度学习网络原理与实战

🍊近期目标:写好专栏的每一篇文章

🍊支持小苏:点赞👍🏼、收藏⭐、留言📩

HRNet源码解析篇

写在前面

Hello,大家好,我是小苏👦🏽👦🏽👦🏽

在上一节中,我已经为大家介绍HRNet的原理部分,其实说起来挺惭愧,因为原理部分介绍的还是比较简单的,我想你仅仅阅读原理部分是很难彻底弄明白HRNet的精髓。那么本节将在上一节的基础上,为大家更细致的讲解HRNet。🧨🧨🧨当然了,本节属于源码解析篇,所有会存在比较多的代码,大家也不用担心看不懂,我都对关键代码做了详细的解释,并画图进一步帮助大家理解,所以大家一定要耐心看下去喔。🎯🎯🎯

这里我还想多说一句,其实写源码解析类博客其实怪难的,因为有时候明明很好表达的内容,用文字为大家展示却有种一拳打在棉花上的无力感,哈哈哈,可能是自己水平不够叭。🧑🧑🧑自己也做过几期视频,感觉效果也还行,感兴趣的可以点击☞☞☞看看,后期可能会考虑写完博客出配套视频的方式为大家介绍知识点。enmmm,说远了,说这些我是想告诉大家,我并不期望通过这一篇文章就能让你把整个HRNet的源码都看透,这是不可能的。但是其中一些关键的地方,如果本文能给你一点启发,那我觉得此篇文章的目的就达到了。此外,大家在阅读代码时,一定不要停留在看的层面,一定要动手调试起来,这样会有不一样的收获。🍋🍋🍋

好了,不说废话了,让我们一起发车,来学学HRNet的源码叭~~~🚖🚖🚖

源码地址:HRNet源码🌱🌱🌱

关键点数据集构建

深度学习中数据才是王道,本文使用的是COCO数据集中的人体关键点检测数据集,对此数据集还不清楚的务必点击下面链接了解详情:

COCO数据集------关键点检测标注文件解析🌼🌼🌼

清楚COCO数据集的格式后,我们一起来看看是如何构建关键点数据集的?首先来说说这里的关键点数据集的构建主要干了什么?其实它就是把原始图像中对人体关键点标注过的图像记录了下来。

我们一点点的来看其是如何实现的,主要定义在CocoKeypoint类中:

python 复制代码
data_root = args.data_path   # data_path:'D://Dataset//coco2017'
data_transform = {
        "train": transforms.Compose([
            transforms.HalfBody(0.3, person_kps_info["upper_body_ids"], person_kps_info["lower_body_ids"]),
            transforms.AffineTransform(scale=(0.65, 1.35), rotation=(-45, 45), fixed_size=fixed_size),
            transforms.RandomHorizontalFlip(0.5, person_kps_info["flip_pairs"]),
            transforms.KeypointToHeatMap(heatmap_hw=heatmap_hw, gaussian_sigma=2, keypoints_weights=kps_weights),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]),
        "val": transforms.Compose([
            transforms.AffineTransform(scale=(1.25, 1.25), fixed_size=fixed_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    }

fixed_size = args.fixed_size  # fixed_size:[256,192]

train_dataset = CocoKeypoint(data_root, "train", transforms=data_transform["train"], fixed_size=args.fixed_size)

上述代码主要展示了传入CocoKeypoint类中的参数,为后面阅读CocoKeypoint类中代码做准备,对本篇transforms不熟悉的可以稍后阅读本文的下一节-->在线数据增强。🥗🥗🥗

下面我们就一步步的来看看CocoKeypoint类到底干了什么?在__init__函数中,先初始化了一系列变量:

python 复制代码
anno_file = f"person_keypoints_{dataset}{years}.json"
self.img_root = os.path.join(root, f"{dataset}{years}")
self.anno_path = os.path.join(root, "annotations", anno_file)
self.fixed_size = fixed_size
self.mode = dataset
self.transforms = transforms
self.coco = COCO(self.anno_path)

我们调试来看看这些值的结果:

这里我重点介绍一下self.coco = COCO(self.anno_path)这句代码,其传入的是self.anno_path参数,即人体关键点检测标注文件------'D://Dataset//coco2017\\annotations\\person_keypoints_train2017.json'我们跳入COCO函数内部调试一下:

首先设置一些字典变量来存储相关信息:

python 复制代码
self.dataset,self.anns,self.cats,self.imgs = dict(),dict(),dict(),dict()
self.imgToAnns, self.catToImgs = defaultdict(list), defaultdict(list)

接着我们会打开标注文件路径并读取得到dataset:

python 复制代码
with open(annotation_file, 'r') as f:
    dataset = json.load(f)

我们来看看dataset的值:

dataset一共有五个字段的值,和我在COCO数据集------关键点检测标注文件解析这篇博客中介绍的是完全一致的。

接着调用createIndex方法为之前定义的字典变量赋值:

python 复制代码
if 'annotations' in self.dataset:
    for ann in self.dataset['annotations']:
        imgToAnns[ann['image_id']].append(ann)
        anns[ann['id']] = ann

这段代码先是遍历数据集中的annotations标签,然后将其image_id作为键,标签作为值构建一个图像id到标签的字典imgToAnns,来看看其遍历一次的结果:

接着是将标签的id作为键,标签作为值构建一个字典,同样看看遍历一次的结果:

注意:这里的id和image_id不一样,image_id是图像的唯一标识,id是目标实例分配的唯一标识符,用于在数据集中唯一标识这个目标实例。一个图像中可能会有多个目标实列,即有多个人。

我们遍历完所有标签,看看imgToAnns和anns中有多少数据:

可以看到imgToAnns共有64115条数据,表示一共有64115张图像存在标注。anns共有262465条数据,表示一共有262465个标注目标实列,也就是标注了262465个人。那么为什么会存在这样的差异呢?因为不是每张图像都会有标注目标实列(一共118287张图像,有标注目标实列的有64115张),也不是每张存在标注目标实列的都只有一个标注目标实列(最少有一个,从imgToAnns数据图中可以看到,image_id为120021的图像有三个标注目标实列,说明标注了3个人,我们也可以来看看这张图像,看看是不是有3个人,如下:

)


python 复制代码
if 'images' in self.dataset:
    for img in self.dataset['images']:
        imgs[img['id']] = img

接着这段代码是遍历数据集中的image图像,然后将id【注意:image中的id指的是image_id,而不是上文说的实列id】作为键,img图像作为值构建字典imgs,来看看遍历一次的结果:

然后来看看遍历完所有数据imgs的结果:

一共有118287条数据,这就是COCO训练集图片的数量。


python 复制代码
if 'categories' in self.dataset:
    for cat in self.dataset['categories']:
        cats[cat['id']] = cat

这段代码是遍历数据集中的categories类别,然后将其id【注意:这里的id指类别的唯一标识符。在人体关键点检测中,这个id都是1,因为我们只会对人体进行标注,而person的类别标识符是1。】作为键,categories信息作为值构建cats字典,来看看遍历一次的结果:

enmmm,dataset中只有一条数据,只能遍历一次,因为只有一个类别id,即id=1。


python 复制代码
if 'annotations' in self.dataset and 'categories' in self.dataset:
    for ann in self.dataset['annotations']:
        catToImgs[ann['category_id']].append(ann['image_id'])

这段代码同样遍历annotations标签,然后将 category_id【注意:这个是类别id,其为1】作为键,image_id作为值构建从类别id到图像id映射的字典 catToImgs,来看一次遍历的结果:

然后来看看遍历完所有数据的catToImgs字典:

一共有262465条数据,即表示有262465张图像有类别id 1(这里包括了重复的图像,比如一个图像中有3个人,那么这里就有三条数据,那么其实这里的262465表示一共标注了262465个person实列)


到这里我们的COCO(self.anno_path)函数的内容就介绍完啦,我们来看看self.coco的值,如下:

其实其就是COCO(self.anno_path)函数中那几个字典变量。

接着我们会把imgs的key进行排序,并转成列表:

python 复制代码
img_ids = list(sorted(self.coco.imgs.keys()))

然后通过det = self.coco将self.coco的值赋给det,并设置一个self.valid_person_list列表用于存储有效的人体关键点信息,并设置一个obj_idx记录目标实列个数。

python 复制代码
self.valid_person_list = []
obj_idx = 0

最后执行下面的代码:

python 复制代码
for img_id in img_ids:
    img_info = self.coco.loadImgs(img_id)[0]
    ann_ids = det.getAnnIds(imgIds=img_id)
    anns = det.loadAnns(ann_ids)
    for ann in anns:
        # only save person class
        if ann["category_id"] != 1:
            print(ann["category_id"])

            # skip objs without keypoints annotation
            if "keypoints" in ann:
                if max(ann["keypoints"]) == 0:
                    continue

                    xmin, ymin, w, h = ann['bbox']
                    # Use only valid bounding boxes
                    if w > 0 and h > 0:
                        info = {
                            "box": [xmin, ymin, w, h],
                            "image_path": os.path.join(self.img_root, img_info["file_name"]),
                            "image_id": img_id,
                            "image_width": img_info['width'],
                            "image_height": img_info['height'],
                            "obj_origin_hw": [h, w],
                            "obj_index": obj_idx,
                            "score": ann["score"] if "score" in ann else 1.
                        }
                        if "keypoints" in ann:
                            keypoints = np.array(ann["keypoints"]).reshape([-1, 3])
                            visible = keypoints[:, 2]
                            keypoints = keypoints[:, :2]
                            info["keypoints"] = keypoints
                            info["visible"] = visible

                            self.valid_person_list.append(info)
                            obj_idx += 1

这段代码干了什么呢,我们一点点来分析:首先是遍历img_ids,第一次拿到第一个image_id=9:

此image_id对应图像如下:

然后执行img_info = self.coco.loadImgs(img_id)[0],loadImgs定义如下:

python 复制代码
def loadImgs(self, ids=[]):
    if _isArrayLike(ids):
        return [self.imgs[id] for id in ids]
    elif type(ids) == int:
        return [self.imgs[ids]]

这个函数主要是根据img_id来加载图像,我们直接来看img_info的结果:

这显示了image_id=9的图像的信息。


接着是ann_ids = det.getAnnIds(imgIds=img_id)这句代码,getAnnIds函数如下:

python 复制代码
def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None):
    imgIds = imgIds if _isArrayLike(imgIds) else [imgIds]
    catIds = catIds if _isArrayLike(catIds) else [catIds]

    if len(imgIds) == len(catIds) == len(areaRng) == 0:
        anns = self.dataset['annotations']
    else:
        if not len(imgIds) == 0:
            lists = [self.imgToAnns[imgId] for imgId in imgIds if imgId in self.imgToAnns]
            anns = list(itertools.chain.from_iterable(lists))
        else:
            anns = self.dataset['annotations']
            anns = anns if len(catIds)  == 0 else [ann for ann in anns if ann['category_id'] in catIds]
            anns = anns if len(areaRng) == 0 else [ann for ann in anns if ann['area'] > areaRng[0] and ann['area'] < areaRng[1]]
            if not iscrowd == None:
                ids = [ann['id'] for ann in anns if ann['iscrowd'] == iscrowd]
            else:
                ids = [ann['id'] for ann in anns]
    return ids

我们注意来看一下这句:lists = [self.imgToAnns[imgId] for imgId in imgIds if imgId in self.imgToAnns],意思是遍历imgIds,然后判断imgId是否存在self.imgToAnns中,我们知道imgToAnns共有64115条数据,表示一共有64115张图像存在标注,也就是说只有存在人物的图像才存在标注,即只有图像存在人物,其imgId才会在imgToAnns中,而我们刚刚image_id=9的图像不存在人物,故lists=[]是空列表,后续返回值会是空列表。因此我们需要换一张存在人物的图像进行展示,当遍历至image_id=36时,图像出现人物,图像如下:

此时lists值如下,为这张图像的标注信息:

接着anns = list(itertools.chain.from_iterable(lists))是将将嵌套的列表(lists)展平成一个单层的列表。

后面的这几句都没起作用:

python 复制代码
anns = anns if len(catIds)  == 0 else [ann for ann in anns if ann['category_id'] in catIds]
anns = anns if len(areaRng) == 0 else [ann for ann in anns if ann['area'] > areaRng[0] and ann['area'] < areaRng[1]]
if not iscrowd == None:
    ids = [ann['id'] for ann in anns if ann['iscrowd'] == iscrowd]

然后通过ids = [ann['id'] for ann in anns]获取到标注目标实列的id,并将其返回给ann_ids,其值如下:

接着是anns = det.loadAnns(ann_ids),将ann_ids传入loadAnns方法中,其定义如下:

python 复制代码
def loadAnns(self, ids=[]):
    if _isArrayLike(ids):
        return [self.anns[id] for id in ids]
    elif type(ids) == int:
        return [self.anns[ids]]

这个函数主要通过ann_ids来加载标注信息,返回的anns如下:

接着是遍历anns,先检查ann["category_id"] 是否为1并检查标注的keypoints关键点是否存在可见关键点。

python 复制代码
for ann in anns:
    # only save person class
    if ann["category_id"] != 1:
        print(ann["category_id"])

    # skip objs without keypoints annotation
    if "keypoints" in ann:
        if max(ann["keypoints"]) == 0:
            continue

然后从标注的bbox中获取xmin, ymin, w, h,并构建info信息,注意这里score不在ann中,其最后值为1。

python 复制代码
xmin, ymin, w, h = ann['bbox']
# Use only valid bounding boxes
if w > 0 and h > 0:
    info = {
        "box": [xmin, ymin, w, h],
        "image_path": os.path.join(self.img_root, img_info["file_name"]),
        "image_id": img_id,
        "image_width": img_info['width'],
        "image_height": img_info['height'],
        "obj_origin_hw": [h, w],
        "obj_index": obj_idx,
        "score": ann["score"] if "score" in ann else 1.
    }

然后将关键点的坐标和可见性分成两个变量表示并加入到info字典中,最后将info添加到self.valid_person_list中,并将obj_idx加1,表示多了一个目标实列 。

python 复制代码
if "keypoints" in ann:
    keypoints = np.array(ann["keypoints"]).reshape([-1, 3])
    visible = keypoints[:, 2]
    keypoints = keypoints[:, :2]
    info["keypoints"] = keypoints
    info["visible"] = visible

self.valid_person_list.append(info)
obj_idx += 1

此循环代码结束,来看看self.valid_person_list.append的值:

【注意:这里的图像只有一个人物,如果图像中包含多个人物原理是一样的,会遍历图像中的各个人物,并把每个人物的信息存放到valid_person_list中】

当我们遍历完所有img_ids数据时,self.valid_person_list.append就存储了所有有效的人体关键点检测的相关信息,最后一共有149813个有效数据。


到这里,关键点检测数据集的构建部分就为大家介绍完了,这部分说难也算不上难,但我认为却是非常重要的一部分,希望大家好好消化一下。🍻🍻🍻

在线数据增强

enmmm,我想大家应该对数据增强有一定的了解叭,比如旋转、剪裁、水平翻转等等,在代码中通常使用transforms.xxx来实现,在之前的博客中,我为大家介绍过一些数据增强,如Faster RCNN中为大家介绍了水平翻转。

这一小节我准备多花点时间来写,因为这部分有的地方是比较难理解,也是蛮重要的。话不多说,我们一起来看看HRNet中使用了哪些数据增强手段:

python 复制代码
data_transform = {
        "train": transforms.Compose([
            transforms.HalfBody(0.3, person_kps_info["upper_body_ids"], person_kps_info["lower_body_ids"]),
            transforms.AffineTransform(scale=(0.65, 1.35), rotation=(-45, 45), fixed_size=fixed_size),
            transforms.RandomHorizontalFlip(0.5, person_kps_info["flip_pairs"]),
            transforms.KeypointToHeatMap(heatmap_hw=heatmap_hw, gaussian_sigma=2, keypoints_weights=kps_weights),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]),
        "val": transforms.Compose([
            transforms.AffineTransform(scale=(1.25, 1.25), fixed_size=fixed_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    }

其训练集和测试集采用不同的数据增强手段,测试集使用的数据增强方法是训练集的子集,所以直接来看训练集中的方法就好了,一共有以下四个:

  • transforms.HalfBody
  • transforms.AffineTransform
  • transforms.RandomHorizontalFlip
  • transforms.KeypointToHeatMap

注意: transforms.ToTensor和transforms.Normalize不属于数据增强,而是处于数据预处理,对这两个不熟悉的可以点击下面的链接了解详情:

pytorch中的transforms.ToTensor和transforms.Normalize理解🍁🍁🍁


下面我将来一个个的为大家介绍这四种数据增强方式,快来和我一起学学叭~~~🥂🥂🥂

在具体介绍每种数据增强方法之前,我先来给大家展示一下本次调试使用的图片,如下:

其shape为(426,640,3),为COCO验证集的第一张图片。

transforms.HalfBody

HalfBody------一半的身体,大家可以猜猜这个数据增强手段干了什么?好叭,不卖关子了,一句话解释它干了什么,就是以一定的概率让人体关键点保留上半部分或者下半部分。你或许会疑惑为什么要把完整的人切分成上半部分和下半部分,这是为了模拟关键点检测中的部分遮挡情况。在实际场景中,人物可能被其他对象或者场景的遮挡,这样的情况会使得关键点检测更加具有挑战性。而我们使用HalfBody数据增强,只使用部分身体进行训练,可以实现类似遮挡的效果。这可以帮助模型学习如何处理部分遮挡的情况,提高模型在真实场景中的鲁棒性。

知道了HalfBody的原理,下面就来看看这个HalfBody类是如何实现的:

首先来看看__init__方法:

python 复制代码
def __init__(self, p: float = 0.3, upper_body_ids=None, lower_body_ids=None):
    assert upper_body_ids is not None
    assert lower_body_ids is not None
    self.p = p
    self.upper_body_ids = upper_body_ids
    self.lower_body_ids = lower_body_ids

这个初始化方法主要定义了一个概率p,upper_body_ids和lower_body_ids。upper_body_ids和lower_body_ids是指人的身体的上部和下部的索引,0,1,2,3,4,5,6,7,8,9,10为上,其它为下:

为了让大家直观的感受,作图如下:

注意:上图索引是从1开始的,代码索引是从0开始的


接着来看__call__方法:

python 复制代码
def __call__(self, image, target):
    if random.random() < self.p:
        kps = target["keypoints"]
        vis = target["visible"]
        upper_kps = []
        lower_kps = []

        # 对可见的keypoints进行归类
        for i, v in enumerate(vis):
            if v > 0.5:
                if i in self.upper_body_ids:
                    upper_kps.append(kps[i])
                else:
                    lower_kps.append(kps[i])

        # 50%的概率选择上或下半身
        if random.random() < 0.5:
            selected_kps = upper_kps
        else:
            selected_kps = lower_kps

        # 如果点数太少就不做任何处理
        if len(selected_kps) > 2:
            selected_kps = np.array(selected_kps, dtype=np.float32)
            xmin, ymin = np.min(selected_kps, axis=0).tolist()  # 寻找x,y坐标的最小值
            xmax, ymax = np.max(selected_kps, axis=0).tolist()  # 寻找x,y坐标的最大值
            w = xmax - xmin
            h = ymax - ymin
            if w > 1 and h > 1:
                # 把w和h适当放大点,要不然关键点处于边缘位置
                xmin, ymin, w, h = scale_box(xmin, ymin, w, h, (1.5, 1.5))
                target["box"] = [xmin, ymin, w, h]

    return image, target

对上述代码做相关解释,首先以一定概率p(0.3)对图像进行HalfBody操作,若满足条件,获取关键点坐标和可见性,如下:

接着按照上半身和下半身对17个关键点进行分类,结果如下:【注意:这里只有15个关键点,因为vis表格有两个0值,表示有两个点没有标注,被if v > 0.5过滤掉了】

上半身一个9个关键点,下半身一共6个关键点,这就是根据upper_body_ids和lower_body_ids来划分的。【注:大家这里要是不理解一定要自己调试看看】

然后会以0.5的概率选择上半身的关键点或者下半身的关键点。如果发现选择的一半身体的关键点个数小于等于2个,则不做任何处理,返回原有的image和target。若关键点个数大于2,则执行以下代码:

python 复制代码
if len(selected_kps) > 2:
    selected_kps = np.array(selected_kps, dtype=np.float32)
    xmin, ymin = np.min(selected_kps, axis=0).tolist()  # 寻找x,y坐标的最小值
    xmax, ymax = np.max(selected_kps, axis=0).tolist()  # 寻找x,y坐标的最大值
    w = xmax - xmin
    h = ymax - ymin
    if w > 1 and h > 1:
    # 把w和h适当放大点,要不然关键点处于边缘位置
    xmin, ymin, w, h = scale_box(xmin, ymin, w, h, (1.5, 1.5))
    target["box"] = [xmin, ymin, w, h]

我先来介绍一下这段代码主要干了什么,其实就是找到新的目标(上半身或下半身)的bbox,我画图为大家解释一下:

这里我们其实是可以得到一个bbox了,但是其太靠近物体边缘了,放大1.5倍,代码如下:

python 复制代码
def scale_box(xmin: float, ymin: float, w: float, h: float, scale_ratio: Tuple[float, float]):
    """根据传入的h、w缩放因子scale_ratio,重新计算xmin,ymin,w,h"""
    s_h = h * scale_ratio[0]
    s_w = w * scale_ratio[1]
    xmin = xmin - (s_w - w) / 2.
    ymin = ymin - (s_h - h) / 2.
    return xmin, ymin, s_w, s_h

同样画个图帮大家理解,如下:

transforms.AffineTransform

AffineTransform------仿射变化,这个是干什么的,我来帮大家解释一下这个仿射变换干了什么,其实就是需要原始图像和目标图像中三个对应点**(代码中使用的是图像中心点,上边界中心点和右边界中心点)**,然后通过这三个点将原始图像变换倒目标图像。

enmmm,大家是不是没怎么明白,别急,我会带大家看看代码,并可视化输出结果,这样大家就能直观的感受到仿射变换到底做了什么了。

首先第一步会调整上一步骤得到的bbox的长宽比,使其符合h:w=256:192,这个是我们输入图片的尺寸,具体代码如下:

python 复制代码
src_xmin, src_ymin, src_xmax, src_ymax = adjust_box(*target["box"], self.fixed_size)

def adjust_box(xmin: float, ymin: float, w: float, h: float, fixed_size: Tuple[float, float]):
    """通过增加w或者h的方式保证输入图片的长宽比固定"""
    xmax = xmin + w
    ymax = ymin + h

    hw_ratio = fixed_size[0] / fixed_size[1]
    if h / w > hw_ratio:
        # 需要在w方向padding
        wi = h / hw_ratio
        pad_w = (wi - w) / 2
        xmin = xmin - pad_w
        xmax = xmax + pad_w
    else:
        # 需要在h方向padding
        hi = w * hw_ratio
        pad_h = (hi - h) / 2
        ymin = ymin - pad_h
        ymax = ymax + pad_h

    return xmin, ymin, xmax, ymax

我们可以看看调整bbox后的图像,如下:

你可能看不出bbox的长宽比,但其就是256:192,其高度h为71.507,宽度w为53.630,精度上有点差别,不相信的大家自己去调试试试喔~~~🍡🍡🍡

接着我们就寻找原图像(bbox)和目标图像的三个点:

python 复制代码
src_w = src_xmax - src_xmin
src_h = src_ymax - src_ymin

src_center = np.array([(src_xmin + src_xmax) / 2, (src_ymin + src_ymax) / 2])
src_p2 = src_center + np.array([0, -src_h / 2])  # top middle
src_p3 = src_center + np.array([src_w / 2, 0])   # right middle

dst_center = np.array([(self.fixed_size[1] - 1) / 2, (self.fixed_size[0] - 1) / 2])
dst_p2 = np.array([(self.fixed_size[1] - 1) / 2, 0])  # top middle
dst_p3 = np.array([self.fixed_size[1] - 1, (self.fixed_size[0] - 1) / 2])

然后对bbox进行缩放和旋转,先是缩放,

python 复制代码
if self.scale is not None:
    scale = random.uniform(*self.scale)
    src_w = src_w * scale
    src_h = src_h * scale
    src_p2 = src_center + np.array([0, -src_h / 2])  # top middle
    src_p3 = src_center + np.array([src_w / 2, 0])   # right middle

我们来看看缩放后的bbox: 【注意这里我调试时的scale取0.7115,所以bbox变小了】

然后是旋转:【注意这里我调试时的angle取-25,所以bbox逆时针旋转了25°】

python 复制代码
if self.rotation is not None:
    angle = random.randint(*self.rotation)  # 角度制
    angle = angle / 180 * math.pi  # 弧度制
    src_p2 = src_center + np.array([src_h / 2 * math.sin(angle), -src_h / 2 * math.cos(angle)])
    src_p3 = src_center + np.array([src_w / 2 * math.cos(angle), src_w / 2 * math.sin(angle)])

我们再来看看旋转后的结果:

最后就是仿射变换了:

python 复制代码
src = np.stack([src_center, src_p2, src_p3]).astype(np.float32)
dst = np.stack([dst_center, dst_p2, dst_p3]).astype(np.float32)

trans = cv2.getAffineTransform(src, dst)  # 计算正向仿射变换矩阵
dst /= 4  # 网络预测的heatmap尺寸是输入图像的1/4
reverse_trans = cv2.getAffineTransform(dst, src)  # 计算逆向仿射变换矩阵,方便后续还原


# 对图像进行仿射变换
resize_img = cv2.warpAffine(img,
                            trans,
                            tuple(self.fixed_size[::-1]),  # [w, h]
                            flags=cv2.INTER_LINEAR)

同样的我们来看看最后的resize_img长什么样,如下:【resize_img的大小是256*192喔】

到这里我们对图像的操作就完成了,不要忘记我们还要对标签进行同样的操作喔,关键点检测的标签就是一个个点嘛,如下:

python 复制代码
if "keypoints" in target:
    kps = target["keypoints"]
    mask = np.logical_and(kps[:, 0] != 0, kps[:, 1] != 0)
    kps[mask] = affine_points(kps[mask], trans)
    target["keypoints"] = kps
    
def affine_points(pt, t):
    ones = np.ones((pt.shape[0], 1), dtype=float)
    pt = np.concatenate([pt, ones], axis=1).T
    new_pt = np.dot(t, pt)
    return new_pt.T
           

最后再来用一个图来总结一下仿射变换都做了什么,如下:

transforms.RandomHorizontalFlip

RandomHorizontalFlip------水平翻转。我想对这个数据增强手段大家都比较熟悉,就是将图片左右进行翻转,其最后实现的效果如下图所示:

和仿射变换一样,要实现水平翻转,我们不仅需要对图片进行水平翻转操作,同样需要对标签进行同步操作,我们分别来看看如何对图片和标签进行水平翻转操作的叭。

  • 图片

    对图片进行水平翻转的操作很简单啦,只需要一行代码就可以了喔,如下:

    python 复制代码
    image = np.ascontiguousarray(np.flip(image, axis=[1]))    # 水平左右翻转
  • 标签

    对标签进行翻转是这步的难点,我当时阅读这部分的代码时弄了很长时间才明白,其实要画一个图大家就能很容易的理解。我们先来看代码叭:【注:我没有复制所有代码过来了,挑了关键的代码】

    python 复制代码
    # Flip horizontal
    keypoints[:, 0] = width - keypoints[:, 0] - 1

    这句代码什么意思呢?其实就是将关键点水平翻转了一下,作图帮大家理解:

    大家会不会认为这样就结束了呢,其实很没有,我们来看看代码中还做了什么,如下:

    python 复制代码
    # matched_parts这些值对应COCO人体关键点,交换人体关键点中对称的点,这个是person_keypoints.json文件中的flip_pairs
    # Change left-right parts
    for pair in self.matched_parts:
        keypoints[pair[0], :], keypoints[pair[1], :] = \
        keypoints[pair[1], :], keypoints[pair[0], :].copy()
    
        visible[pair[0]], visible[pair[1]] = \
        visible[pair[1]], visible[pair[0]].copy()

    这步交换了人体关键点中堆成的关键点,为什么要这么做,我当时就是这里疑惑了好久好久,我来画个图帮大家理解一下:

这样你可能还看不出端倪,我在画出水平翻转后的图像,如下:

你会发现如果单纯的将两个点对应过来,左右关系会对调,因此需要把标签进行左右互换。【大家这里如果觉得不好理解的话,可以自己动动手,画画图,相信你会有所收获】

transforms.KeypointToHeatMap

KeypointToHeatMap------将关键点映射为热力图。我们在理论部分说到,HRNet是基于热力图实现关键点检测,不清楚的可以去原理详解篇寻找寻找答案。那么其是怎么将关键点映射成热力图的呢,我们一起来看看代码是怎么实现的叭。

首先,先来看看其__init__函数:

python 复制代码
def __init__(self,
             heatmap_hw: Tuple[int, int] = (256 // 4, 192 // 4),
             gaussian_sigma: int = 2,
             keypoints_weights=None):
    self.heatmap_hw = heatmap_hw
    self.sigma = gaussian_sigma
    self.kernel_radius = self.sigma * 3
    self.use_kps_weights = False if keypoints_weights is None else True
    self.kps_weights = keypoints_weights

    # generate gaussian kernel(not normalized)
    kernel_size = 2 * self.kernel_radius + 1
    kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
    x_center = y_center = kernel_size // 2
    for x in range(kernel_size):
        for y in range(kernel_size):
            kernel[y, x] = np.exp(-((x - x_center) ** 2 + (y - y_center) ** 2) / (2 * self.sigma ** 2))
            # print(kernel)

            self.kernel = kernel

这段主要定义了存储热力图的宽度和高度、高斯标准差和关键点权重等信息,然后生成了一个大小为13*13的高斯核kernel(中间的值大,往四周扩散值越来越小),如下图所示:

接着我们来看__call__函数:

python 复制代码
def __call__(self, image, target):
    kps = target["keypoints"]
    num_kps = kps.shape[0]
    kps_weights = np.ones((num_kps,), dtype=np.float32)
    if "visible" in target:
        visible = target["visible"]
        kps_weights = visible

        heatmap = np.zeros((num_kps, self.heatmap_hw[0], self.heatmap_hw[1]), dtype=np.float32)
        heatmap_kps = (kps / 4 + 0.5).astype(np.int)  # round
        for kp_id in range(num_kps):
            v = kps_weights[kp_id]
            if v < 0.5:
                # 如果该点的可见度很低,则直接忽略
                continue

                x, y = heatmap_kps[kp_id]
                ul = [x - self.kernel_radius, y - self.kernel_radius]  # up-left x,y
                br = [x + self.kernel_radius, y + self.kernel_radius]  # bottom-right x,y
                # 如果以xy为中心kernel_radius为半径的辐射范围内与heatmap没交集,则忽略该点(该规则并不严格)
                if ul[0] > self.heatmap_hw[1] - 1 or \
                ul[1] > self.heatmap_hw[0] - 1 or \
                br[0] < 0 or \
                br[1] < 0:
                    # If not, just return the image as is
                    kps_weights[kp_id] = 0
                    continue

                    # Usable gaussian range
                    # 计算高斯核有效区域(高斯核坐标系)
                    g_x = (max(0, -ul[0]), min(br[0], self.heatmap_hw[1] - 1) - ul[0])
                    g_y = (max(0, -ul[1]), min(br[1], self.heatmap_hw[0] - 1) - ul[1])
                    # image range
                    # 计算heatmap中的有效区域(heatmap坐标系)
                    img_x = (max(0, ul[0]), min(br[0], self.heatmap_hw[1] - 1))
                    img_y = (max(0, ul[1]), min(br[1], self.heatmap_hw[0] - 1))

                    if kps_weights[kp_id] > 0.5:
                        # 将高斯核有效区域复制到heatmap对应区域
                        heatmap[kp_id][img_y[0]:img_y[1] + 1, img_x[0]:img_x[1] + 1] = \
                        self.kernel[g_y[0]:g_y[1] + 1, g_x[0]:g_x[1] + 1]

                        if self.use_kps_weights:
                            kps_weights = np.multiply(kps_weights, self.kps_weights)

                            plot_heatmap(image, heatmap, kps, kps_weights)

                            target["heatmap"] = torch.as_tensor(heatmap, dtype=torch.float32)
                            target["kps_weights"] = torch.as_tensor(kps_weights, dtype=torch.float32)

                            return image, target

我给大家解释一下可能难理解的地方:

python 复制代码
heatmap_kps = (kps / 4 + 0.5).astype(np.int)

这句是将关键点的坐标映射到热力图上,因为最终的热力图相较于原图像下采样了4倍,所以要除以4,这里加上0.5是起到一个四舍五入的作用,因为后面要将坐标转为int格式。

python 复制代码
ul = [x - self.kernel_radius, y - self.kernel_radius]  # up-left x,y
br = [x + self.kernel_radius, y + self.kernel_radius]  # bottom-right x,y

这两句是找到某个关键点对应热力图的左上角(ul)和右下角(br)的坐标,kernel_radius是高斯核的半径,如下图所示,hw坐标系表示热力图坐标,中间的⚪表示关键点在热力图上的坐标,坐标为(x,y):

python 复制代码
# 如果以xy为中心kernel_radius为半径的辐射范围内与heatmap没交集,则忽略该点(该规则并不严格)
if ul[0] > self.heatmap_hw[1] - 1 or \
        ul[1] > self.heatmap_hw[0] - 1 or \
        br[0] < 0 or \
        br[1] < 0:
    # If not, just return the image as is
    kps_weights[kp_id] = 0
    continue

这句是看看以xy为中心kernel_radius为半径的辐射范围内(就是上图中的正方形区域内)与heatmap(就是上图的hw坐标系,当然其h=64,w=48,并不是无线延长的坐标系)有没有交集,若无交集,则将kps_weights[kp_id]置为0。

python 复制代码
# Usable gaussian range
# 计算高斯核有效区域(高斯核坐标系)
g_x = (max(0, -ul[0]), min(br[0], self.heatmap_hw[1] - 1) - ul[0])
g_x = (max(0, -ul[1]), min(br[1], self.heatmap_hw[0] - 1) - ul[1])
# image range
# 计算heatmap中的有效区域(heatmap坐标系)
img_x = (max(0, ul[0]), min(br[0], self.heatmap_hw[1] - 1))
img_y = (max(0, ul[1]), min(br[1], self.heatmap_hw[0] - 1))

这几句分别计算高斯核有效区域和heatmap中的有效区域,为下一步将将高斯核有效区域复制到heatmap对应区域做准备:

python 复制代码
if kps_weights[kp_id] > 0.5:
    # 将高斯核有效区域复制到heatmap对应区域
    heatmap[kp_id][img_y[0]:img_y[1] + 1, img_x[0]:img_x[1] + 1] = \
        self.kernel[g_y[0]:g_y[1] + 1, g_x[0]:g_x[1] + 1]

这几句到底实现了什么呢,其实就是把高斯核kernel复制到热力图中,至于复制到什么位置,复制多少,就看g_x、g_x、img_x和img_y了。我调试帮助大家理解一下,比如现在g_x=(0,12)、g_y=(0,12)、img_x=(25,37)和img_y=(12,24)。

g_x[0]:g_x[1]+1=0:12+1、g_y[0]:g_y[1]+1=0:12+1表示复制kernel的x方向(0,12+1)范围内的值和y方向(0,12+1)范围内,你看kernel的shape你会发现,其大小为13*13,那么这个(0,12+1)就是复制整个kernel数组**(这里刚好是整个数组,你调试的话会有不同的结果)**:

那么把这个数组复制到哪里呢,其实就是热力图的对应区域,这是就用到了img_x=(25,37)和img_y=(12,24),将其复制到热力图w方向(25,37+1)和h方向(12,24+1)的位置,如下图所示:


这里展示一下图片和产生热力图的结果,如下图所示:【注:由于不是同一次调试的结果,所以这里的图像和之前的有所差异】


最后我还想说一个小点,就是kps_weights这个值,表示的是关键点的权重,如果没有指定这个参数,那么其就默认是关键点的可见性,如果指定了这个参数,其会让原来的可见性乘这个指定的参数,在HRNet中,这个kps_weights默认如下:

小结

HRNet中的在线数据增强方式到这里就为大家介绍完啦,我觉得这部分还是非常重要的,大家可以去认真的学习一下喔,不明白的可以先调试调试,实在搞不懂欢迎评论区和我交流探讨。🥂🥂🥂

网络结构搭建

HRNet的网络结构我在原理详解篇已经为大家介绍过了,也简略的为大家展示了一些代码,但是没用具体介绍网络的详细结构。这里呢,我也不打算介绍了,因为我认为网络搭建部分真的是比较简单的,就像搭积木一样,一层一层的,只要你拿起代码对照着网络结构图调试一遍就会非常清晰了。所以这里大家一定要动起小手来喔!!!🌼🌼🌼

网络训练和预测

我们一起来看看训练阶段的代码,主要看训练一个epoch的情况就好啦,即train_one_epoch函数,首先有一个热身训练的代码:

python 复制代码
lr_scheduler = None
if epoch == 0 and warmup is True:  # 当训练第一轮(epoch=0)时,启用warmup训练方式,可理解为热身训练
    warmup_factor = 1.0 / 1000
    warmup_iters = min(1000, len(data_loader) - 1)

    lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)

关于此部分代码可以从我的这篇博客-->poly学习率策略源码详解中查看详情,对这种学习率调整策略有详细解释,链接如下:

深度学习语义分割篇------DeeplabV3原理详解+源码实战🌱🌱🌱

接着就来说说for循环遍历数据集的过程,使用的for循环如下:

python 复制代码
for i, [images, targets] in enumerate(metric_logger.log_every(data_loader, print_freq, header)):

我们可以调试进入log_every函数中,注意到log_every函数中有一个yield obj,yield是python中的关键字,是一个生成器,每次log_every运行到yield obj时都会暂停执行下面的代码,而是将obj返回给调用方,这样做的目的是节省内存。

这么说我觉得大家听的还是云里雾里,我画一个图解释一下代码的运行流程:

其按顺序依次执行①②③④⑤,在执行完③时,obj会传给①,得到image和target。知道了这一点,那么剩下的内容就比较简单啦,这里就不在过多叙述咯,不清楚的大家一定要调试调试喔。


接下来就是预测过程,我们会通过网络得到输出结果,其尺寸为(1,17,64,48),1表示batch为1,后面我们需要对这个输出做一些后处理操作,使其能够将预测关键点映射到原图上。关于预测过程,我也在HRNet原理详解篇进行了介绍,不清楚的一定要去看看喔🍚🍚🍚

总结

呼呼呼~~~终于写完啦,源码解析篇就到这里结束啦,整个HRNet到这里也结束咯,如果有任何不明白的地方欢迎和我一起探讨,共同进步喔。🥂🥂🥂

参考链接

HRNet论文🍁🍁🍁

HRNet网络简介🍁🍁🍁

HRNet源码🍁🍁🍁

如若文章对你有所帮助,那就🛴🛴🛴

相关推荐
Jumbuck_101 个月前
YOLOv8-pose+streamlit 实现人体关键点检测/姿态估计系统
python·深度学习·yolo·计算机视觉·姿态估计·关键点检测·健身
Zero_one_ws1 个月前
计算机视觉四大任务模型汇总
人工智能·深度学习·计算机视觉·图像分类·关键点检测·图像目标检测·目标分割
阿利同学2 个月前
自动泊车系统中的YOLOv8 pose关键点车位线检测
yolo·自动泊车·关键点检测·车位检测·yolov8 pose·自定义模型训练·联系 qq1309399183
阿利同学2 个月前
手势识别&手势控制系统-OpenCV&Python(源码和教程)
人工智能·python·opencv·手势识别·手势控制·关键点检测
Struart_R2 个月前
REGTR: End-to-end Point Cloud Correspondences with Transformers 论文解读
人工智能·transformer·三维视觉·点云配准·姿态估计
Struart_R3 个月前
SIRA-PCR: Sim-to-Real Adaptation for 3D Point Cloud Registration 论文解读
人工智能·三维视觉·点云配准·姿态估计
钓了猫的鱼儿6 个月前
关键点检测——面部情绪数据集
目标检测·关键点检测·面部情绪识别数据集
吃鱼不卡次6 个月前
YOLOv8深度剖析专栏导航
目标检测·yolov8·实例分割·关键点检测·旋转目标检测·训练、验证、预测、导出模型·训练原理
阿利同学7 个月前
姿态估计-人脸识别mesh-3d手势识别-3d目标检测-背景分割-人脸关键点
人工智能·目标检测·3d·3d目标检测·姿态估计·关键点检测·3d手势识别