Skip to content

Commit

Permalink
remove tfm dependency
Browse files Browse the repository at this point in the history
  • Loading branch information
rvankoert committed Jun 17, 2024
1 parent 60ee295 commit 216f883
Show file tree
Hide file tree
Showing 3 changed files with 312 additions and 6 deletions.
1 change: 0 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -11,5 +11,4 @@ blinker==1.7.0
fpdf==1.7.2
scikit-image==0.22.0
prometheus-client==0.20.0
tf-models-official==2.14.1
xlsxwriter==3.2.0
9 changes: 4 additions & 5 deletions src/data/augment_layers.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,7 @@
import numpy as np
from skimage.filters import threshold_otsu, threshold_sauvola
import tensorflow as tf
import tensorflow_models as tfm

import gaussian_filter2d

class ShearXLayer(tf.keras.layers.Layer):
def __init__(self, binary=False, **kwargs):
Expand Down Expand Up @@ -547,9 +546,9 @@ def call(self, inputs, training=None):
blur_factor = 1
else:
blur_factor = round(random.uniform(0.1, 2), 1)
return tfm.vision.augment.gaussian_filter2d(inputs,
filter_shape=(11, 11),
sigma=blur_factor)
return gaussian_filter2d(inputs,
filter_shape=(11, 11),
sigma=blur_factor)


class RandomWidthLayer(tf.keras.layers.Layer):
Expand Down
308 changes: 308 additions & 0 deletions src/data/gaussian_filter2d.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,308 @@
# Copyright 2019 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.
# ==============================================================================

### COPIED ####
### FROM ###
### https://github.com/tensorflow/addons/blob/v0.20.0/tensorflow_addons/image/filters.py#L218-L302 ###
### DO NOT REMOVE THIS NOTICE ###


import tensorflow as tf
from tensorflow_addons.image import utils as img_utils
from tensorflow_addons.utils import keras_utils
from tensorflow_addons.utils.types import TensorLike

from typing import Optional, Union, List, Tuple, Iterable


