Tensorflow2 imgaug pipline

[johnny-mueller]

I have already tried to answer the question on the known support pages but have not found a satisfying answer yet, so I will try it here.

I want to convert data augmentation with imgaug into Tensorflow. The current state is that I have a working form in python but due to the conversion of the bounding boxes the code does not seem to work. Is it possible to implement a data augmentation pipline in Tensorflow with tf.data.Dataset for object detection (with transformation of bounding boxes)?

With the function call 3 parameters are passed - a batch of images, a batch of several bounding boxes per image and the image size.

In the following you will find the code that is not yet working in Tensorflow 2 but in python.

Function call:

train_dataset = train_dataset.map(lambda x, y: (image_batch_augmentation(x, y, image_size)))

@tf.function
def tf_image_batch_augmentation(batch_images, batch_bbox, image_shape):

    def create_BoundingBox(bbox, w, h):
        return BoundingBox(bbox[0]*h, bbox[1]*w, bbox[2]*h, bbox[3]*w, tf.cast(bbox[4], tf.int32))

    def dissolve_BoundingBox_to_array(bbox, image_size):
        # returns a numpy array with normalized bounding box coordinates
        return np.asarray([bbox.x1, bbox.y1, bbox.x2, bbox.y2, bbox.label]).astype('float32') / image_size
 
    def add_zerso_BoundingBox(batch):
        return np.concatenate((batch, np.zeros(shape=((100-len(batch)),5), dtype='float32')))

    w = image_shape
    h = image_shape

    # Create BoundingBox-Object from imgaug framework
    bbox = [[create_BoundingBox(bbox, float(w), float(h)) for bbox in batch if sum(bbox) != 0] for batch in batch_bbox]
    # Proprocess BoundingBoxes for Batch Processing
    bbox = [BoundingBoxesOnImage(batch, shape=(int(w),int(h))) for batch in bbox]

    # Activate deterministic so that the image and the corresponding bounding boxes are transformed identically
    seq_det = seq.to_deterministic()

    # Data augmentation step
    images_aug, bbsoi_aug = seq_det(images=batch_images.numpy(), bounding_boxes=bbox)

    bbsoi_aug = [batch.remove_out_of_image().clip_out_of_image()  for batch in bbsoi_aug]

    bbsoi = np.asarray([[dissolve_BoundingBox_to_array(box, w) for box in batch]  for batch in bbsoi_aug])
    bbsoi = ([add_zerso_BoundingBox(batch) for batch in bbsoi])
    
    return images_aug, bbsoi