def choose_part(self, sample):
keypoint_name = ""
keypoints = list(extract_keypoints(sample['target'], self.dataset)[0].keys())
if len(keypoints) > 0:
if self.mode == "all":
keypoint_name = random.choice(keypoints + list(part_mapping[self.dataset].keys()))
elif self.mode == "keypoint":
keypoint_name = random.choice(keypoints)
else: # Fallback to parts
parts = list(part_mapping[self.dataset].keys())
filtered_parts = [part for part in parts if part not in (list(advanced_parts[self.dataset].keys()))]
keypoint_name = random.choice(filtered_parts)
return keypoint_name
def __call__(self, sample):
img = sample['image']
keypoint_list = extract_keypoints(sample['target'], self.dataset)
img, center_pos, widths = cutout_img(img, keypoint_list, self.part, mean_coloring=self.mean_coloring, dataset = self.dataset, width=self.width)
sample['image'] = img
if self.part in part_mapping[self.dataset].keys() and self.remove_anns:
keypoints_arr = sample['target']
joints_vis = sample['joints_vis']
keypoints_arr, joints_vis = remove_annotation(keypoints_arr, joints_vis, center_pos, widths)
sample['target'] = keypoints_arr
sample['joints_vis'] = joints_vis
return sample
def __call__(self, sample):
keypoints = list(
extract_keypoints(sample['target'], self.dataset)[0].keys())
num_iter = min(len(keypoints), self.num_keypoints)
for i in range(num_iter):
chance = random.random()
if chance < self.probability:
keypoint_name = random.choice(keypoints)
keypoints.remove(keypoint_name)
chosen_transform = random.choice(self.transforms)
chosen_transform.part = keypoint_name
sample = chosen_transform(sample)
return sample
def __call__(self, sample):
img = sample['image']
keypoint_list = extract_keypoints(sample['target'], self.dataset)
self.part = random.choice(self.parts)
if self.check_if_has_part(keypoint_list[0], self.part):
img, center_pos, widths = self.perform_cutmix(img, keypoint_list)
sample['image'] = img
if self.part in part_mapping[self.dataset].keys() and self.remove_anns:
keypoints_arr = sample['target']
joints_vis = sample['joints_vis']
keypoints_arr, joints_vis = remove_annotation(keypoints_arr, joints_vis, center_pos, widths)
sample['target'] = keypoints_arr
sample['joints_vis'] = joints_vis
return sample