lightweight处理关键点的顺序

dataset = CocoTrainDataset(prepared_train_labels, train_images_folder,
stride, sigma, path_thickness,
transform=transforms.Compose([
ConvertKeypoints(),
Scale(),
Rotate(pad=(128, 128, 128)),
CropPad(pad=(128, 128, 128)),
Flip()]))
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

for epochId in range(current_epoch, 280):
scheduler.step()
total_losses = [0, 0] * (num_refinement_stages + 1) # heatmaps loss, paf loss per stage
batch_per_iter_idx = 0
for batch_data in train_loader:
if batch_per_iter_idx == 0:
optimizer.zero_grad()

images = batch_data['image'].cuda()
keypoint_masks = batch_data['keypoint_mask'].cuda()
paf_masks = batch_data['paf_mask'].cuda()
keypoint_maps = batch_data['keypoint_maps'].cuda()
paf_maps = batch_data['paf_maps'].cuda()

这一部分先看CocoTrainDataset。CocoTrainDataset里发现convertkeypoints对读取的关键点顺序做更改自己做heatmap和paf。

class ConvertKeypoints:
def __call__(self, sample):
label = sample['label']
h, w, _ = sample['image'].shape
keypoints = label['keypoints']
for keypoint in keypoints: # keypoint[2] == 0: occluded, == 1: visible, == 2: not in image
if keypoint[0] == keypoint[1] == 0:
keypoint[2] = 2
if (keypoint[0] < 0
or keypoint[0] >= w
or keypoint[1] < 0
or keypoint[1] >= h):
keypoint[2] = 2
for other_label in label['processed_other_annotations']:
keypoints = other_label['keypoints']
for keypoint in keypoints:
if keypoint[0] == keypoint[1] == 0:
keypoint[2] = 2
if (keypoint[0] < 0
or keypoint[0] >= w
or keypoint[1] < 0
or keypoint[1] >= h):
keypoint[2] = 2
label['keypoints'] = self._convert(label['keypoints'], w, h)

for other_label in label['processed_other_annotations']:
other_label['keypoints'] = self._convert(other_label['keypoints'], w, h)
return sample

def _convert(self, keypoints, w, h):
# Nose, Neck, R hand, L hand, R leg, L leg, Eyes, Ears
reorder_map = [1, 7, 9, 11, 6, 8, 10, 13, 15, 17, 12, 14, 16, 3, 2, 5, 4]
converted_keypoints = list(keypoints[i - 1] for i in reorder_map)
converted_keypoints.insert(1, [(keypoints[5][0] + keypoints[6][0]) / 2,
(keypoints[5][1] + keypoints[6][1]) / 2, 0]) # Add neck as a mean of shoulders
if keypoints[5][2] == 2 or keypoints[6][2] == 2:
converted_keypoints[1][2] = 2
elif keypoints[5][2] == 1 and keypoints[6][2] == 1:
converted_keypoints[1][2] = 1
if (converted_keypoints[1][0] < 0
or converted_keypoints[1][0] >= w
or converted_keypoints[1][1] < 0
or converted_keypoints[1][1] >= h):
converted_keypoints[1][2] = 2
return converted_keypoints

从这个函数看出来调换的keypoint位置自己做的脖子用未调换的做出来。然后放到1位置。