From 3bd45a894951c89d3e866c0a35cb27c3414e0b93 Mon Sep 17 00:00:00 2001 From: biol-jsh Date: Fri, 12 Aug 2022 07:30:24 -0600 Subject: [PATCH 1/2] Fix fliplr augmentation for multi animal The current implementation of fliplr breaks when keypoints are hidden since only labels for visible points are forwarded to the pipeline. --- .../datasets/pose_multianimal_imgaug.py | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py index 7681f86e27..70d9cf617b 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py @@ -251,7 +251,6 @@ def get_batch(self): batch_images = [] batch_joints = [] joint_ids = [] - inds_visible = [] data_items = [] for i in range(self.batch_size): data_item = self.data[img_idx[i]] @@ -266,18 +265,17 @@ def get_batch(self): if self.has_gt: Joints = data_item.joints kpts = np.zeros((self._n_kpts * self._n_animals, 2)) + for j in range(self._n_animals): for n, x, y in Joints.get(j, []): kpts[j * self._n_kpts + int(n)] = x, y - joint_id = [ - Joints[person_id][:, 0].astype(int) for person_id in Joints.keys() - ] - joint_ids.append(joint_id) + batch_joints.append(kpts) - inds_visible.append(np.flatnonzero(np.all(kpts != 0, axis=1))) + joint_id = np.array(list(range(self._n_kpts))*self._n_animals) + joint_ids.append(joint_id) batch_images.append(image) - return batch_images, joint_ids, batch_joints, inds_visible, data_items + return batch_images, joint_ids, batch_joints, data_items def get_targetmaps_update( self, joint_ids, joints, data_items, sm_size, scale, @@ -343,7 +341,7 @@ def calc_target_and_scoremap_sizes(self): def next_batch(self, plotting=False): while True: - batch_images, joint_ids, batch_joints, inds_visible, data_items = self.get_batch() + batch_images, joint_ids, batch_joints, data_items = self.get_batch() # Scale is sampled only once (per batch) to transform all of the images into same size. target_size, sm_size = self.calc_target_and_scoremap_sizes() @@ -357,8 +355,10 @@ def next_batch(self, plotting=False): # Discard keypoints whose coordinates lie outside the cropped image batch_joints_valid = [] joint_ids_valid = [] - for joints, ids, visible in zip(batch_joints, joint_ids, inds_visible): - joints = joints[visible] + for joints, ids in zip(batch_joints, joint_ids): + #invisible joints are represented by nans + mask = ~np.isnan(joints[:,0]) + joints = joints[mask,:] inside = np.logical_and.reduce( ( joints[:, 0] < image_shape[1], @@ -368,13 +368,7 @@ def next_batch(self, plotting=False): ) ) batch_joints_valid.append(joints[inside]) - temp = [] - start = 0 - for array in ids: - end = start + array.size - temp.append(array[inside[start:end]]) - start = end - joint_ids_valid.append(temp) + joint_ids_valid.append([ids[inside]]) # If you would like to check the augmented images, script for saving # the images with joints on: From df615337aa7b1b5abaacb5fe2861b2040568d907 Mon Sep 17 00:00:00 2001 From: biol-jsh Date: Mon, 15 Aug 2022 03:03:35 -0600 Subject: [PATCH 2/2] Update deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py Co-authored-by: Jessy Lauer <30733203+jeylau@users.noreply.github.com> --- .../datasets/pose_multianimal_imgaug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py index 70d9cf617b..f3f2ce5992 100644 --- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py +++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py @@ -264,7 +264,7 @@ def get_batch(self): ) if self.has_gt: Joints = data_item.joints - kpts = np.zeros((self._n_kpts * self._n_animals, 2)) + kpts = np.full((self._n_kpts * self._n_animals, 2), np.nan) for j in range(self._n_animals): for n, x, y in Joints.get(j, []):