watershed(分水岭算法)

为了检测出图中轮廓：

 算法效果比起slic差不少

借鉴文章，这个讲的很不错

OpenCv学习笔记5--图像分割之分水岭算法_wx612f24149b7fc的技术博客_51CTO博客https://blog.51cto.com/u_15351425/3725366 

import cv2
import numpy as np
import os
import os.path as osp
from tqdm import tqdm
check_list = ['copper','bg','check','dust']
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

'''
            分水岭算法 图像分割 
                
            分水岭算法实现图像自动分割的步骤：

                1 图像灰度化、Canny边缘检测
                2 查找轮廓，并且把轮廓信息按照不同的编号绘制到watershed的第二个参数markers上，相当于标记注水点。
                3 watershed分水岭算法
                4 绘制分割出来的区域，然后使用随机颜色填充，再跟源图像融合，以得到更好的显示效果。
                cv2.watershed(image, markers) → None
'''
def watershed(img_path,end_path):
    # 读入图片
    img = cv2.imread(img_path)
    # 32位有符号整数类型
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # canny边缘检测 函数返回一副二值图，其中包含检测出的边缘。
    canny = cv2.Canny(gray_img, 20, 150)
    contours, hierarchy = cv2.findContours(canny, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    marks = np.zeros(img.shape[:2], np.int32)
    imageContours = np.zeros(img.shape[:2], np.uint8)
    # 轮廓颜色
    compCount = 0
    index = 0
    for index in range(len(contours)):

        # 对marks进行标记，对不同区域的轮廓使用不同的亮度绘制，相当于设置注水点，有多少个轮廓，就有多少个注水点
        # 图像上不同线条的灰度值是不同的，底部略暗，越往上灰度越高
        marks = cv2.drawContours(marks, contours, index, (index, index, index), 1, 8, hierarchy)
    markerShows = cv2.convertScaleAbs(marks)
    # 使用分水岭算法
    marks = cv2.watershed(img, marks)
    afterWatershed = cv2.convertScaleAbs(marks)
    ###分水岭算法之后,让水漫起来,并且把堤坝即分水岭绘制为绿色
    img[marks == -1] = [ 0, 255, 0]
    
    cv2.imwrite(end_path,img)
 

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/watershed'
    ori_list = get_ori_list(ori_folder)
    
    for pic_name in tqdm(ori_list):
        end_path = osp.join(end_folder,osp.basename(pic_name))
        img_path = osp.join(ori_folder,osp.basename(pic_name))  
        watershed(img_path,end_path)