mirror of
https://github.com/wassname/scikit-image.git
synced 2026-07-22 13:00:09 +08:00
81 lines
2.8 KiB
Cython
81 lines
2.8 KiB
Cython
import numpy as np
|
|
cimport numpy as np
|
|
|
|
from itertools import product
|
|
|
|
|
|
def quickshift(np.ndarray[dtype=np.float_t, ndim=3, mode="c"] image, sigma=5, tau=10):
|
|
"""Computes quickshift clustering in RGB-(x,y) space.
|
|
|
|
Parameters
|
|
----------
|
|
image: ndarray, [width, height, channels]
|
|
Input image
|
|
sigma: float
|
|
Width of Gaussian kernel used in smoothing the
|
|
sample density. Higher means less clusters.
|
|
tau: float
|
|
Cut-off point for data distances.
|
|
Higher means less clusters.
|
|
|
|
Returns
|
|
-------
|
|
segment_mask: ndarray, [width, height]
|
|
Integer mask indicating segment labels.
|
|
"""
|
|
|
|
# We compute the distances twice since otherwise
|
|
# we might get crazy memory overhead (width * height * windowsize**2)
|
|
|
|
# TODO do smoothing beforehand?
|
|
# TODO manage borders somehow?
|
|
|
|
# window size for neighboring pixels to consider
|
|
if sigma < 1:
|
|
raise ValueError("Sigma should be >= 1")
|
|
cdef int w = int(2 * sigma)
|
|
|
|
cdef int width = image.shape[0]
|
|
cdef int height = image.shape[1]
|
|
|
|
cdef np.ndarray[dtype=np.float_t, ndim=2] densities = np.zeros((width, height))
|
|
|
|
# compute densities
|
|
for x, y in product(xrange(width), xrange(height)):
|
|
current_pixel = image[x, y, :]
|
|
for xx, yy in product(xrange(-w / 2, w / 2 + 1), repeat=2):
|
|
x_, y_ = x + xx, y + yy
|
|
if 0 <= x_ < width and 0 <= y_ < height:
|
|
dist = np.sum((current_pixel - image[x_, y_, :])**2) + (x - x_)**2 + (y - y_)**2
|
|
densities[x, y] += np.exp(-dist / sigma)
|
|
|
|
# this will break ties that otherwise would give us headache
|
|
|
|
densities += np.random.normal(scale=0.00001, size=(width, height))
|
|
# default parent to self:
|
|
parent = np.arange(width * height).reshape(width, height)
|
|
dist_parent = np.zeros((width, height))
|
|
# find nearest node with higher density
|
|
for x, y in product(xrange(width), xrange(height)):
|
|
current_density = densities[x, y]
|
|
current_pixel = image[x, y, :]
|
|
closest = np.inf
|
|
for xx, yy in product(xrange(-w / 2, w / 2 + 1), repeat=2):
|
|
x_, y_ = x + xx, y + yy
|
|
if 0 <= x_ < width and 0 <= y_ < height:
|
|
if densities[x_, y_] > current_density:
|
|
dist = np.sum((current_pixel - image[x_, y_, :])**2) + (x - x_)**2 + (y - y_)**2
|
|
if dist < closest:
|
|
closest = dist
|
|
parent[x, y] = x_ * width + y_
|
|
dist_parent[x, y] = closest
|
|
|
|
dist_parent = dist_parent.ravel()
|
|
flat = parent.ravel()
|
|
flat[dist_parent > tau] = np.arange(width * height)[dist_parent > tau]
|
|
old = np.zeros_like(flat)
|
|
while (old != flat).any():
|
|
old = flat
|
|
flat = flat[flat]
|
|
return flat.reshape(parent.shape)
|