KerasAug

Description

KerasAug is a library that includes pure TF/Keras preprocessing and augmentation layers, providing support for various data types such as images, labels, bounding boxes, segmentation masks, and more.

Note The image on the left provides the visualization of the layers in KerasAug; the image on the right provides the visualization of the YOLOV8 pipeline using KerasAug

KerasAug aims to provide fast and user-friendly preprocessing and augmentation layers, facilitating seamless integration with TensorFlow, Keras, and KerasCV.

KerasAug is:

• built entirely using TensorFlow, TensorFlow Probability, Keras and KerasCV
• supporting various data types, including images, labels, bounding boxes, segmentation masks, and more.
• compatible with GPU (partially compatible with TPU/XLA)
• seamlessly integrating with the tf.data and tf.keras.Model API
• cosistent with officially published implementations

Why Use KerasAug Rather than KerasCV for Preprocessing/Augmentation?

1. KerasAug is generally faster than KerasCV

   See benchmarks for more details

2. The APIs of KerasAug are highly stable compared to KerasCV

   While KerasCV may encounter issues running examples/tutorial/advanced_pipeline.py, KerasAug is capable of running the same script successfully.

3. The implementations in KerasAug are consistent with popular libraries like torchvision, Albumentations and ultralytics

4. KerasAug provides the functionality of sanitizing bounding boxes

   RandomAffine with bounding_box_min_area_ratio and bounding_box_max_aspect_ratio

   You can use keras_aug.utils.bounding_box.sanitize_bounding_boxes to customize your self-defined preprocessing/augmentation layer

Installation

pip install keras-aug keras-cv tensorflow tensorflow_probability --upgrade

Warning KerasAug is NOT compatible with keras-cv < 0.5.0.

Quickstart

import keras_aug
import keras_cv
import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow import keras

# Create a preprocessing pipeline using KerasAug
BATCH_SIZE = 16
NUM_CLASSES = 3
augmenter = keras.Sequential(
    [
        keras_aug.layers.RandomFlip(),
        keras_aug.layers.RandAugment(
            value_range=(0, 255),
            augmentations_per_image=3,
            magnitude=15,  # [0, 30]
            magnitude_stddev=0.15,
        ),
        keras_aug.layers.CutMix(),
    ]
)


def preprocess_data(images, labels, augment=False):
    labels = tf.one_hot(labels, NUM_CLASSES)
    inputs = {"images": images, "labels": labels}
    outputs = augmenter(inputs) if augment else inputs
    return outputs["images"], outputs["labels"]


train_dataset, test_dataset = tfds.load(
    "rock_paper_scissors", as_supervised=True, split=["train", "test"]
)
train_dataset = (
    train_dataset.batch(BATCH_SIZE)
    .map(
        lambda x, y: preprocess_data(x, y, augment=True),
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    .prefetch(tf.data.AUTOTUNE)
)
test_dataset = (
    test_dataset.batch(BATCH_SIZE)
    .map(preprocess_data, num_parallel_calls=tf.data.AUTOTUNE)
    .prefetch(tf.data.AUTOTUNE)
)

# Create a model using a pretrained backbone
backbone = keras_cv.models.EfficientNetV2Backbone.from_preset(
    "efficientnetv2_b0_imagenet"
)
model = keras_cv.models.ImageClassifier(
    backbone=backbone,
    num_classes=NUM_CLASSES,
    activation="softmax",
)
model.compile(
    loss="categorical_crossentropy",
    optimizer=keras.optimizers.Adam(learning_rate=1e-5),
    metrics=["accuracy"],
)

# Train your model
model.fit(
    train_dataset,
    validation_data=test_dataset,
    epochs=8,
)

# KerasCV Quickstart
...
Epoch 8/8
158/158 [==============================] - 39s 242ms/step - loss: 0.7930 - accuracy: 0.7171 - val_loss: 0.2488 - val_accuracy: 0.9946

# KerasAug Quickstart
...
Epoch 8/8
158/158 [==============================] - 34s 215ms/step - loss: 0.7680 - accuracy: 0.7567 - val_loss: 0.2639 - val_accuracy: 1.0000

KerasAug runs faster (215ms/step vs. 242ms/step) than KerasCV and achieves better performance.

YOLOV8 Training Pipeline Demonstration

See https://kerasaug.readthedocs.io/en/latest/get_started/tutorial.html

Benchmark

Please refer to benchmarks/README.md for more details.

KerasAug is generally faster than KerasCV.

Unit: FPS (frames per second)

Citing KerasAug

KerasAug is an extension of KerasCV and it would be preferable to acknowledge and cite both libraries jointly.

@misc{wood2022kerascv,
  title={KerasCV},
  author={Wood, Luke and Tan, Zhenyu and Stenbit, Ian and Bischof, Jonathan and Zhu, Scott and Chollet, Fran\c{c}ois and others},
  year={2022},
  howpublished={\url{https://github.com/keras-team/keras-cv}},
}

@misc{chiu2023kerasaug,
  title={KeraAug},
  author={Hongyu, Chiu},
  year={2023},
  howpublished={\url{https://github.com/james77777778/keras-aug}},
}