关键点检测一：HRNet数据预处理（MPII）
关键点检测一：HRNet数据预处理（MPII）
- 前言
- HRNet源码
- 数据预处理
-
- 目录位置
- 源码分析
-
- JointsDataset类
-
- `__init__`
- `_get_db`
- `half_body_transform`
- `__getitem__`
- `get_affine_transform`
- `generate_target`
前言
最近在做考场行为分析的一个项目，其中我负责的是使用关键点检测算法来进行考生异常行为检测。之前只接触过分类算法，写与看的代码也只限于分类任务。而检测任务工程量太大，因此在看官方源码时非常的吃力，因此希望写博客来记录一下。
HRNet源码
GitHub地址
HRNet的项目主页
在这里想提一句，HRNet是由中科大的团队提出的，我在看HRNet的源码时觉得读的很享受，或许是中国人懂中国人吧，不仅是论文原文还是源码，读的过程中发现遣词造句与代码逻辑都非常符合中国人的思维。
- 看到评论区有人问标注文件，所以重新编辑一下：
链接：标注文件网盘链接
提取码：87nl - 相关内容：关键点检测一：HRNet训练MPII跑通代码记录
数据预处理
其实每一份开源代码，大多数情况下与模型相关的代码很少，大部分代码其实是在进行数据预处理和train、evaluate、save模型，因此HRNet源码读的过程中，model相关代码读的比较顺利，毕竟原文把思路也写的很清楚明了，所以本文也不想再分析model，而是从数据预处理开始。
目录位置
数据预处理代码位于lib/dataset，关键点检测主要的数据集是MPII和COCO，本文以MPII数据集为例说明关键点检测常用的预处理方法和注意事项。不用COCO是因为用了COCO很多API，不如MPII简单。
源码分析
mpii和coco文件的主类都是继承JointsDataset.py里的JointsDataset类，因此先从JointsDataset类开始。
JointsDataset类
__init__
定义了很多相关的参数，这些参数来自lib/config文件夹下三个参数配置文件
_get_db
这个函数是将所有的标注信息格式化后返回。由于mpii和coco数据集的标注格式不一样，因此无法用统一的代码进行读取。本文就是以mpii为例，因此为了便于分析，在JointsDataset类分析的地方直接把mpii继承之后实现的部分放过来。
# 该部分代码在`lib/dataset/mpii`
def _get_db(self):
# create train/val split
# 请先按照Readme将数据准备好
file_name = os.path.join(
self.root, 'annot', self.image_set+'.json'
)
with open(file_name) as anno_file:
anno = json.load(anno_file)
gt_db = []
for a in anno:
image_name = a['image']
# mpii标注中的center和scale是指：
# H * W的原图像中，bbox的框原来应该是四个坐标确定，这里是用center和scale两个值来表示
# bbox的center即为center， 而bbox在mpii中默认是正方形，边长（宽） = scale * 200，这个200是官方定的
c = np.array(a['center'], dtype=np.float)
s = np.array([a['scale'], a['scale']], dtype=np.float)
# Adjust center/scale slightly to avoid cropping limbs
# 因为mpii直接默认bbox为正方形，因此可能真正的bbox是矩形，调成正方形后可能会把人体某些部分给裁掉，所以直接把正方形扩大
if c[0] != -1:
c[1] = c[1] + 15 * s[1]
s = s * 1.25
# MPII uses matlab format, index is based 1,
# we should first convert to 0-based index
c = c - 1
# 用到的都只有前两维
joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
if self.image_set != 'test':
joints = np.array(a['joints'])
joints[:, 0:2] = joints[:, 0:2] - 1
joints_vis = np.array(a['joints_vis'])
assert len(joints) == self.num_joints, \
'joint num diff: {} vs {}'.format(len(joints),
self.num_joints)
joints_3d[:, 0:2] = joints[:, 0:2]
joints_3d_vis[:, 0] = joints_vis[:]
joints_3d_vis[:, 1] = joints_vis[:]
image_dir = 'images.zip@' if self.data_format == 'zip' else 'images'
gt_db.append(
{
'image': os.path.join(self.root, image_dir, image_name),
'center': c,
'scale': s,
'joints_3d': joints_3d,
'joints_3d_vis': joints_3d_vis,
'filename': '',
'imgnum': 0,
}
)
return gt_db
half_body_transform
这个函数我觉得主要是用来数据增强的时候使用，也就是说，并不是所有的数据都是全身的关节，为了增强模型的鲁棒性，也应当适当加一些半身的图像进行训练。
def half_body_transform(self, joints, joints_vis):
# 首先获得上半身和下半身的关节id，这些关节必须都是可见的
upper_joints = []
lower_joints = []
for joint_id in range(self.num_joints):
if joints_vis[joint_id][0] > 0: # 这些关节必须都是可见的
if joint_id in self.upper_body_ids:
upper_joints.append(joints[joint_id])
else:
lower_joints.append(joints[joint_id])
# 根据概率决定是上半身还是下半身
if np.random.randn() < 0.5 and len(upper_joints) > 2:
selected_joints = upper_joints
else:
selected_joints = lower_joints \
if len(lower_joints) > 2 else upper_joints
if len(selected_joints) < 2:
return None, None
selected_joints = np.array(selected_joints, dtype=np.float32)
center = selected_joints.mean(axis=0)[:2] # 计算选出来的关节的坐标中心
# 通过右下与左上得到半身区域的宽和高来得到scale
left_top = np.amin(selected_joints, axis=0)
right_bottom = np.amax(selected_joints, axis=0)
w = right_bottom[0] - left_top[0]
h = right_bottom[1] - left_top[1]
# 保证是正方形
if w > self.aspect_ratio * h:
h = w * 1.0 / self.aspect_ratio
elif w < self.aspect_ratio * h:
w = h * self.aspect_ratio
scale = np.array(
[
w * 1.0 / self.pixel_std,
h * 1.0 / self.pixel_std
],
dtype=np.float32
)
# 适当放大，避免裁剪到人
scale = scale * 1.5
return center, scale
__getitem__
def __getitem__(self, idx):
db_rec = copy.deepcopy(self.db[idx])
# 读idx图像及其标注信息
image_file = db_rec['image']
filename = db_rec['filename'] if 'filename' in db_rec else ''
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
if self.data_format == 'zip':
from utils import zipreader
data_numpy = zipreader.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
else:
data_numpy = cv2.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
if self.color_rgb:
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
if data_numpy is None:
logger.error('=> fail to read {}'.format(image_file))
raise ValueError('Fail to read {}'.format(image_file))
joints = db_rec['joints_3d']
joints_vis = db_rec['joints_3d_vis']
c = db_rec['center']
s = db_rec['scale']
score = db_rec['score'] if 'score' in db_rec else 1
r = 0
# 训练则需要数据增强：flip和rotate
if self.is_train:
# 是否用半身
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body # = 8
and np.random.rand() < self.prob_half_body): # = 0.0
c_half_body, s_half_body = self.half_body_transform(
joints, joints_vis
)
if c_half_body is not None and s_half_body is not None:
c, s = c_half_body, s_half_body
sf = self.scale_factor # 0.25
rf = self.rotation_factor # 30
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf) # 0.75 - 1.25
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
if random.random() <= 0.6 else 0 # 0 / -60 - 60
if self.flip and random.random() <= 0.5: # 水平翻转
data_numpy = data_numpy[:, ::-1, :] #原图像的w方向翻转
joints, joints_vis = fliplr_joints(
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)# 该函数在`lib/utils/transforms.py，把标注坐标进行翻转
c[0] = data_numpy.shape[1] - c[0] - 1
trans = get_affine_transform(c, s, r, self.image_size)# 把原bbox先缩放到image_size，再按box中心旋转r°
input = cv2.warpAffine(
data_numpy,
trans,
(int(self.image_size[0]), int(self.image_size[1])),
flags=cv2.INTER_LINEAR)
if self.transform:
input = self.transform(input)
for i in range(self.num_joints):
if joints_vis[i, 0] > 0.0:
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)# 对原图的transform都要记得把对应的标注也要transform
target, target_weight = self.generate_target(joints, joints_vis)
target = torch.from_numpy(target)
target_weight = torch.from_numpy(target_weight)
meta = {
'image': image_file,
'filename': filename,
'imgnum': imgnum,
'joints': joints,
'joints_vis': joints_vis,
'center': c,
'scale': s,
'rotation': r,
'score': score
}
return input, target, target_weight, meta
def fliplr_joints(joints, joints_vis, width, matched_parts):
# Flip horizontal
joints[:, 0] = width - joints[:, 0] - 1 # x坐标变为 w - x - 1
# Change left-right parts
for pair in matched_parts:
joints[pair[0], :], joints[pair[1], :] = \
joints[pair[1], :], joints[pair[0], :].copy()
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
return joints*joints_vis, joints_vis # flip后的joint为什么还有和vis相乘我还是没搞懂？？？
get_affine_transform
源码的这个函数我真的看不懂，于是我把stacked hourglass network源码里进行缩放和旋转的部分代替了源码的这个函数，发现两种方法对图像的效果是一样的，所以下面我说明的是stacked hourglass network源码里的做法。这个函数我也看了特别久，原因在于之前我对仿射变换了解很少，所以建议先学习一下仿射变换以及常见的仿射变换矩阵再来看这个函数就会简单得多。
def get_affine_transform(center, scale, res, rot=0):
# Generate transformation matrix
# 首先是缩放到res尺寸
# 缩放矩阵本来应该就是[[W,0][0,H]]，但是为什么还有第三行和第三列那两个数我想了很久才想明白
h = 200 * scale[0]
t = np.zeros((3, 3))
t[0, 0] = float(res[1]) / h
t[1, 1] = float(res[0]) / h
t[0, 2] = res[1] * (-float(center[0]) / h + .5)# 把中心变到原点
t[1, 2] = res[0] * (-float(center[1]) / h + .5)# 把中心变到原点
t[2, 2] = 1
if not rot == 0:
rot = -rot # To match direction of rotation from cropping
rot_mat = np.zeros((3,3))
rot_rad = rot * np.pi / 180
sn,cs = np.sin(rot_rad), np.cos(rot_rad)
rot_mat[0,:2] = [cs, -sn]
rot_mat[1,:2] = [sn, cs]
rot_mat[2,2] = 1
# Need to rotate around center
t_mat = np.eye(3)
t_mat[0,2] = -res[1]/2
t_mat[1,2] = -res[0]/2
t_inv = t_mat.copy()
t_inv[:2,2] *= -1
t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
t = np.dot(rot_mat, np.dot(t_mat, t))
return t
为了更好的展示每个设置的作用，我首先把下面这两行注释掉并且把r = 0，结果如下图所示，左边是注释前的，右边是注释后的。区别在于中心点的位置。
t[0, 2] = res[1] * (-float(center[0]) / h + .5)# 把中心变到原点
t[1, 2] = res[0] * (-float(center[1]) / h + .5)# 把中心变到原点
我再把r = 10，结果如下图所示，左边是注释前的，右边是注释后的。区别感觉在于旋转中心点的位置。注意：r > 0，是按逆时针旋转的。
总结：
transform是对bbox进行的，不是对原图像，因此要注意center的位置，要进行相应的平移把bbox移到想要进行的transform对应的初始坐标处。
- 缩放与平移：
res的shape是(H, W)
- 旋转
为什么要把中心移来移去？缩放变换的矩阵中心随意，只要把对应的W和H确定好就行，但是旋转就有中心一说了。我们想要缩放后的框按中心点旋转，旋转矩阵常见的起始点是远点，也即
[cos, -sin]
[sin, cos]
所以把缩放后的框移到对应的位置，就可以利用这个矩阵进行旋转了，当然也可以不移动，但是对应的旋转矩阵就要进行相应的变换，我只是解释一下源码的做法。
这里需要注意一下，图像的坐标轴和我们平时画的不一样（y轴的方向不一样）所以上面的矩阵在我们正常的坐标系里是逆时针，但在图像的坐标轴里，是我们人眼认知的顺时针。所以解释了源码里的这一句：
rot = -rot # To match direction of rotation from cropping
源码使用的旋转矩阵是正常情况下的逆时针旋转矩阵，那么会使图像顺时针转动，但是源码想要图像逆时针旋转r°，所以就把rot = -rot就变成了逆时针。
def affine_transform(pt, t):
#把对应的gt也进行相应的transform
# 2,3 * 3, --->2,
new_pt = np.array([pt[0], pt[1], 1.]).T
new_pt = np.dot(t, new_pt)
return new_pt[:2]
generate_target
关键点检测主流做法还是以热图作为ground truth，通过MSE进行优化。
def generate_target(self, joints, joints_vis):
'''
:param joints: [num_joints, 3]
:param joints_vis: [num_joints, 3]
:return: target, target_weight(1: visible, 0: invisible)
'''
target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
target_weight[:, 0] = joints_vis[:, 0]
assert self.target_type == 'gaussian', \
'Only support gaussian map now!'
if self.target_type == 'gaussian':
# 生成heatmap_size大小的高斯热图
target = np.zeros((self.num_joints,
self.heatmap_size[1],
self.heatmap_size[0]),
dtype=np.float32)
tmp_size = self.sigma * 3 # 高斯半径的大小
for joint_id in range(self.num_joints):
feat_stride = self.image_size / self.heatmap_size
mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
# Check that any part of the gaussian is in-bounds
ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
or br[0] < 0 or br[1] < 0:
# If not, just return the image as is
target_weight[joint_id] = 0
continue
# # Generate gaussian
size = 2 * tmp_size + 1
x = np.arange(0, size, 1, np.float32)
y = x[:, np.newaxis]
x0 = y0 = size // 2
# The gaussian is not normalized, we want the center value to equal 1
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))
# Usable gaussian range
g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
# Image range
img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])
v = target_weight[joint_id]
if v > 0.5:
target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
if self.use_different_joints_weight:
target_weight = np.multiply(target_weight, self.joints_weight)
return target, target_weight
