toys/py-replace-bg/main.py

#!/usr/bin/env python3
# -*- coding: utf-8; mode: python; tab-width: 4; indent-tabs-mode: nil -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
#           DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
#                   Version 2, December 2004
#
# Copyright (C) 2004 Sam Hocevar <sam@hocevar.net>
#
# Everyone is permitted to copy and distribute verbatim or modified
# copies of this license document, and changing it is allowed as long
# as the name is changed.
#
#           DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE
#  TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION
#
# 0. You just DO WHAT THE FUCK YOU WANT TO.
#
# Copyright (c) 2020- donkey <anjingyu_ws@foxmail.com>

import sys
import os

import cv2
import mediapipe as mp
import numpy as np
from tqdm import tqdm
import time
import matplotlib.pyplot as plt

def look_img(img):
    img_rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    plt.imshow(img_rgb)
    plt.show()

def show_test(file_path: str):
    img = cv2.imread(file_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=0)
    results = seg.process(img_rgb)
    mask = results.segmentation_mask

    plt.imshow(mask)
    plt.show()

    mask = mask > 0.5
    plt.imshow(mask)
    plt.show()

def show_test2(file_path):
    mp_drawing = mp.solutions.drawing_utils
    mp_selfie_segmentation = mp.solutions.selfie_segmentation
    IMAGE_FILES = []
    BG_COLOR = (0, 255, 0)
    MASK_COLOR = (255, 255, 255)
    with mp_selfie_segmentation.SelfieSegmentation(model_selection=0) as selfie_segmentation:
        image = cv2.imread(file_path)
        image_height, image_width, _ = image.shape
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = selfie_segmentation.process(image)
        # 在背景图像上绘制分割图
        # 为了改善边界周围的分割，可以考虑在 results.segmentation_mask进行双边过滤
        condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.5
        #生成纯色图像,白色的mask图纸
        #fg_image = np.zeros(image.shape, dtype=np.uint8)
        #fg_image[:] = MASK_COLOR
        fg_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        bg_image = np.zeros(image.shape, dtype=np.uint8)
        bg_image[:] = MASK_COLOR
        output_image = np.where(condition, fg_image, bg_image)
        cv2.imshow('output_image', output_image)
        cv2.imwrite('selfie0.png', output_image)

if __name__ == '__main__':
    args = sys.argv[1:]
    if len(args) == 0:
        print("Please give a picture path as the argument.")
        exit(0)

    file_path = args[0]
    if not os.path.isfile(args[0]):
        print(f"Please ensure the file path <{args[0]}> is valid.")
        exit(-1)

    # show_test(file_path)
    show_test2(file_path)