def next_batch(self):
while True:
batch_images, joint_ids, batch_joints, data_items, sm_size, target_size = self.get_batch()
pipeline = self.build_augmentation_pipeline(
height=target_size[0],
width=target_size[1],
apply_prob=0.5)
batch_images, batch_joints = pipeline(images=batch_images,
keypoints=batch_joints)
#If you would like to check the augmented images, script for saving
#the images with joints on:
#import imageio
#for i in range(self.batch_size):
# joints = batch_joints[i]
# kps = KeypointsOnImage([Keypoint(x=joint[0], y=joint[1]) for joint in joints], shape=batch_images[i].shape)
# im = kps.draw_on_image(batch_images[i])
# imageio.imwrite('some_location/augmented/'+str(i)+'.png', im)
image_shape = arr(batch_images).shape[1:3]
batch = {Batch.inputs: arr(batch_images).astype(np.float64)}
if self.has_gt:
scmap_update = self.get_scmap_update(joint_ids, batch_joints, data_items, sm_size, image_shape)
batch.update(scmap_update)
batch = {key: data_to_input_batch(data) for (key, data) in batch.items()}
batch[Batch.data_item] = data_items
return batch
def next_batch(self, plotting=False):
while True:
(
batch_images,
joint_ids,
batch_joints,
data_items,
sm_size,
target_size,
) = self.get_batch()
pipeline = self.build_augmentation_pipeline(height=target_size[0],
width=target_size[1],
apply_prob=0.5)
batch_images, batch_joints = pipeline(images=batch_images,
keypoints=batch_joints)
# If you would like to check the augmented images, script for saving
# the images with joints on:
if plotting:
for i in range(self.batch_size):
joints = batch_joints[i]
kps = KeypointsOnImage(
[Keypoint(x=joint[0], y=joint[1]) for joint in joints],
shape=batch_images[i].shape,
)
im = kps.draw_on_image(batch_images[i])
# imageio.imwrite(data_items[i].im_path.split('/')[-1],im)
imageio.imwrite(
os.path.join(self.cfg["project_path"],
str(i) + ".png"), im)
image_shape = arr(batch_images).shape[1:3]
batch = {Batch.inputs: arr(batch_images).astype(np.float64)}
if self.has_gt:
targetmaps = self.get_targetmaps_update(
joint_ids, batch_joints, data_items, sm_size, image_shape)
batch.update(targetmaps)
# if returndata:
# return batch_images,batch_joints,targetmaps
batch = {
key: data_to_input_batch(data)
for (key, data) in batch.items()
}
batch[Batch.data_item] = data_items
return batch
