diff --git a/deepposekit/augment/FlipAxis.py b/deepposekit/augment/FlipAxis.py index 08529c7..2095275 100644 --- a/deepposekit/augment/FlipAxis.py +++ b/deepposekit/augment/FlipAxis.py @@ -14,14 +14,16 @@ # limitations under the License. import numpy as np +import imgaug.augmenters as iaa +import six.moves as sm +import h5py + from deepposekit.io.BaseGenerator import BaseGenerator -from imgaug.augmenters import meta -from imgaug import parameters as iap __all__ = ["FlipAxis"] -class FlipAxis(meta.Augmenter): +class FlipAxis(iaa.Fliplr): """ Flips the input image and keypoints across an axis. A generalized class for flipping images and keypoints @@ -37,82 +39,103 @@ class FlipAxis(meta.Augmenter): This can be a deepposekit.io.BaseGenerator for annotations or an array of integers specifying which keypoint indices to swap. - - p: int, default 0.5 - The probability that an image is flipped axis: int, default 0 Axis over which images are flipped axis=0 flips up-down (np.flipud) axis=1 flips left-right (np.fliplr) - - seed: None or int or np.random.RandomState, default None - The random state for the augmenter. - + name: None or str, default None Name given to the Augmenter object. The name is used in print(). If left as None, will print 'UnnamedX' - + deterministic: bool, default False If set to true, each batch will be augmented the same way. + random_state: None or int or np.random.RandomState, default None + The random state for the augmenter. Attributes ---------- - p: int - The probability that an image is flipped - axis: int The axis to reflect the image. swap_index: array The keypoint indices to swap when the image is flipped - """ - def __init__(self, swap_index, p=0.5, axis=0, seed=None, name=None, deterministic=False): - super(FlipAxis, self).__init__(seed=seed, name=name, random_state="deprecated", deterministic=deterministic) - self.p = iap.handle_probability_param(p, "p") + def __init__( + self, + swap_index, + p=0.5, + axis=0, + name=None, + deterministic=False, + random_state=None, + ): + + super(FlipAxis, self).__init__( + p=p, name=name, deterministic=deterministic, random_state=random_state + ) + self.axis = axis if isinstance(swap_index, BaseGenerator): if hasattr(swap_index, "swap_index"): self.swap_index = swap_index.swap_index elif isinstance(swap_index, np.ndarray): self.swap_index = swap_index - - - def _augment_batch_(self, batch, random_state, parents, hooks): - samples = self.p.draw_samples((batch.nb_rows,), - random_state=random_state) - for i, sample in enumerate(samples): - if sample >= 0.5: - - if batch.images is not None: - if self.axis == 0: - batch.images[i] = np.flipud(batch.images[i]) - if self.axis == 1: - batch.images[i] = np.fliplr(batch.images[i]) - - if batch.keypoints is not None: - kpsoi = batch.keypoints[i] - if self.axis == 0: - height = kpsoi.shape[0] - for kp in kpsoi.keypoints: - kp.y = (height-1) - kp.y + def _augment_images(self, images, random_state, parents, hooks): + """ Augments the images + + Handles the augmentation over a specified axis + + Returns + ------- + images: array + Array of augmented images. + + """ + nb_images = len(images) + samples = self.p.draw_samples((nb_images,), random_state=random_state) + for i in sm.xrange(nb_images): + if samples[i] == 1: + if self.axis == 1: + images[i] = np.fliplr(images[i]) + elif self.axis == 0: + images[i] = np.flipud(images[i]) + self.samples = samples + return images + + def _augment_keypoints(self, keypoints_on_images, random_state, parents, hooks): + """ Augments the keypoints + + Handles the augmentation over a specified axis + and swaps the keypoint labels using swap_index. + For example, the left leg will be swapped with the right leg + This is accomplished by reordering the keypoints. + + Returns + ------- + keypoints_on_images: array + Array of new coordinates of the keypoints. + + """ + nb_images = len(keypoints_on_images) + samples = self.p.draw_samples((nb_images,), random_state=random_state) + for i, keypoints_on_image in enumerate(keypoints_on_images): + if samples[i] == 1: + for keypoint in keypoints_on_image.keypoints: if self.axis == 1: - width = kpsoi.shape[1] - for kp in kpsoi.keypoints: - kp.x = (width-1) - kp.x - swapped = kpsoi.keypoints.copy() - for r in range(len(kpsoi.keypoints)): - idx = self.swap_index[r] - if idx >= 0: - kpsoi.keypoints[r] = swapped[idx] - - return batch - - def get_parameters(self): - """See :func:`~imgaug.augmenters.meta.Augmenter.get_parameters`.""" - return [self.p, self.axis, self.swap_index] \ No newline at end of file + width = keypoints_on_image.shape[1] + keypoint.x = (width - 1) - keypoint.x + elif self.axis == 0: + height = keypoints_on_image.shape[0] + keypoint.y = (height - 1) - keypoint.y + swapped = keypoints_on_image.keypoints.copy() + for r in range(len(keypoints_on_image.keypoints)): + idx = self.swap_index[r] + if idx >= 0: + keypoints_on_image.keypoints[r] = swapped[idx] + return keypoints_on_images diff --git a/deepposekit/callbacks.py b/deepposekit/callbacks.py index 19313dc..31bcf2c 100644 --- a/deepposekit/callbacks.py +++ b/deepposekit/callbacks.py @@ -113,7 +113,7 @@ def on_epoch_end(self, epoch, logs): euclidean = evaluation_dict["euclidean"] confidence = evaluation_dict["confidence"] if self.filepath is not None: - with h5py.File(self.filepath) as h5file: + with h5py.File(self.filepath, 'r+') as h5file: values = { "val_loss": np.array([logs.get("val_loss")]).reshape(1), "loss": np.array([logs.get("loss")]).reshape(1), @@ -125,8 +125,8 @@ def on_epoch_end(self, epoch, logs): for key, value in values.items(): data = h5file["logs"][key] value = np.array(value) - data.resize(tuple(value.shape)) if data.shape[0] == 0: + data.resize(tuple(value.shape)) data[:] = value else: data.resize(data.shape[0] + 1, axis=0) @@ -134,21 +134,22 @@ def on_epoch_end(self, epoch, logs): euclidean = euclidean.flatten() confidence = confidence.flatten() + y_visible = ~np.isnan(y_error[...,0].flatten()) if self.confidence_threshold: mask = confidence >= confidence_threshold euclidean = euclidean[mask] confidence = confidence[mask] - - keypoint_percentile = np.percentile( - [euclidean, confidence], [0, 5, 25, 50, 75, 95, 100], axis=1 - ).T - euclidean_perc, confidence_perc = keypoint_percentile - - euclidean_mean, confidence_mean = np.mean([euclidean, confidence], axis=1) - - logs["euclidean"] = euclidean_mean - logs["confidence"] = confidence_mean + y_visible = y_visible[mask] + + if ~y_visible.all(): + euclidean = euclidean[y_visible] + confidence[~y_visible] = 1-confidence[~y_visible] + + euclidean_perc = np.percentile(euclidean, [0, 5, 25, 50, 75, 95, 100]) + euclidean_mean = euclidean.mean() + confidence_perc = np.percentile(confidence, [0, 5, 25, 50, 75, 95, 100]) + confidence_mean = confidence.mean() if self.verbose: print( diff --git a/deepposekit/utils/keypoints.py b/deepposekit/utils/keypoints.py index cb89186..f38d39e 100644 --- a/deepposekit/utils/keypoints.py +++ b/deepposekit/utils/keypoints.py @@ -185,6 +185,7 @@ def imgaug_to_numpy(keypoints): def keypoint_errors(y_true, y_pred): + y_true[y_true<-100] = np.NaN # Non-visible points have keypoint location -99999 before augmentation y_error = y_true - y_pred euclidean = np.sqrt(np.sum(y_error ** 2, axis=-1))