OpenCV python GrabCut算法提取前景背景图片

OpenCV python GrabCut算法提取前景背景图片

处理图片:[lena_color.jpg]

手动mask[mask.jpg]

import numpy as np
import cv2
import matplotlib.pyplot as plt


def main():

    # 1.导入图片
    img_src = cv2.imread("lena_color.jpg")
    img_rgb = cv2.cvtColor(img_src, cv2.COLOR_BGR2RGB)

    # 2.创建掩模图片
    mask = np.zeros(img_src.shape[:2], np.uint8)
    bg_model = np.zeros((1, 65), np.float64)
    fg_model = np.zeros((1, 65), np.float64)

    # 3.执行区域 前景背景分割
    rect = (50, 50, 400, 500)
    cv2.grabCut(img_src, mask, rect, bg_model, fg_model, 5, cv2.GC_INIT_WITH_RECT)

    # 4.导入手动掩模图片
    mask2 = cv2.imread("mask.jpg", 0)
    mask2_show = cv2.imread("mask.jpg", -1)
    m2_rgb = cv2.cvtColor(mask2_show, cv2.COLOR_BGR2RGB)

    # 5.执行mask 前景背景分割
    mask[mask2 == 0] = 0
    mask[mask2 == 255] = 1
    mask, bg_model, fg_model = cv2.grabCut(img_src, mask, None, bg_model, fg_model, 5, cv2.GC_INIT_WITH_MASK)

    mask = np.where((mask == 2) | (mask == 0), 0, 1).astype("uint8")

    # 6.分离前景背景图片
    img_gc = img_src * mask[:, :, np.newaxis]
    img_gc = cv2.cvtColor(img_gc, cv2.COLOR_BGR2RGB)

    # 7.显示结果
    plt.figure("显示结果", figsize=(12, 6))
    plt.subplot(131)
    plt.imshow(m2_rgb)
    plt.axis("off")

    plt.subplot(132)
    plt.imshow(img_gc)
    plt.axis("off")

    plt.subplot(133)
    plt.imshow(mask)
    plt.axis("off")

    plt.show()


if __name__ == '__main__':
    main()

处理结果图片