def __getitem__(self, idx):
"""
Given the id of an image, return the image and the corresponding
confidence maps and the parts affinity fields.
"""
# Get a specific image id from the list of image ids.
img_index = self.img_indices[idx]
# Load the image
img_name = os.path.join(self.img_dir, get_image_name(img_index))
img = cv2.imread(img_name).transpose(1, 0, 2) / 255
original_shape = img.shape[:2]
# Resize image to 400x400 dimensions.
img = cv2.resize(normalize(img), (img_size, img_size))
# Get the annotation id of the annotaions about the image.
annotations_indices = self.coco.getAnnIds(img_index)
# Load the annotations from the annotaion ids.
annotations = self.coco.loadAnns(annotations_indices)
keypoints = []
mask = np.zeros(
(img_size // transform_scale, img_size // transform_scale),
np.uint8)
for annotation in annotations:
if annotation['num_keypoints'] != 0:
keypoints.append(annotation['keypoints'])
mask = mask | cv2.resize(
self.coco.annToMask(annotation),
(img_size // transform_scale, img_size // transform_scale))
# Add neck joints to the list of keypoints
keypoints = add_neck_joint(keypoints)
# Adjust keypoints according to resized images.
keypoints = adjust_keypoints(keypoints, original_shape)
conf_maps = generate_confidence_maps(keypoints)
paf = generate_paf(keypoints)
paf = paf.reshape(paf.shape[0], paf.shape[1],
paf.shape[2] * paf.shape[3])
return img, conf_maps, paf, mask.transpose()
def gen_data(all_keypoints, batch_size=64, val=False, affine_transform=True):
"""
Generate batches of training data.
Inputs:
all_keypoints:
"""
batch_count = len(all_keypoints.keys()) // batch_size
count = 0
# Loop over all keypoints in batches
for batch in range(1, batch_count * batch_size + 1, batch_size):
count += 1
images = np.zeros((batch_size, im_width, im_height, 3), dtype=np.uint8)
# Loop over all individual indices in a batch
for image_id in range(batch, batch + batch_size):
img_name = get_image_name(image_id - 1)
if val:
img = cv2.imread(
os.path.join(dataset_dir, 'new_val2017', img_name))
else:
img = cv2.imread(
os.path.join(dataset_dir, 'new_train2017', img_name))
images[image_id % batch] = img
conf_maps = generate_confidence_maps(
all_keypoints, range(batch - 1, batch + batch_size - 1))
pafs = generate_paf(all_keypoints,
range(batch - 1, batch + batch_size - 1))
yield images, conf_maps, pafs
# Handle cases where the total size is not a multiple of batch_size
if len(all_keypoints.keys()) % batch_size != 0:
start_index = batch_size * batch_count
final_index = list(all_keypoints.keys())[-1]
# print(final_index + 1 - start_index)
images = np.zeros(
(final_index + 1 - start_index, im_width, im_height, 3),
dtype=np.uint8)
for image_id in range(start_index, final_index + 1):
img_name = get_image_name(image_id)
if val:
img = cv2.imread(
os.path.join(dataset_dir, 'new_val2017', img_name))
else:
img = cv2.imread(
os.path.join(dataset_dir, 'new_train2017', img_name))
images[image_id % batch_size] = img
conf_maps = generate_confidence_maps(
all_keypoints, range(start_index, final_index + 1))
pafs = generate_paf(all_keypoints, range(start_index, final_index + 1))
yield images, conf_maps, pafs