先看数据处理部分(datasets.py)
1.CocoKeypoints类中的getitem
取的index实际上就是一个随机值,也就是说从train数据集中随机取一张图片,因为论文中用的coco数据集,所以使用的都是coco官方提供的一些api,方便提取标注、图片信息。
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``.
"""
image_id = self.ids[index]
ann_ids = self.coco.getAnnIds(imgIds=image_id, catIds=self.cat_ids)
anns = self.coco.loadAnns(ann_ids)
anns = copy.deepcopy(anns)
image_info = self.coco.loadImgs(image_id)[0]
self.log.debug(image_info)
with open(os.path.join(self.root, image_info['file_name']), 'rb') as f:
image = Image.open(f).convert('RGB')
meta_init = {
'dataset_index': index,
'image_id': image_id,
'file_name': image_info['file_name'],
}
image, anns, meta = self.preprocess(image, anns, None)
if isinstance(image, list):
return self.multi_image_processing(image, anns, meta, meta_init)
return self.single_image_processing(image, anns, meta, meta_init)
然后就是跳到multi_image_processing或是single_image_processing中,我们可以看出multi_image_processing也是基于single_image_processing
def multi_image_processing(self, image_list, anns_list, meta_list, meta_init):
return list(zip(*[
self.single_image_processing(image, anns, meta, meta_init)
for image, anns, meta in zip(image_list, anns_list, meta_list)
]))
所以我们只看single_image_processing部分就行。
2.single_image_processing
这段代码中最关键的是get_ground_truth,获取GT。
def single_image_processing(self, image, anns, meta, meta_init):
meta.update(meta_init)
# transform image
original_size = image.size
image = self.image_transform(image)
assert image.size(2) == original_size[0]
assert image.size(1) == original_size[1]
# mask valid
valid_area = meta['valid_area']
utils.mask_valid_area(image, valid_area)
self.log.debug(meta)
heatmaps, pafs = self.get_ground_truth(anns)
heatmaps = torch.from_numpy(
heatmaps.transpose((2, 0, 1)).astype(np.float32))
pafs = torch.from_numpy(pafs.transpose((2, 0, 1)).astype(np.float32))
return image, heatmaps, pafs
3.get_ground_truth
这是整体代码,下面我们分开来说
def get_ground_truth(self, anns):
grid_y = int(self.input_y / self.stride)
grid_x = int(self.input_x / self.stride)
channels_heat = (self.HEATMAP_COUNT + 1)
channels_paf = 2 * len(self.LIMB_IDS)
heatmaps = np.zeros((int(grid_y), int(grid_x), channels_heat))
pafs = np.zeros((int(grid_y), int(grid_x), channels_paf))
keypoints = []
for ann in anns:
single_keypoints = np.array(ann['keypoints']).reshape(17,3)
single_keypoints = self.add_neck(single_keypoints)
keypoints.append(single_keypoints)
keypoints = np.array(keypoints)
keypoints = self.remove_illegal_joint(keypoints)
# confidance maps for body parts
for i in range(self.HEATMAP_COUNT):
joints = [jo[i] for jo in keypoints]#每一种关节点
for joint in joints:#遍历每一个点
if joint[2] > 0.5:#1是标注被遮挡 2是标注且没被遮挡
center = joint[:2]#点坐标
gaussian_map = heatmaps[:, :, i]
heatmaps[:, :, i] = putGaussianMaps(
center, gaussian_map,
7.0, grid_y, grid_x, self.stride)
# pafs
for i, (k1, k2) in enumerate(self.LIMB_IDS):
# limb
count = np.zeros((int(grid_y), int(grid_x)), dtype=np.uint32) # 表示该位置是否被计算了多次(计算的数量)
for joint in keypoints:
if joint[k1, 2] > 0.5 and joint[k2, 2] > 0.5:
centerA = joint[k1, :2]
centerB = joint[k2, :2]
vec_map = pafs[:, :, 2 * i:2 * (i + 1)] #每一个躯干位置,选择x和y两个方向
pafs[:, :, 2 * i:2 * (i + 1)], count = putVecMaps(
centerA=centerA,
centerB=centerB,
accumulate_vec_map=vec_map,
count=count, grid_y=grid_y, grid_x=grid_x, stride=self.stride
)
# background
heatmaps[:, :, -1] = np.maximum(
1 - np.max(heatmaps[:, :, :self.HEATMAP_COUNT], axis=2),
0.
)
return heatmaps, pafs
我们的input_y和input_x=368,经过三次下采样也就是2^3=8,得到的特征图大小为46*46(368/8=46),即grid_x =grid_y=46
HEATMAP_COUNT=18,其中标注中有17个关键点,还有一个neck是通过左右肩膀自己手动计算出来的。
len(self.LIMB_IDS)=19,18个关键点有19种连接方式,(怎么连接人家已经给好了)。
所以这里面 channels_heat =19,加的1是表示背景。 channels_paf =38,乘2表示用x和y两个方向表示一个向量。
然后初始化heatmaps 和pafs ,分别初始化为(46,46,19)和(46,46,38)的全零矩阵。
grid_y = int(self.input_y / self.stride)
grid_x = int(self.input_x / self.stride)
channels_heat = (self.HEATMAP_COUNT + 1)
channels_paf = 2 * len(self.LIMB_IDS)
heatmaps = np.zeros((int(grid_y), int(grid_x), channels_heat))
pafs = np.zeros((int(grid_y), int(grid_x), channels_paf))
这里其实就是将标注中的keypoints 提取出来保存下来,并且还加上我们计算的neck 关键点。最后再过滤一遍keypoints 数据
(数据中可能有标注错误的情况,如果标注出来的位置x,y都超过了特征图大小范围(超出图片了),我们就把他过滤掉)
注:这里要提一句
在ann标注中,关键点(x,y,k)除了有x,y的坐标值以外还有一个k,k属于0,1,2中的一个。
k=0表示原图标注中就没有这个关键点(相对应的x,y值也是随机值,没有用,后面会过滤掉)
k=1表示关键点被遮挡,但是给出关键点的坐标
k=2表示正常的关键点,没有被遮挡
我们只用k=1和2的关键点。
所以这里的np.array(ann[‘keypoints’]).reshape(17,3),中的17 表示标注中提供的17个关键点,3表示关键点对应的x,y,k
keypoints = []
for ann in anns:
single_keypoints = np.array(ann['keypoints']).reshape(17,3)
single_keypoints = self.add_neck(single_keypoints)
keypoints.append(single_keypoints)
keypoints = np.array(keypoints)
keypoints = self.remove_illegal_joint(keypoints)
然后开始构建heatmaps:
遍历第i个关键点:
从keypoints中取出每个人的关键点,将每个人第i个关键点存放在joints中。(这里的keypoints包含了所有人的所有关键点)
遍历第j个joints:
只取k=1或2的joint
保存joints的x,y
基于x,y生成高斯图(每次要基于之前的高斯图,因为重叠的部分要叠加)
# confidance maps for body parts
for i in range(self.HEATMAP_COUNT):
joints = [jo[i] for jo in keypoints]#每一种关节点
for joint in joints:#遍历每一个点
if joint[2] > 0.5:#1是标注被遮挡 2是标注且没被遮挡
center = joint[:2]#点坐标
gaussian_map = heatmaps[:, :, i]
heatmaps[:, :, i] = putGaussianMaps(
center, gaussian_map,
7.0, grid_y, grid_x, self.stride)
然后开始构建pafs:
遍历第i个limb,同时拿到k1,k2信息(由k1,k2构成一个limb,LIMB_IDS里有指定都有哪些关键点相连):
定义一个count (46,46),用于记录每个点都包含几个limb信息(用于后面计算重叠部分,向量的重叠不是叠加,是取平均)
遍历第j个keypoints:(第j个人)
先看这个人的k1,k2关键点的k是否属于1或2
保存k1的x,y
保存k2的x,y
对于每个limb构建2个pafs
由k1的x,y和k2的x,y计算pafs
# pafs
for i, (k1, k2) in enumerate(self.LIMB_IDS):
# limb
count = np.zeros((int(grid_y), int(grid_x)), dtype=np.uint32) # 表示该位置是否被计算了多次(计算的数量)
for joint in keypoints:
if joint[k1, 2] > 0.5 and joint[k2, 2] > 0.5:
centerA = joint[k1, :2]
centerB = joint[k2, :2]
vec_map = pafs[:, :, 2 * i:2 * (i + 1)] #每一个躯干位置,选择x和y两个方向
pafs[:, :, 2 * i:2 * (i + 1)], count = putVecMaps(
centerA=centerA,
centerB=centerB,
accumulate_vec_map=vec_map,
count=count, grid_y=grid_y, grid_x=grid_x, stride=self.stride
)
。。。。未完待续