opencv mask后处理方法

对图像选定值进行膨胀处理

import cv2
import os
import shutil

import numpy as np

src = r'D:\Projects\Dataset\gray_mask'
dst = r'D:\Projects\Dataset\erode_mask'

if os.path.exists(dst):
    shutil.rmtree(dst)
os.mkdir(dst)


def IsConnectedArea(mask):
    num_labels, labels = cv2.connectedComponents(mask, connectivity=8)
    if num_labels == 2:
        return True
    return False


def RemoveSmallConnectedArea(mask):
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)
    if num_labels < 2:
        return mask
    id = -1
    max = -1
    for i in range(1, num_labels):
        if stats[i][4] > max:
            id = i
            max = stats[i][4]
    # print(max / mask.size)
    if max / mask.size < 0.08:
        return mask * 0
    return ((labels == id) * mask * 1.5).astype(np.uint8)


show = False

for i, path in enumerate(os.listdir(src)):
    print(i,path)
    mask_path = os.path.join(src, path)
    image_path = mask_path.replace('gray_mask', 'image')

    mask = cv2.imread(mask_path)
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    image = cv2.imread(image_path)

    mask = np.array(mask)
    score = np.sum(mask == 129) / np.sum(mask > -1)
    if 0.25 < score < 0.5:
        mask_129= mask * (mask == 129)
        mask_255= mask * (mask == 255)
        if IsConnectedArea(mask):
            mask_129_erode = cv2.dilate(gums, cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15)), iterations=5)
            mask_129_erode = cv2.add(gums_erode, mask_255)
            # cv2.imshow('mask_129_erode', mask_129_erode)
            dilate_boundary = cv2.subtract(gums_erode, mask)

            dilate_boundary2 = RemoveSmallConnectedArea(dilate_boundary)
            # cv2.imshow('dilate_boundary2', dilate_boundary2)
            mask2 = cv2.add(dilate_boundary2, mask)

            mask2 = cv2.cvtColor(mask2, cv2.COLOR_GRAY2BGR)
            image2 = image * (mask2 > 0)
            # score=np.sum(dilate_boundary2 >0) / np.sum(dilate_boundary2 > -1)
            # print(score)
            if np.sum(dilate_boundary2) > 0:
                shutil.copy(image_path, os.path.join(dst, path.split('.')[0] + "[0].png"))
                shutil.copy(mask_path, os.path.join(dst, path.split('.')[0] + "[1].png"))
                cv2.imwrite(os.path.join(dst, path.split('.')[0] + "[2].png"), dilate_boundary)
                cv2.imwrite(os.path.join(dst, path.split('.')[0] + "[3].png"), mask2)
                cv2.imwrite(os.path.join(dst, path.split('.')[0] + "[4].png"), image2)
            if show:
                cv2.imshow('mask', mask)
                cv2.imshow('image', image)
                cv2.imshow('dilate_boundary', dilate_boundary)
                cv2.imshow('mask2', mask2)
                cv2.imshow('image2', image2)
                cv2.waitKey(0)

你可能感兴趣的:(opencv,人工智能,计算机视觉)