opencv实现超像素分割

实现效果图：

同时还使用了 mask图，要识别的区域为白色，背景为黑色

import cv2
import numpy as np
import os.path as osp
import os
import numpy as np
from tqdm import tqdm
from skimage import morphology
from skimage import segmentation


def get_ori_list(ori_folder):
    img_list = os.listdir(ori_folder)
    ori_list = []
    for img_name in img_list:
        flag = 0 
        for sample in check_list:
            if sample in img_name:
                flag=1
                break
        if flag==0:
            ori_list.append(osp.join(ori_folder,img_name))
        if len(ori_list)>20:
            break
    return ori_list

check_list = ['copper','bg','check','dust']
"""超像素由一系列位置相邻且颜色、亮度、纹理等特征相似的像素点组成的小区域。
这些小区域大多保留了进一步进行图像分割的有效信息，且一般不会破坏图像中物体的边界信息，
用少量的超像素代替大量像素表达图像特征，降低了图像处理的复杂度，
一般作为分割算法的预处理步骤。"""


def get_reverse(pic_matrix):
    where_0 = np.where(pic_matrix == 0)
    where_255 = np.where(pic_matrix == 255)
    
    pic_matrix[where_0] = 255
    pic_matrix[where_255] = 0

    return pic_matrix



def use_slic_by_SLIC(img_path,mask_path,end_path):
    """https://scikit-image.org/docs/dev/auto_examples/segmentation/plot_mask_slic.html#sphx-glr-auto-examples-segmentation-plot-mask-slic-py"""
    image = cv2.imread(img_path)
    #image = cv2.resize(image,(256,256))
    mask_image = cv2.imread(mask_path)
    #mask_image = get_reverse(mask_image)
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    # mask_image = cv2.dilate(mask_image, kernel)
    #mask_image = cv2.resize(mask_image,(256,256))
    img_R = mask_image[:,:,0]
    mask = img_R>220
    # 生成扁平的盘状结构元素 skimage.morphology.disk
    mask = morphology.opening(mask, morphology.disk(10))
    #mask是一个只包含True和False的ndarray，它的shape和data一致，
    # segmentation.slic在Color-（x，y，z）空间中使用k-means聚类来分割图像。
    m_slic = segmentation.slic(image, n_segments=20000, mask=mask)
    slic_image = segmentation.mark_boundaries(image, m_slic,outline_color=(0,1,1))
    cv2.imwrite(end_path,slic_image*255)  

    return slic_image*255






if __name__ == '__main__':
    
    ori_folder = '/cloud_disk/users/huh/dataset/PCB/ddrnet_23_slim/pre_process_img'
    end_folder = '/cloud_disk/users/huh/pcb/script/slic_result/1_cv2'
    ori_list = get_ori_list(ori_folder)
    for img_path in tqdm(ori_list):
        end_path = osp.join(end_folder,osp.basename(img_path))
        mask_path = osp.join(ori_folder,osp.basename(img_path[:-4])+'_bg_mask.jpg')
        try:
            slic_image = use_slic_by_SLIC(img_path,mask_path,end_path)
            # slic_image = cv2.imread(osp.join('/cloud_disk/users/huh/pcb/script/slic_result/1_without_small',osp.basename(img_path)))
            #slic_image = cv2.resize(slic_image,(512,512))
            #remove_small_objects(mask_path,end_path,img_path,slic_image)
        except ValueError:
            print(end_path)