# -- encoding:utf-8 --

import numpy as np
import cv2 as cv
from skimage import util
from skimage import segmentation
import matplotlib.pyplot as plt
np.set_printoptions(threshold=np.inf)


def _generate_segments(im_orig, scale, sigma, min_size):
    """
        segmenmt smallest regions by the algorithm of Felzenswalb and
        Huttenlocher
    """

    # open the Image
    """
    segmentation.felzenszwalb() 参数：
        image,  输入对象
        scale=1,  如果设置比较大 ---> 少且大的分割
        sigma=0.8,  （高斯核宽度的标准差），用于分割之前的图片平滑的预处理。
        min_size：一般用于限制区域框的面积大小。
    """

    im_mask = segmentation.felzenszwalb(util.img_as_float(im_orig),
                                        scale=scale, sigma=sigma,
                                        min_size=min_size)

    return im_orig, im_mask


if __name__ == '__main__':
    # path = "./images/000009.jpg"
    path = "./images/small.png"
    img = cv.imread(path)
    img = cv.cvtColor(img, cv.COLOR_BGR2RGB)

    # 进行图像分割
    old_img, img_mask = _generate_segments(img, scale=50, sigma=2, min_size=10)
    print(np.mean(img == old_img))
    print(img_mask)

    # 展示一下
    img_mask_max = np.max(img_mask)
    img_mask_min = np.min(img_mask)
    img_mask = ((img_mask - img_mask_min) / (img_mask_max - img_mask_min) * 255).astype(np.uint8)
    plt.imshow(img_mask)
    plt.show()

    # cv.imshow('img_mask', img_mask)
    # cv.waitKey(0)
    # cv.destroyAllWindows()