Bump tornado from 6.4 to 6.4.1

augmentations/augmentation.py

@@ -0,0 +1,12 @@
+from abc import abstractmethod
+
+import tensorflow as tf
+
+
+class Augmentation:
+    @abstractmethod
+    def __call__(self, image: tf.Tensor) -> tf.Tensor:
+        """
+        Both the input and output should be between [0, 1]
+        """
+        raise NotImplementedError()

augmentations/auto_augment.py

@@ -0,0 +1,18 @@
+import tensorflow as tf
+from timm.data.auto_augment import rand_augment_transform
+
+from augmentations.augmentation import Augmentation
+
+
+class AutoAugment(Augmentation):
+    def __init__(self):
+        self.transform = rand_augment_transform(
+            config_str="rand-m9-mstd0.5-inc1", hparams={}
+        )
+
+    def __call__(self, image):
+        image_as_pil = tf.keras.preprocessing.image.array_to_img(image)
+        augmented_image = self.transform(image_as_pil)
+        return tf.convert_to_tensor(
+            tf.keras.preprocessing.image.img_to_array(augmented_image) / 255
+        )

augmentations/native_aspect_ratio_resize.py

@@ -0,0 +1,34 @@
+import tensorflow as tf
+
+from augmentations.augmentation import Augmentation
+
+
+class NativeAspectRatioResize(Augmentation):
+    def __init__(self, square_size, patch_size):
+        """
+        The image will be
+         1. Rescaled so that the area is smaller or equal square_size^2
+         2. Cropped so that the sides are multiples of patch_size
+        """
+        self.square_size = square_size
+        self.patch_size = patch_size
+
+    def __call__(self, image):
+        # extract the true height and width of the image (they are None when implicit)
+        height, width, _ = tf.shape(image)[0], tf.shape(image)[1], tf.shape(image)[2]
+        # compute the sqrt of the aspect ratio
+        sqrt_ratio = tf.cast(tf.sqrt(height / width), tf.float32)
+        # compute the new height and width
+        height = tf.cast(224 * sqrt_ratio, tf.int32)
+        width = tf.cast(224**2 / tf.cast(height, tf.float32), tf.int32)
+        # resize the image, now the num pixels is ~= 224^2
+        image = tf.image.resize(image, [height, width])
+
+        target_height = height - (height % self.patch_size)
+        target_width = width - (width % self.patch_size)
+        offset_height = (height - target_height) // 2
+        offset_width = (width - target_width) // 2
+        image = tf.image.crop_to_bounding_box(
+            image, offset_height, offset_width, target_height, target_width
+        )
+        return image

augmentations/random_horizontal_flip.py

@@ -0,0 +1,8 @@
+import tensorflow as tf
+
+from augmentations.augmentation import Augmentation
+
+
+class RandomHorizontalFlip(Augmentation):
+    def __call__(self, image):
+        return tf.image.random_flip_left_right(image)

augmentations/random_resized_crop.py

@@ -0,0 +1,80 @@
+import tensorflow as tf
+
+from augmentations.augmentation import Augmentation
+
+
+# Port of PyTorch's RandomResizedCrop
+# https://pytorch.org/vision/main/_modules/torchvision/transforms/transforms.html#RandomResizedCrop
+class RandomResizedCrop(Augmentation):
+    def __init__(self, size, scale, ratio):
+        self.size = tf.constant([size, size])
+        self.scale = tf.constant(scale)
+        self.ratio = tf.constant(ratio)
+        self.log_ratio = tf.math.log(ratio)
+
+    def get_params(self, image):
+        original_height, original_width, num_channels = image.shape
+        assert num_channels == 3
+        area = original_height * original_width
+
+        for _ in range(10):
+            target_area = area * tf.random.uniform(
+                shape=(),
+                minval=self.scale[0],
+                maxval=self.scale[1],
+            )
+            aspect_ratio = tf.math.exp(
+                tf.random.uniform(
+                    shape=(), minval=self.log_ratio[0], maxval=self.log_ratio[1]
+                )
+            )
+
+            new_width = tf.cast(
+                tf.math.round(tf.math.sqrt(target_area * aspect_ratio)), tf.int32
+            )
+            new_height = tf.cast(
+                tf.math.round(tf.math.sqrt(target_area / aspect_ratio)), tf.int32
+            )
+
+            if 0 < new_width <= original_width and 0 < new_height <= original_height:
+                i = tf.random.uniform(
+                    shape=(),
+                    minval=0,
+                    maxval=original_height - new_height + 1,
+                    dtype=tf.int32,
+                )
+                j = tf.random.uniform(
+                    shape=(),
+                    minval=0,
+                    maxval=original_width - new_width + 1,
+                    dtype=tf.int32,
+                )
+
+                return i, j, new_height, new_width
+
+        # Fallback to central crop
+        in_ratio = float(original_width) / float(original_height)
+        if in_ratio < tf.math.reduce_min(self.ratio):
+            new_width = original_width
+            new_height = tf.cast(
+                tf.math.round(new_width / tf.math.reduce_min(self.ratio)), tf.int32
+            )
+        elif in_ratio > tf.math.reduce_max(self.ratio):
+            new_height = original_height
+            new_width = tf.cast(
+                tf.math.round(new_height * tf.math.reduce_max(self.ratio)), tf.int32
+            )
+        else:  # whole image
+            new_width = original_width
+            new_height = original_height
+        i = (original_height - new_height) // 2
+        j = (original_width - new_width) // 2
+        return i, j, new_height, new_width
+
+    def __call__(self, image):
+        i, j, new_height, new_width = self.get_params(image)
+        crop = image[i : i + new_height, j : j + new_width]
+        resized_crop = tf.image.resize(
+            crop, self.size, method=tf.image.ResizeMethod.BICUBIC
+        )
+        return resized_crop

