OpenCV分割烟丝、烟梗、背景

代码：

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage import morphology
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus']=False

def find_tobacoo_index(png_path):
    # 找出黑图中所有值为 1 的坐标点索引
    image_png = cv2.imread(png_path)
    one_png = image_png[:, :, 2]
    index_1 = np.where(one_png == 1)
    index_zip = [*zip(index_1[0], index_1[1])]

    return index_zip

def process_tobacco_stem(jpg_path,tobacoo_index):
    image_jpg = cv2.imread(jpg_path)
    yuv = cv2.cvtColor(image_jpg, cv2.COLOR_BGR2YUV)
    v = yuv[:, :, 2]
    v_binary = ~np.logical_not(v > 124)
    yuv_img_one = morphology.remove_small_objects(v_binary, 200, 1, True)
    v_binary1 = yuv_img_one.astype("int")
    # 循环过后单通道数组中背景为 0，梗为 1，丝为 2
    for i in tobacoo_index:
        v_binary1[i] = 2

    return v_binary1, image_jpg

def mix_tobacoo_stem(stem,image):
    result_image = np.zeros_like(image)
    result_image[:, :, 0] = stem
    result_image[:, :, 1] = stem
    result_image[:, :, 2] = stem

    return result_image

def show_image(image,times):
    one = image[:, :, 0]
    # 烟丝 --> 红色
    tobacco = (one == 2).astype("int")
    tobacco[tobacco == 1] = 255
    # 梗 --> 绿色
    stem = (one == 1).astype("int")
    stem[stem == 1] = 255
    # R 通道只有烟丝为 255，G 通道只有烟梗为 255，B 通道所有数字为 0
    image_show = np.zeros_like(image)
    image_show[:, :, 0] = tobacco
    image_show[:, :, 1] = stem
    unique = np.unique(image)
    plt.imshow(image_show)
    plt.title("第{}张 *** nuique：{}".format(times, unique))
    plt.show()


def main():
    png_path = "./0328_data/pngtoay"
    jpg_path = "./0328_data/jpjtoday"
    save = r"D:\pythonItems\convert_colors\0328_data\01png"
    png_dir = os.listdir(png_path)
    jpg_dir = os.listdir(jpg_path)

    for i in range(len(png_dir)):
        pngPath = png_path + "\\" + png_dir[i]
        jpgPath = jpg_path + "\\" + jpg_dir[i]
        index = find_tobacoo_index(pngPath)
        stem,image = process_tobacco_stem(jpgPath,index)
        save_image = mix_tobacoo_stem(stem,image)
        save_path = save + "\\" + png_dir[i]
        unique1 = np.unique(save_image)
        print("888",unique1)
        cv2.imwrite(save_path, save_image)
        print(f"{i} *** shape:{save_image.shape} *** uniqu:{np.unique(save_image)}")
        show_image(save_image,i + 1)


if __name__=="__main__":
    main()

原图：

标注烟丝的图片：

处理后：

背景为0烟丝烟梗为1烟丝为2