python基于mask检测人体并输出黑白剪影

使用mask Rcnn实现人物黑白剪影

因为涉及到这块内容，就想着分享出来希望能给大家一个参考。
使用的是python+mask，不会用的可以看看气球那个例子（没耐心的不看也行，反正也不复杂）。
生成的是和图片同名的文件夹，里面包含了处理好的每张图片，如果想放到一起修改一下存储路径即可。

原理简介

基于的是mask原理进行识别，然后将“人”这个分类单独提取出来，再进行黑白渲染，生成剪影。
以此类推，如果想对其他物体进行剪影描绘，只需更改需要凸显的分类即可，相当简单粗暴。

准备

需安装opencv，numpy等，这部分参考mask安装要求，下面只贴出处理用的代码部分。
coco的模型可以直接下载使用

代码分析

import cv2
import numpy as np
import os
from samples import coco
from mrcnn import utils
from mrcnn import model as modellib
from datetime import datetime


# Load the pre-trained model data
ROOT_DIR = os.getcwd()
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)


# Change the config infermation
class InferenceConfig(coco.CocoConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1


config = InferenceConfig()

# COCO dataset object names
model = modellib.MaskRCNN(
    mode="inference", model_dir=MODEL_DIR, config=config
)
model.load_weights(COCO_MODEL_PATH, by_name=True)
class_names = [
    'BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
    'bus', 'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
    'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
    'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
    'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard',
    'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
    'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
    'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
    'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
    'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
    'teddy bear', 'hair drier', 'toothbrush'
]

# ==========================================================
# This function is used to change the colorful background information

def apply_mask_black(image, mask):
    image[:, :, 0] = np.where(
        mask == 0,
        0,
        255
    )
    image[:, :, 1] = np.where(
        mask == 0,
        0,
        255
    )
    image[:, :, 2] = np.where(
        mask == 0,
        0,
        255
    )
    return image


# This function is used to show the object detection result in original image.
def display_instances(image, boxes, masks, ids, names, scores):
    # max_area will save the largest object for all the detection results
    max_area = 0

    # n_instances saves the amount of all objects
    n_instances = boxes.shape[0]

    if not n_instances:
        print('NO INSTANCES TO DISPLAY')
    else:
        assert boxes.shape[0] == masks.shape[-1] == ids.shape[0]
    mask = np.array([0])
    # mask = False
    for i in range(n_instances):
        if not np.any(boxes[i]):
            continue

        # compute the square of each object
        y1, x1, y2, x2 = boxes[i]
        square = (y2 - y1) * (x2 - x1)

        # use label to select person object from all the 80 classes in COCO dataset
        label = names[ids[i]]

        if label == 'person':
            # save the largest object in the image as main character
            # other people will be regarded as background
            if square > max_area:
                max_area = square
                mask = masks[:, :, i]
            else:
                continue
        else:
            continue

        # apply mask for the image
    # by mistake you put apply_mask inside for loop or you can write continue in if also

    image = apply_mask_black(image, mask)
    return image if not mask.sum() ==0 else mask
    # return image

def img_black(image_path):
    images_list = os.listdir(image_path)
    out_path = /home/nvidia/xxx/
    for name in sorted(images_list):
        original_image = os.path.join(image_path, name)
        image = cv2.imread(original_image)
        # gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        results = model.detect([image], verbose=0)
        r = results[0]
        frame = display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])

        if frame.sum() == 0:
            continue
        else:
            # 生成每个图片对应的同名文件夹，里面放有该图片
            output_dir = out_path+name+'/'
            os.mkdir(output_dir)
            cv2.imwrite(output_dir+name,frame)

if __name__ == '__main__':
    time = datetime.now()
    # images path
    image_path = '/home/pcm/Mask-GaitSet/IMG/images'
    img_black(image_path)
    print('Evaluation complete. Cost:', datetime.now() - time)

最终结果

下面用的是网上随便找的一张图片进行举例，前两张是原图，后二张出来的图片就是剪影效果（应该不算侵权吧？如果有懂的同学麻烦告诉我）



效果貌似还可以的说，但是这玩意要想效果好点，还是自己训练个模型会更准点。

=2022.7更新==
补充一下论文：
论文地址
demo使用下载地址
背景置灰demo