# Image segmentation based on the paper:
# Felzenszwalb, Pedro F.; Huttenlocher, Daniel P.: "Efficient graph-based image segmentation", International Journal of Computer Vision, 59 (2): 167–181, 2004.


import matplotlib.pyplot as plt
import numpy as np

from skimage.data import astronaut
from skimage.segmentation import felzenszwalb
from skimage.segmentation import mark_boundaries
from skimage.util import img_as_float
import skimage.io

#img = img_as_float(astronaut()[::2, ::2])
img = skimage.io.imread('lighthouse512.png')

segments_fz = felzenszwalb(img, scale=100, sigma=0.5, min_size=50)

print(f'Number of segments: {len(np.unique(segments_fz))}')

fig = plt.figure(figsize=(10, 10))

plt.imshow(mark_boundaries(img, segments_fz))

plt.tight_layout()
plt.show()

Number of segments: 577

# show the segments using different, random colors 

# set seed to get repeatable results
np.random.seed(33)

gray_img = img.mean(axis=2) 
labels = segments_fz

# get unique components
unique_labels = np.unique(labels)

# assign random colors (RGB in [0,1])
rand_colors = np.random.rand(len(unique_labels), 3)

# create color image
h, w = gray_img.shape
color_img = np.zeros((h, w, 3))

for i, lbl in enumerate(unique_labels):
    mask = (labels == lbl)
    for c in range(3):
        color_img[:, :, c][mask] = rand_colors[i, c] * (gray_img[mask] / 255.0)

plt.figure(figsize=(10,10))
plt.imshow(color_img)
#plt.imsave('lighthouse_seg1.png', color_img )
plt.show()