/
_felzenszwalb.py
77 lines (64 loc) · 2.98 KB
/
_felzenszwalb.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import numpy as np
from .._shared.utils import warn
from ._felzenszwalb_cy import _felzenszwalb_grey
def felzenszwalb(image, scale=1, sigma=0.8, min_size=20):
"""Computes Felsenszwalb's efficient graph based image segmentation.
Produces an oversegmentation of a multichannel (i.e. RGB) image
using a fast, minimum spanning tree based clustering on the image grid.
The parameter ``scale`` sets an observation level. Higher scale means
less and larger segments. ``sigma`` is the diameter of a Gaussian kernel,
used for smoothing the image prior to segmentation.
The number of produced segments as well as their size can only be
controlled indirectly through ``scale``. Segment size within an image can
vary greatly depending on local contrast.
For RGB images, the algorithm computes a separate segmentation for each
channel and then combines these. The combined segmentation is the
intersection of the separate segmentations on the color channels.
Parameters
----------
image : (width, height, 3) or (width, height) ndarray
Input image.
scale : float
Free parameter. Higher means larger clusters.
sigma : float
Width of Gaussian kernel used in preprocessing.
min_size : int
Minimum component size. Enforced using postprocessing.
Returns
-------
segment_mask : (width, height) ndarray
Integer mask indicating segment labels.
References
----------
.. [1] Efficient graph-based image segmentation, Felzenszwalb, P.F. and
Huttenlocher, D.P. International Journal of Computer Vision, 2004
"""
if image.ndim == 2:
# assume single channel image
return _felzenszwalb_grey(image, scale=scale, sigma=sigma,
min_size=min_size)
elif image.ndim != 3:
raise ValueError("Felzenswalb segmentation can only operate on RGB and"
" grey images, but input array of ndim %d given."
% image.ndim)
# assume we got 2d image with multiple channels
n_channels = image.shape[2]
if n_channels != 3:
warn("Got image with %d channels. Is that really what you"
" wanted?" % image.shape[2])
segmentations = []
# compute quickshift for each channel
for c in range(n_channels):
channel = np.ascontiguousarray(image[:, :, c])
s = _felzenszwalb_grey(channel, scale=scale, sigma=sigma,
min_size=min_size)
segmentations.append(s)
# put pixels in same segment only if in the same segment in all images
# we do this by combining the channels to one number
n0 = segmentations[0].max() + 1
n1 = segmentations[1].max() + 1
segmentation = (segmentations[0] + segmentations[1] * n0
+ segmentations[2] * n0 * n1)
# make segment labels consecutive numbers starting at 0
labels = np.unique(segmentation, return_inverse=True)[1]
return labels.reshape(image.shape[:2])