Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow_addons/image/distance_transform.py: 47%

1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Distance transform ops.""" 

16 

17import tensorflow as tf 

18from tensorflow_addons.image import utils as img_utils 

19from tensorflow_addons.utils.resource_loader import LazySO 

20from tensorflow_addons.utils.types import TensorLike 

21 

22from typing import Optional, Type 

23 

24_image_so = LazySO("custom_ops/image/_image_ops.so") 

25 

26tf.no_gradient("Addons>EuclideanDistanceTransform") 

27 

28 

29def euclidean_dist_transform( 

30 images: TensorLike, 

31 dtype: Type[tf.dtypes.DType] = tf.float32, 

32 name: Optional[str] = None, 

33) -> tf.Tensor: 

34 """Applies euclidean distance transform(s) to the image(s). 

35 

36 Based on [Distance Transforms of Sampled Functions] 

37 (http://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf). 

38 

39 Args: 

40 images: A tensor of shape `(num_images, num_rows, num_columns, num_channels)` 

41 (NHWC), or `(num_rows, num_columns, num_channels)` (HWC) or 

42 `(num_rows, num_columns)` (HW). 

43 dtype: `tf.dtypes.DType` of the output tensor. 

44 name: The name of the op. 

45 

46 Returns: 

47 Image(s) with the type `dtype` and same shape as `images`, with the 

48 transform applied. If a tensor of all ones is given as input, the 

49 output tensor will be filled with the max value of the `dtype`. 

50 

51 Raises: 

52 TypeError: If `image` is not tf.uint8, or `dtype` is not floating point. 

53 """ 

54 

55 with tf.name_scope(name or "euclidean_distance_transform"): 

56 image_or_images = tf.convert_to_tensor(images, name="images") 

57 

58 if image_or_images.dtype.base_dtype != tf.uint8: 

59 raise TypeError("Invalid dtype %s. Expected uint8." % image_or_images.dtype) 

60 

61 images = img_utils.to_4D_image(image_or_images) 

62 original_ndims = img_utils.get_ndims(image_or_images) 

63 

64 if dtype not in [tf.float16, tf.float32, tf.float64]: 

65 raise TypeError("`dtype` must be float16, float32 or float64") 

66 

67 output = _image_so.ops.addons_euclidean_distance_transform(images, dtype) 

68 

69 return img_utils.from_4D_image(output, original_ndims)