_quickshift.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104
  1. import numpy as np
  2. from .._shared.filters import gaussian
  3. from .._shared.utils import _supported_float_type
  4. from ..color import rgb2lab
  5. from ..util import img_as_float
  6. from ._quickshift_cy import _quickshift_cython
  7. def quickshift(
  8. image,
  9. ratio=1.0,
  10. kernel_size=5,
  11. max_dist=10,
  12. return_tree=False,
  13. sigma=0,
  14. convert2lab=True,
  15. rng=42,
  16. *,
  17. channel_axis=-1,
  18. ):
  19. """Segment image using quickshift clustering in Color-(x,y) space.
  20. Produces an oversegmentation of the image using the quickshift mode-seeking
  21. algorithm.
  22. Parameters
  23. ----------
  24. image : (M, N, C) ndarray
  25. Input image. The axis corresponding to color channels can be specified
  26. via the `channel_axis` argument.
  27. ratio : float, optional, between 0 and 1
  28. Balances color-space proximity and image-space proximity.
  29. Higher values give more weight to color-space.
  30. kernel_size : float, optional
  31. Width of Gaussian kernel used in smoothing the
  32. sample density. Higher means fewer clusters.
  33. max_dist : float, optional
  34. Cut-off point for data distances.
  35. Higher means fewer clusters.
  36. return_tree : bool, optional
  37. Whether to return the full segmentation hierarchy tree and distances.
  38. sigma : float, optional
  39. Width for Gaussian smoothing as preprocessing. Zero means no smoothing.
  40. convert2lab : bool, optional
  41. Whether the input should be converted to Lab colorspace prior to
  42. segmentation. For this purpose, the input is assumed to be RGB.
  43. rng : {`numpy.random.Generator`, int}, optional
  44. Pseudo-random number generator.
  45. By default, a PCG64 generator is used (see :func:`numpy.random.default_rng`).
  46. If `rng` is an int, it is used to seed the generator.
  47. The PRNG is used to break ties, and is seeded with 42 by default.
  48. channel_axis : int, optional
  49. The axis of `image` corresponding to color channels. Defaults to the
  50. last axis.
  51. Returns
  52. -------
  53. segment_mask : (M, N) ndarray
  54. Integer mask indicating segment labels.
  55. Notes
  56. -----
  57. The authors advocate to convert the image to Lab color space prior to
  58. segmentation, though this is not strictly necessary. For this to work, the
  59. image must be given in RGB format.
  60. References
  61. ----------
  62. .. [1] Quick shift and kernel methods for mode seeking,
  63. Vedaldi, A. and Soatto, S.
  64. European Conference on Computer Vision, 2008
  65. """
  66. image = img_as_float(np.atleast_3d(image))
  67. float_dtype = _supported_float_type(image.dtype)
  68. image = image.astype(float_dtype, copy=False)
  69. if image.ndim > 3:
  70. raise ValueError("Only 2D color images are supported")
  71. # move channels to last position as expected by the Cython code
  72. image = np.moveaxis(image, source=channel_axis, destination=-1)
  73. if convert2lab:
  74. if image.shape[-1] != 3:
  75. raise ValueError("Only RGB images can be converted to Lab space.")
  76. image = rgb2lab(image)
  77. if kernel_size < 1:
  78. raise ValueError("`kernel_size` should be >= 1.")
  79. image = gaussian(image, sigma=[sigma, sigma, 0], mode='reflect', channel_axis=-1)
  80. image = np.ascontiguousarray(image * ratio)
  81. segment_mask = _quickshift_cython(
  82. image,
  83. kernel_size=kernel_size,
  84. max_dist=max_dist,
  85. return_tree=return_tree,
  86. rng=rng,
  87. )
  88. return segment_mask