def _pad(
image: TensorLike,
filter_shape: Union[List[int], Tuple[int]],
mode: str = "CONSTANT",
constant_values: TensorLike = 0,
) -> tf.Tensor:
"""Explicitly pad a 4-D image.
Equivalent to the implicit padding method offered in `tf.nn.conv2d` and
`tf.nn.depthwise_conv2d`, but supports non-zero, reflect and symmetric
padding mode. For the even-sized filter, it pads one more value to the
right or the bottom side.
Args:
image: A 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
filter_shape: A `tuple`/`list` of 2 integers, specifying the height
and width of the 2-D filter.
mode: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
The type of padding algorithm to use, which is compatible with
`mode` argument in `tf.pad`. For more details, please refer to
https://www.tensorflow.org/api_docs/python/tf/pad.
constant_values: A `scalar`, the pad value to use in "CONSTANT"
padding mode.
"""
if mode.upper() not in {"REFLECT", "CONSTANT", "SYMMETRIC"}:
raise ValueError(
'padding should be one of "REFLECT", "CONSTANT", or "SYMMETRIC".'
)
constant_values = tf.convert_to_tensor(constant_values, image.dtype)
filter_height, filter_width = filter_shape
pad_top = (filter_height - 1) // 2
pad_bottom = filter_height - 1 - pad_top
pad_left = (filter_width - 1) // 2
pad_right = filter_width - 1 - pad_left
paddings = [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
return tf.pad(image, paddings, mode=mode, constant_values=constant_values)


@tf.function
def mean_filter2d(
image: TensorLike,
filter_shape: Union[int, Iterable[int]] = (3, 3),
padding: str = "REFLECT",
constant_values: TensorLike = 0,
name: Optional[str] = None,
) -> tf.Tensor:
"""Perform mean filtering on image(s).
Args:
image: Either a 2-D `Tensor` of shape `[height, width]`,
a 3-D `Tensor` of shape `[height, width, channels]`,
or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
the height and width of the 2-D mean filter. Can be a single integer
to specify the same value for all spatial dimensions.
padding: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
The type of padding algorithm to use, which is compatible with
`mode` argument in `tf.pad`. For more details, please refer to
https://www.tensorflow.org/api_docs/python/tf/pad.
constant_values: A `scalar`, the pad value to use in "CONSTANT"
padding mode.
name: A name for this operation (optional).
Returns:
2-D, 3-D or 4-D `Tensor` of the same dtype as input.
Raises:
ValueError: If `image` is not 2, 3 or 4-dimensional,
if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
or if `filter_shape` is invalid.
"""
with tf.name_scope(name or "mean_filter2d"):
image = tf.convert_to_tensor(image, name="image")
original_ndims = img_utils.get_ndims(image)
image = img_utils.to_4D_image(image)

filter_shape = keras_utils.normalize_tuple(filter_shape, 2, "filter_shape")

# Keep the precision if it's float;
# otherwise, convert to float32 for computing.
orig_dtype = image.dtype
if not image.dtype.is_floating:
image = tf.dtypes.cast(image, tf.dtypes.float32)

# Explicitly pad the image
image = _pad(image, filter_shape, mode=padding, constant_values=constant_values)

# Filter of shape (filter_width, filter_height, in_channels, 1)
# has the value of 1 for each element.
area = tf.constant(filter_shape[0] * filter_shape[1], dtype=image.dtype)
filter_shape += (tf.shape(image)[-1], 1)
kernel = tf.ones(shape=filter_shape, dtype=image.dtype)

output = tf.nn.depthwise_conv2d(
image, kernel, strides=(1, 1, 1, 1), padding="VALID"
)

output /= area

output = img_utils.from_4D_image(output, original_ndims)
return tf.dtypes.cast(output, orig_dtype)


@tf.function
def median_filter2d(
image: TensorLike,
filter_shape: Union[int, Iterable[int]] = (3, 3),
padding: str = "REFLECT",
constant_values: TensorLike = 0,
name: Optional[str] = None,
) -> tf.Tensor:
"""Perform median filtering on image(s).
Args:
image: Either a 2-D `Tensor` of shape `[height, width]`,
a 3-D `Tensor` of shape `[height, width, channels]`,
or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
the height and width of the 2-D median filter. Can be a single integer
to specify the same value for all spatial dimensions.
padding: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
The type of padding algorithm to use, which is compatible with
`mode` argument in `tf.pad`. For more details, please refer to
https://www.tensorflow.org/api_docs/python/tf/pad.
constant_values: A `scalar`, the pad value to use in "CONSTANT"
padding mode.
name: A name for this operation (optional).
Returns:
2-D, 3-D or 4-D `Tensor` of the same dtype as input.
Raises:
ValueError: If `image` is not 2, 3 or 4-dimensional,
if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
or if `filter_shape` is invalid.
"""
with tf.name_scope(name or "median_filter2d"):
image = tf.convert_to_tensor(image, name="image")
original_ndims = img_utils.get_ndims(image)
image = img_utils.to_4D_image(image)

filter_shape = keras_utils.normalize_tuple(filter_shape, 2, "filter_shape")

image_shape = tf.shape(image)
batch_size = image_shape[0]
height = image_shape[1]
width = image_shape[2]
channels = image_shape[3]

# Explicitly pad the image
image = _pad(image, filter_shape, mode=padding, constant_values=constant_values)

area = filter_shape[0] * filter_shape[1]

floor = (area + 1) // 2
ceil = area // 2 + 1

patches = tf.image.extract_patches(
image,
sizes=[1, filter_shape[0], filter_shape[1], 1],
strides=[1, 1, 1, 1],
rates=[1, 1, 1, 1],
padding="VALID",
)

patches = tf.reshape(patches, shape=[batch_size, height, width, area, channels])

patches = tf.transpose(patches, [0, 1, 2, 4, 3])

# Note the returned median is casted back to the original type
# Take [5, 6, 7, 8] for example, the median is (6 + 7) / 2 = 3.5
# It turns out to be int(6.5) = 6 if the original type is int
top = tf.nn.top_k(patches, k=ceil).values
if area % 2 == 1:
median = top[:, :, :, :, floor - 1]
else:
median = (top[:, :, :, :, floor - 1] + top[:, :, :, :, ceil - 1]) / 2

output = tf.cast(median, image.dtype)
output = img_utils.from_4D_image(output, original_ndims)
return output


def _get_gaussian_kernel(sigma, filter_shape):
"""Compute 1D Gaussian kernel."""
sigma = tf.convert_to_tensor(sigma)
x = tf.range(-filter_shape // 2 + 1, filter_shape // 2 + 1)
x = tf.cast(x**2, sigma.dtype)
x = tf.nn.softmax(-x / (2.0 * (sigma**2)))
return x


def _get_gaussian_kernel_2d(gaussian_filter_x, gaussian_filter_y):
"""Compute 2D Gaussian kernel given 1D kernels."""
gaussian_kernel = tf.matmul(gaussian_filter_x, gaussian_filter_y)
return gaussian_kernel


@tf.function
def gaussian_filter2d(
image: TensorLike,
filter_shape: Union[int, Iterable[int]] = (3, 3),
sigma: Union[List[float], Tuple[float], float] = 1.0,
padding: str = "REFLECT",
constant_values: TensorLike = 0,
name: Optional[str] = None,
) -> TensorLike:
"""Perform Gaussian blur on image(s).
Args:
image: Either a 2-D `Tensor` of shape `[height, width]`,
a 3-D `Tensor` of shape `[height, width, channels]`,
or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
the height and width of the 2-D gaussian filter. Can be a single
integer to specify the same value for all spatial dimensions.
sigma: A `float` or `tuple`/`list` of 2 floats, specifying
the standard deviation in x and y direction the 2-D gaussian filter.
Can be a single float to specify the same value for all spatial
dimensions.
padding: A `string`, one of "REFLECT", "CONSTANT", or "SYMMETRIC".
The type of padding algorithm to use, which is compatible with
`mode` argument in `tf.pad`. For more details, please refer to
https://www.tensorflow.org/api_docs/python/tf/pad.
constant_values: A `scalar`, the pad value to use in "CONSTANT"
padding mode.
name: A name for this operation (optional).
Returns:
2-D, 3-D or 4-D `Tensor` of the same dtype as input.
Raises:
ValueError: If `image` is not 2, 3 or 4-dimensional,
if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
if `filter_shape` is invalid,
or if `sigma` is invalid.
"""
with tf.name_scope(name or "gaussian_filter2d"):
if isinstance(sigma, (list, tuple)):
if len(sigma) != 2:
raise ValueError("sigma should be a float or a tuple/list of 2 floats")
else:
sigma = (sigma,) * 2

if any(s < 0 for s in sigma):
raise ValueError("sigma should be greater than or equal to 0.")

image = tf.convert_to_tensor(image, name="image")
sigma = tf.convert_to_tensor(sigma, name="sigma")

original_ndims = img_utils.get_ndims(image)
image = img_utils.to_4D_image(image)

# Keep the precision if it's float;
# otherwise, convert to float32 for computing.
orig_dtype = image.dtype
if not image.dtype.is_floating:
image = tf.cast(image, tf.float32)

channels = tf.shape(image)[3]
filter_shape = keras_utils.normalize_tuple(filter_shape, 2, "filter_shape")

sigma = tf.cast(sigma, image.dtype)
gaussian_kernel_x = _get_gaussian_kernel(sigma[1], filter_shape[1])
gaussian_kernel_x = gaussian_kernel_x[tf.newaxis, :]

gaussian_kernel_y = _get_gaussian_kernel(sigma[0], filter_shape[0])
gaussian_kernel_y = gaussian_kernel_y[:, tf.newaxis]

gaussian_kernel_2d = _get_gaussian_kernel_2d(
gaussian_kernel_y, gaussian_kernel_x
)
gaussian_kernel_2d = gaussian_kernel_2d[:, :, tf.newaxis, tf.newaxis]
gaussian_kernel_2d = tf.tile(gaussian_kernel_2d, [1, 1, channels, 1])

image = _pad(image, filter_shape, mode=padding, constant_values=constant_values)

output = tf.nn.depthwise_conv2d(
input=image,
filter=gaussian_kernel_2d,
strides=(1, 1, 1, 1),
padding="VALID",
)
output = img_utils.from_4D_image(output, original_ndims)
return tf.cast(output, orig_dtype)

0 comments on commit 216f883

Please sign in to comment.