Skip to content

Fix fliplr augmentation for multi animal #1946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
biol-jsh wants to merge 2 commits into from
Closed

Fix fliplr augmentation for multi animal #1946

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3bd45a8
Fix fliplr augmentation for multi animal
biol-jsh Aug 12, 2022
df61533
Update deeplabcut/pose_estimation_tensorflow/datasets/pose_multianima…
biol-jsh Aug 15, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 12 additions & 18 deletions deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand All @@ -265,19 +264,18 @@ 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, []):
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,
Expand Down Expand Up @@ -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()
Expand All @@ -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],
Expand All @@ -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:
Expand Down