requirements.txt

@@ -99,7 +99,7 @@ termcolor==2.4.0
 toolz==0.12.1
 torch==2.3.0
 torchvision==0.18.0
-tornado==6.4
+tornado==6.4.1
 tqdm==4.66.2
 traitlets==5.14.3
 triton==2.3.0
@@ -109,4 +109,4 @@ wcwidth==0.2.13
 Werkzeug==3.0.2
 wrapt==1.16.0
 zipp==3.18.1
-tf-models-official==2.16.0
+timm==1.0.3

tests/augmentations/test_augmentation.py

@@ -0,0 +1,29 @@
+import pytest
+import tensorflow as tf
+
+from augmentations.auto_augment import AutoAugment
+from augmentations.native_aspect_ratio_resize import NativeAspectRatioResize
+from augmentations.random_horizontal_flip import RandomHorizontalFlip
+
+augmentations = [
+    AutoAugment(),
+    NativeAspectRatioResize(square_size=224, patch_size=14),
+    RandomHorizontalFlip(),
+    # BICUBIC interpolation can make the output be outside of [0, 1]
+    # RandomResizedCrop(size=224, scale=(0.4, 1.0), ratio=(0.75, 1.33)),
+]
+
+
+@pytest.mark.parametrize("augmentation", augmentations)
+def test_output_dtype_and_range(augmentation):
+    # The output should be a float32 in the range [0, 1]
+    input_image = tf.random.uniform(shape=(1000, 2000, 3), minval=0, maxval=1)
+    assert tf.math.reduce_min(input_image) >= 0
+    assert tf.math.reduce_max(input_image) <= 1
+    assert input_image.dtype == tf.float32
+
+    output_image = augmentation(input_image)
+
+    assert tf.math.reduce_min(output_image) >= 0
+    assert tf.math.reduce_max(output_image) <= 1
+    assert output_image.dtype == tf.float32

tests/augmentations/test_native_aspect_ratio_resize.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+from augmentations.native_aspect_ratio_resize import NativeAspectRatioResize
+
+
+def test_native_aspect_ratio_resize():
+    transformation = NativeAspectRatioResize(224, 14)
+    image = tf.random.uniform(shape=(1000, 3000, 3), minval=0, maxval=1)
+    new_image = transformation(image)
+
+    assert tf.math.reduce_prod(new_image.shape) <= 3 * 224**2
+    assert new_image.shape[0] % 14 == 0
+    assert new_image.shape[1] % 14 == 0

tests/augmentations/test_random_horizontal_flip.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+from augmentations.random_horizontal_flip import RandomHorizontalFlip
+
+
+def test_random_horizontal_flip():
+    image = tf.random.uniform(shape=(224, 224, 3), minval=0, maxval=1)
+    random_horizontal_flip = RandomHorizontalFlip()
+    new_image = random_horizontal_flip(image)
+
+    assert tf.reduce_all(new_image == tf.reverse(image, axis=[1])) or tf.reduce_all(
+        new_image == image
+    )

tests/augmentations/test_random_resized_crop.py

@@ -0,0 +1,16 @@
+import pytest
+import tensorflow as tf
+
+from augmentations.random_resized_crop import RandomResizedCrop
+
+
+@pytest.mark.parametrize("original_size", [(1000, 2000), (300, 100)])
+@pytest.mark.parametrize("size", [224, 448])
+@pytest.mark.parametrize("scale", [(0.4, 1.0), (0.1, 0.2), (2.0, 3.0)])
+@pytest.mark.parametrize("ratio", [(0.75, 1.33), (0.1, 0.2), (2.0, 3.0)])
+def test_random_resized_crop(original_size, size, scale, ratio):
+    original_image = tf.random.uniform(shape=(*original_size, 3), minval=0, maxval=1)
+    random_resized_crop = RandomResizedCrop(size, scale, ratio)
+    transformed_image = random_resized_crop(original_image)
+
+    assert transformed_image.shape == (size, size, 3)