official/vision/ops/augment_test.py
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for autoaugment."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import random
from unittest import mock
from absl.testing import parameterized
import numpy as np
import tensorflow as tf, tf_keras
from official.vision.configs import common as configs
from official.vision.ops import augment
def get_dtype_test_cases():
return [
('uint8', tf.uint8),
('int32', tf.int32),
('float16', tf.float16),
('float32', tf.float32),
]
@parameterized.named_parameters(get_dtype_test_cases())
class TransformsTest(parameterized.TestCase, tf.test.TestCase):
"""Basic tests for fundamental transformations."""
def test_to_from_4d(self, dtype):
for shape in [(10, 10), (10, 10, 10), (10, 10, 10, 10)]:
original_ndims = len(shape)
image = tf.zeros(shape, dtype=dtype)
image_4d = augment.to_4d(image)
self.assertEqual(4, tf.rank(image_4d))
self.assertAllEqual(image, augment.from_4d(image_4d, original_ndims))
def test_transform(self, dtype):
image = tf.constant([[1, 2], [3, 4]], dtype=dtype)
self.assertAllEqual(
augment.transform(image, transforms=[1] * 8), [[4, 4], [4, 4]])
def test_translate(self, dtype):
image = tf.constant(
[[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], dtype=dtype)
translations = [-1, -1]
translated = augment.translate(image=image, translations=translations)
expected = [[1, 0, 1, 1], [0, 1, 0, 0], [1, 0, 1, 1], [1, 0, 1, 1]]
self.assertAllEqual(translated, expected)
def test_translate_shapes(self, dtype):
translation = [0, 0]
for shape in [(3, 3), (5, 5), (224, 224, 3)]:
image = tf.zeros(shape, dtype=dtype)
self.assertAllEqual(image, augment.translate(image, translation))
def test_translate_invalid_translation(self, dtype):
image = tf.zeros((1, 1), dtype=dtype)
invalid_translation = [[[1, 1]]]
with self.assertRaisesRegex(TypeError, 'rank 1 or 2'):
_ = augment.translate(image, invalid_translation)
def test_rotate(self, dtype):
image = tf.reshape(tf.cast(tf.range(9), dtype), (3, 3))
rotation = 90.
transformed = augment.rotate(image=image, degrees=rotation)
expected = [[2, 5, 8], [1, 4, 7], [0, 3, 6]]
self.assertAllEqual(transformed, expected)
def test_rotate_shapes(self, dtype):
degrees = 0.
for shape in [(3, 3), (5, 5), (224, 224, 3)]:
image = tf.zeros(shape, dtype=dtype)
self.assertAllEqual(image, augment.rotate(image, degrees))
def test_random_cutout_video(self, dtype):
for num_channels in (1, 2, 3):
video = tf.ones((2, 2, 2, num_channels), dtype=dtype)
video = augment.cutout_video(video)
num_zeros = np.sum(video == 0)
self.assertGreater(num_zeros, 0)
def test_cutout_video_with_fixed_shape(self, dtype):
tf.random.set_seed(0)
video = tf.ones((10, 10, 10, 1), dtype=dtype)
video = augment.cutout_video(video, mask_shape=tf.constant([2, 2, 2]))
num_zeros = np.sum(video == 0)
self.assertEqual(num_zeros, 8)
class AutoaugmentTest(tf.test.TestCase, parameterized.TestCase):
AVAILABLE_POLICIES = [
'v0',
'test',
'simple',
'reduced_cifar10',
'svhn',
'reduced_imagenet',
'detection_v0',
'vit',
]
def test_autoaugment(self):
"""Smoke test to be sure there are no syntax errors."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
for policy in self.AVAILABLE_POLICIES:
augmenter = augment.AutoAugment(augmentation_name=policy)
aug_image = augmenter.distort(image)
self.assertEqual((224, 224, 3), aug_image.shape)
def test_autoaugment_with_bboxes(self):
"""Smoke test to be sure there are no syntax errors with bboxes."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 4), dtype=tf.float32)
for policy in self.AVAILABLE_POLICIES:
augmenter = augment.AutoAugment(augmentation_name=policy)
aug_image, aug_bboxes = augmenter.distort_with_boxes(image, bboxes)
self.assertEqual((224, 224, 3), aug_image.shape)
self.assertEqual((2, 4), aug_bboxes.shape)
def test_randaug(self):
"""Smoke test to be sure there are no syntax errors."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
augmenter = augment.RandAugment()
aug_image = augmenter.distort(image)
self.assertEqual((224, 224, 3), aug_image.shape)
def test_randaug_with_bboxes(self):
"""Smoke test to be sure there are no syntax errors with bboxes."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 4), dtype=tf.float32)
augmenter = augment.RandAugment()
aug_image, aug_bboxes = augmenter.distort_with_boxes(image, bboxes)
self.assertEqual((224, 224, 3), aug_image.shape)
self.assertEqual((2, 4), aug_bboxes.shape)
def test_randaug_build_for_detection(self):
"""Smoke test to be sure there are no syntax errors built for detection."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 4), dtype=tf.float32)
augmenter = augment.RandAugment.build_for_detection()
self.assertCountEqual(augmenter.available_ops, [
'AutoContrast', 'Equalize', 'Invert', 'Posterize', 'Solarize', 'Color',
'Contrast', 'Brightness', 'Sharpness', 'Cutout', 'SolarizeAdd',
'Rotate_BBox', 'ShearX_BBox', 'ShearY_BBox', 'TranslateX_BBox',
'TranslateY_BBox'
])
aug_image, aug_bboxes = augmenter.distort_with_boxes(image, bboxes)
self.assertEqual((224, 224, 3), aug_image.shape)
self.assertEqual((2, 4), aug_bboxes.shape)
def test_all_policy_ops(self):
"""Smoke test to be sure all augmentation functions can execute."""
prob = 1
magnitude = 10
replace_value = [128] * 3
cutout_const = 100
translate_const = 250
image = tf.ones((224, 224, 3), dtype=tf.uint8)
bboxes = None
for op_name in augment.NAME_TO_FUNC.keys() - augment.REQUIRE_BOXES_FUNCS:
func, _, args = augment._parse_policy_info(op_name, prob, magnitude,
replace_value, cutout_const,
translate_const)
image, bboxes = func(image, bboxes, *args)
self.assertEqual((224, 224, 3), image.shape)
self.assertIsNone(bboxes)
def test_all_policy_ops_with_bboxes(self):
"""Smoke test to be sure all augmentation functions can execute."""
prob = 1
magnitude = 10
replace_value = [128] * 3
cutout_const = 100
translate_const = 250
image = tf.ones((224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 4), dtype=tf.float32)
for op_name in augment.NAME_TO_FUNC:
func, _, args = augment._parse_policy_info(op_name, prob, magnitude,
replace_value, cutout_const,
translate_const)
image, bboxes = func(image, bboxes, *args)
self.assertEqual((224, 224, 3), image.shape)
self.assertEqual((2, 4), bboxes.shape)
def test_autoaugment_video(self):
"""Smoke test with video to be sure there are no syntax errors."""
image = tf.zeros((2, 224, 224, 3), dtype=tf.uint8)
for policy in self.AVAILABLE_POLICIES:
augmenter = augment.AutoAugment(augmentation_name=policy)
aug_image = augmenter.distort(image)
self.assertEqual((2, 224, 224, 3), aug_image.shape)
def test_autoaugment_video_with_boxes(self):
"""Smoke test with video to be sure there are no syntax errors."""
image = tf.zeros((2, 224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 2, 4), dtype=tf.float32)
for policy in self.AVAILABLE_POLICIES:
augmenter = augment.AutoAugment(augmentation_name=policy)
aug_image, aug_bboxes = augmenter.distort_with_boxes(image, bboxes)
self.assertEqual((2, 224, 224, 3), aug_image.shape)
self.assertEqual((2, 2, 4), aug_bboxes.shape)
def test_randaug_video(self):
"""Smoke test with video to be sure there are no syntax errors."""
image = tf.zeros((2, 224, 224, 3), dtype=tf.uint8)
augmenter = augment.RandAugment()
aug_image = augmenter.distort(image)
self.assertEqual((2, 224, 224, 3), aug_image.shape)
def test_all_policy_ops_video(self):
"""Smoke test to be sure all video augmentation functions can execute."""
prob = 1
magnitude = 10
replace_value = [128] * 3
cutout_const = 100
translate_const = 250
image = tf.ones((2, 224, 224, 3), dtype=tf.uint8)
bboxes = None
for op_name in augment.NAME_TO_FUNC.keys() - augment.REQUIRE_BOXES_FUNCS:
func, _, args = augment._parse_policy_info(op_name, prob, magnitude,
replace_value, cutout_const,
translate_const)
image, bboxes = func(image, bboxes, *args)
self.assertEqual((2, 224, 224, 3), image.shape)
self.assertIsNone(bboxes)
def test_all_policy_ops_video_with_bboxes(self):
"""Smoke test to be sure all video augmentation functions can execute."""
prob = 1
magnitude = 10
replace_value = [128] * 3
cutout_const = 100
translate_const = 250
image = tf.ones((2, 224, 224, 3), dtype=tf.uint8)
bboxes = tf.ones((2, 2, 4), dtype=tf.float32)
for op_name in augment.NAME_TO_FUNC:
func, _, args = augment._parse_policy_info(op_name, prob, magnitude,
replace_value, cutout_const,
translate_const)
if op_name in {
'Rotate_BBox',
'ShearX_BBox',
'ShearY_BBox',
'TranslateX_BBox',
'TranslateY_BBox',
'TranslateY_Only_BBoxes',
}:
with self.assertRaises(ValueError):
func(image, bboxes, *args)
else:
image, bboxes = func(image, bboxes, *args)
self.assertEqual((2, 224, 224, 3), image.shape)
self.assertEqual((2, 2, 4), bboxes.shape)
def _generate_test_policy(self):
"""Generate a test policy at random."""
op_list = list(augment.NAME_TO_FUNC.keys())
size = 6
prob = [round(random.uniform(0., 1.), 1) for _ in range(size)]
mag = [round(random.uniform(0, 10)) for _ in range(size)]
policy = []
for i in range(0, size, 2):
policy.append([(op_list[i], prob[i], mag[i]),
(op_list[i + 1], prob[i + 1], mag[i + 1])])
return policy
def test_custom_policy(self):
"""Test autoaugment with a custom policy."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
augmenter = augment.AutoAugment(policies=self._generate_test_policy())
aug_image = augmenter.distort(image)
self.assertEqual((224, 224, 3), aug_image.shape)
def test_autoaugment_three_augment(self):
"""Test three augmentation."""
image = tf.random.normal(shape=(224, 224, 3), dtype=tf.float32)
augmenter = augment.AutoAugment(augmentation_name='deit3_three_augment')
aug_image = augmenter.distort(image)
self.assertEqual((224, 224, 3), aug_image.shape)
self.assertFalse(tf.math.reduce_all(image == aug_image))
@parameterized.named_parameters(
{'testcase_name': '_OutOfRangeProb',
'sub_policy': ('Equalize', 1.1, 3), 'value': '1.1'},
{'testcase_name': '_OutOfRangeMag',
'sub_policy': ('Equalize', 0.9, 11), 'value': '11'},
)
def test_invalid_custom_sub_policy(self, sub_policy, value):
"""Test autoaugment with out-of-range values in the custom policy."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
policy = self._generate_test_policy()
policy[0][0] = sub_policy
augmenter = augment.AutoAugment(policies=policy)
with self.assertRaisesRegex(
tf.errors.InvalidArgumentError,
r'Expected \'tf.Tensor\(False, shape=\(\), dtype=bool\)\' to be true. '
r'Summarized data: ({})'.format(value)):
augmenter.distort(image)
def test_invalid_custom_policy_ndim(self):
"""Test autoaugment with wrong dimension in the custom policy."""
policy = [[('Equalize', 0.8, 1), ('Shear', 0.8, 4)],
[('TranslateY', 0.6, 3), ('Rotate', 0.9, 3)]]
policy = [[policy]]
with self.assertRaisesRegex(
ValueError,
r'Expected \(:, :, 3\) but got \(1, 1, 2, 2, 3\).'):
augment.AutoAugment(policies=policy)
def test_invalid_custom_policy_shape(self):
"""Test autoaugment with wrong shape in the custom policy."""
policy = [[('Equalize', 0.8, 1, 1), ('Shear', 0.8, 4, 1)],
[('TranslateY', 0.6, 3, 1), ('Rotate', 0.9, 3, 1)]]
with self.assertRaisesRegex(
ValueError,
r'Expected \(:, :, 3\) but got \(2, 2, 4\)'):
augment.AutoAugment(policies=policy)
def test_invalid_custom_policy_key(self):
"""Test autoaugment with invalid key in the custom policy."""
image = tf.zeros((224, 224, 3), dtype=tf.uint8)
policy = [[('AAAAA', 0.8, 1), ('Shear', 0.8, 4)],
[('TranslateY', 0.6, 3), ('Rotate', 0.9, 3)]]
augmenter = augment.AutoAugment(policies=policy)
with self.assertRaisesRegex(KeyError, '\'AAAAA\''):
augmenter.distort(image)
class RandomErasingTest(tf.test.TestCase, parameterized.TestCase):
def test_random_erase_replaces_some_pixels(self):
image = tf.zeros((224, 224, 3), dtype=tf.float32)
augmenter = augment.RandomErasing(probability=1., max_count=10)
aug_image = augmenter.distort(image)
self.assertEqual((224, 224, 3), aug_image.shape)
self.assertNotEqual(0, tf.reduce_max(aug_image))
class MixupAndCutmixTest(tf.test.TestCase, parameterized.TestCase):
def test_mixup_and_cutmix_smoothes_labels(self):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal((batch_size, 224, 224, 3), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
num_classes=num_classes, label_smoothing=label_smoothing)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
def test_mixup_changes_image(self):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal((batch_size, 224, 224, 3), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
mixup_alpha=1., cutmix_alpha=0., num_classes=num_classes)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
self.assertFalse(tf.math.reduce_all(images == aug_images))
def test_cutmix_changes_image(self):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal((batch_size, 224, 224, 3), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
mixup_alpha=0., cutmix_alpha=1., num_classes=num_classes)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
self.assertFalse(tf.math.reduce_all(images == aug_images))
def test_mixup_and_cutmix_smoothes_labels_with_videos(self):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal((batch_size, 8, 224, 224, 3), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
num_classes=num_classes, label_smoothing=label_smoothing)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
@parameterized.product(num_channels=[3, 4])
def test_mixup_changes_video(self, num_channels: int):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal(
(batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
mixup_alpha=1., cutmix_alpha=0., num_classes=num_classes)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
self.assertFalse(tf.math.reduce_all(images == aug_images))
@parameterized.product(num_channels=[3, 4])
def test_cutmix_changes_video(self, num_channels: int):
batch_size = 12
num_classes = 1000
label_smoothing = 0.1
images = tf.random.normal(
(batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
labels = tf.range(batch_size)
augmenter = augment.MixupAndCutmix(
mixup_alpha=0., cutmix_alpha=1., num_classes=num_classes)
aug_images, aug_labels = augmenter.distort(images, labels)
self.assertEqual(images.shape, aug_images.shape)
self.assertEqual(images.dtype, aug_images.dtype)
self.assertEqual([batch_size, num_classes], aug_labels.shape)
self.assertAllLessEqual(aug_labels, 1. - label_smoothing +
2. / num_classes) # With tolerance
self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
1e4) # With tolerance
self.assertFalse(tf.math.reduce_all(images == aug_images))
class SSDRandomCropTest(tf.test.TestCase, parameterized.TestCase):
@parameterized.named_parameters(
dict(
testcase_name='filter first one',
bboxes=[[0, 0, 1, 1], [0, 0, 0.5, 0.5]],
crop_box=[[[0, 0, 0.5, 1]]],
min_box_overlap=0.6,
expected=[[0, 0, 0, 0], [0, 0, 0.5, 0.5]],
),
dict(
testcase_name='empty box list',
bboxes=tf.zeros([0, 4], dtype=tf.float32),
crop_box=[[[0, 0, 1, 1]]],
min_box_overlap=0.5,
expected=tf.zeros([0, 4], dtype=tf.float32),
),
)
def test_filter_boxes_by_ioa(
self, bboxes, crop_box, min_box_overlap, expected
):
new_bboxes = augment.filter_boxes_by_ioa(
bboxes=tf.constant(bboxes, dtype=tf.float32),
crop_box=tf.constant(crop_box, dtype=tf.float32),
min_box_overlap=min_box_overlap,
)
self.assertAllClose(expected, new_bboxes)
@parameterized.named_parameters(
dict(
testcase_name='whole image and box',
bboxes=[[0, 0, 1, 1], [0.1, 0.2, 0.8, 0.5]],
ori_image_size=[200, 600],
new_image_size=[100, 200],
offset=[70, 100],
expected=[[0, 0, 1, 1], [0, 0.1, 0.9, 1]],
),
dict(
testcase_name='zero size boxes',
bboxes=tf.zeros([1, 4], dtype=tf.float32),
ori_image_size=[200, 600],
new_image_size=[100, 200],
offset=[70, 100],
expected=tf.zeros([1, 4], dtype=tf.float32),
),
dict(
testcase_name='empty box list',
bboxes=tf.zeros([0, 4], dtype=tf.float32),
ori_image_size=[200, 600],
new_image_size=[100, 200],
offset=[70, 100],
expected=tf.zeros([0, 4], dtype=tf.float32),
),
)
def test_crop_normalized_boxes(
self, bboxes, ori_image_size, new_image_size, offset, expected
):
got = augment.crop_normalized_boxes(
bboxes=tf.constant(bboxes, dtype=tf.float32),
ori_image_size=tf.constant(ori_image_size, dtype=tf.int32),
new_image_size=tf.constant(new_image_size, dtype=tf.int32),
offset=tf.constant(offset, dtype=tf.int32),
)
self.assertAllClose(expected, got)
@parameterized.named_parameters(
dict(
testcase_name='uint8 image',
image=tf.zeros([320, 256, 3], dtype=tf.uint8),
),
dict(
testcase_name='float32 image',
image=tf.zeros([320, 256, 3], dtype=tf.float32),
),
)
def test_distort_with_boxes_output_shape(self, image):
bboxes = tf.constant([[0, 0, 0.5, 0.5], [0.5, 0.5, 1.0, 1.0]])
augmenter = augment.SSDRandomCrop()
new_image, new_bboxes = augmenter.distort_with_boxes(
image=image,
bboxes=bboxes,
)
self.assertDTypeEqual(new_image, image.dtype)
self.assertDTypeEqual(new_bboxes, bboxes.dtype)
self.assertShapeEqual(new_bboxes, bboxes)
self.assertAllGreaterEqual(new_bboxes, 0)
self.assertAllLessEqual(new_bboxes, 1)
def test_distort_with_empty_bboxes(self):
image = tf.zeros([320, 256, 3], dtype=tf.uint8)
bboxes = tf.zeros([0, 4], dtype=tf.float32)
augmenter = augment.SSDRandomCrop()
new_image, new_bboxes = augmenter.distort_with_boxes(
image=image,
bboxes=bboxes,
)
self.assertDTypeEqual(new_image, image.dtype)
self.assertDTypeEqual(new_bboxes, bboxes.dtype)
self.assertShapeEqual(new_bboxes, bboxes)
@parameterized.named_parameters(
dict(
testcase_name='uint8 image',
image=tf.zeros([320, 256, 3], dtype=tf.uint8),
),
dict(
testcase_name='float32 image',
image=tf.zeros([320, 256, 3], dtype=tf.float32),
),
)
def test_distort_with_boxes_run_as_tf_function(self, image):
bboxes = tf.constant([[0, 0, 0.5, 0.5], [0.5, 0.5, 1.0, 1.0]])
augmenter = augment.SSDRandomCrop()
aug_function = tf.function(augmenter.distort_with_boxes)
new_image, new_bboxes = aug_function(image=image, bboxes=bboxes)
self.assertDTypeEqual(new_image, image.dtype)
self.assertDTypeEqual(new_bboxes, bboxes.dtype)
self.assertShapeEqual(new_bboxes, bboxes)
self.assertAllGreaterEqual(new_bboxes, 0)
self.assertAllLessEqual(new_bboxes, 1)
def test_distort_with_boxes_run_as_tf_function_empty_bboxes(self):
image = tf.zeros([320, 256, 3], dtype=tf.uint8)
bboxes = tf.zeros([0, 4], dtype=tf.float32)
augmenter = augment.SSDRandomCrop()
aug_function = tf.function(augmenter.distort_with_boxes)
new_image, new_bboxes = aug_function(image=image, bboxes=bboxes)
self.assertDTypeEqual(new_image, image.dtype)
self.assertDTypeEqual(new_bboxes, bboxes.dtype)
self.assertShapeEqual(new_bboxes, bboxes)
def test_distort_with_boxes_filter_and_crop(self):
augmenter = augment.SSDRandomCrop(
params=[
configs.SSDRandomCropParam(
min_object_covered=0.0,
min_box_overlap=0.5,
prob_to_apply=1.0,
)
],
)
image = tf.zeros([320, 256, 3], dtype=tf.uint8)
bboxes = tf.constant(
[
[0., 0., 1., 1.], # filtered by low box overlap
[0.25, 0.75, 0.5, 2.], # kept with box clipped
[0.25, 0.48, 0.5, 0.75], # kept with box clipped
],
dtype=tf.float32,
)
with mock.patch.object(
tf.image, 'sample_distorted_bounding_box', autospec=True
) as mock_sample_box:
# crop box is an upper right box
offset = tf.constant([0, 128, 0], dtype=tf.int32)
new_image_size = tf.constant([160, 128, -1], dtype=tf.int32)
crop_box = tf.constant([[[0, 0.5, 0.5, 1.0]]], dtype=tf.float32)
mock_sample_box.return_value = offset, new_image_size, crop_box
new_image, new_bboxes = augmenter.distort_with_boxes(
image=image,
bboxes=bboxes,
)
self.assertAllClose(tf.zeros([160, 128, 3], dtype=tf.uint8), new_image)
self.assertAllClose(
tf.constant(
[[0., 0., 0., 0.], [0.5, 0.5, 1., 1.], [0.5, 0., 1., 0.5]],
dtype=tf.float32,
),
new_bboxes,
)
if __name__ == '__main__':
tf.test.main()