Spaces:
Build error
Build error
File size: 3,187 Bytes
4bb934b |
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 78 79 80 81 82 83 84 |
import numpy as np
from .disjoint_set import DisjointSet
from .utils import smoothen, difference, get_random_rgb_image
def segment_graph(num_vertices, num_edges, edges, c):
edges[0 : num_edges, :] = edges[edges[0 : num_edges, 2].argsort()]
u = DisjointSet(num_vertices)
threshold = np.zeros(shape=num_vertices, dtype=float)
for i in range(num_vertices):
threshold[i] = c
for i in range(num_edges):
pedge = edges[i, :]
a = u.find(pedge[0])
b = u.find(pedge[1])
if a != b:
if (pedge[2] <= threshold[a]) and (pedge[2] <= threshold[b]):
u.join(a, b)
a = u.find(a)
threshold[a] = pedge[2] + (c / u.size(a))
return u
def segment(in_image, sigma, k, min_size):
height, width, band = in_image.shape
smooth_red_band = smoothen(in_image[:, :, 0], sigma)
smooth_green_band = smoothen(in_image[:, :, 1], sigma)
smooth_blue_band = smoothen(in_image[:, :, 2], sigma)
# build graph
edges_size = width * height * 4
edges = np.zeros(shape=(edges_size, 3), dtype=object)
num = 0
for y in range(height):
for x in range(width):
if x < width - 1:
edges[num, 0] = int(y * width + x)
edges[num, 1] = int(y * width + (x + 1))
edges[num, 2] = difference(
smooth_red_band, smooth_green_band,
smooth_blue_band, x, y, x + 1, y
)
num += 1
if y < height - 1:
edges[num, 0] = int(y * width + x)
edges[num, 1] = int((y + 1) * width + x)
edges[num, 2] = difference(
smooth_red_band, smooth_green_band,
smooth_blue_band, x, y, x, y + 1
)
num += 1
if (x < width - 1) and (y < height - 2):
edges[num, 0] = int(y * width + x)
edges[num, 1] = int((y + 1) * width + (x + 1))
edges[num, 2] = difference(
smooth_red_band, smooth_green_band,
smooth_blue_band, x, y, x + 1, y + 1
)
num += 1
if (x < width - 1) and (y > 0):
edges[num, 0] = int(y * width + x)
edges[num, 1] = int((y - 1) * width + (x + 1))
edges[num, 2] = difference(
smooth_red_band, smooth_green_band,
smooth_blue_band, x, y, x + 1, y - 1
)
num += 1
u = segment_graph(width * height, num, edges, k)
for i in range(num):
a = u.find(edges[i, 0])
b = u.find(edges[i, 1])
if (a != b) and ((u.size(a) < min_size) or (u.size(b) < min_size)):
u.join(a, b)
num_cc = u.num_sets()
output = np.zeros(shape=(height, width, 3))
colors = np.zeros(shape=(height * width, 3))
for i in range(height * width):
colors[i, :] = get_random_rgb_image()
for y in range(height):
for x in range(width):
comp = u.find(y * width + x)
output[y, x, :] = colors[comp, :]
return output
|