複数の人が写りこんだ写真から最も人らしい人を抽出する

「最も人らしい人」(=「人である確率が最も高い物体」)を抽出する。
f:id:touch-sp:20190815162512p:plain

コード

import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from gluoncv import model_zoo, data, utils

net = model_zoo.get_model('mask_rcnn_fpn_resnet101_v1d_coco', pretrained=True)

im_fname = 'example.jpg'

x, orig_img = data.transforms.presets.rcnn.load_test(im_fname)

ids, scores, bboxes, masks = [xx[0].asnumpy() for xx in net(x)]

person_index = np.where(ids[:,0]==0)

if len(person_index[0]) != 0:
    
    index = person_index[0][0]
    width, height = orig_img.shape[1], orig_img.shape[0]

    #if gluoncv==0.4
    masks_ = utils.viz.expand_mask(masks[index:index+1], bboxes[index:index+1], (width, height), scores[index:index+1])
        
    ##if gluoncv==0.5 beta
    #masks_, _ = utils.viz.expand_mask(masks[index:index+1], bboxes[index:index+1], (width, height), scores[index:index+1])

    masks_ = masks_[0]
    
    index_non = np.where(masks_ == 0)
    
    orig_img[:,:,0][index_non] = 255
    orig_img[:,:,1][index_non] = 255
    orig_img[:,:,2][index_non] = 255
    
    img_size = Image.open(im_fname).size

    save_img = Image.fromarray(orig_img).resize(img_size)
    save_img.save('result_' + im_fname)

動機

どうしてこのようなことをしようと思ったか?
顔画像データセット「Labeled Faces in the Wild」を使ってGANで顔を描こうと思ったが、画像に複数人が写っていたり背景がまちまちであったりしたのでこのような前処理を行うと精度が上がるのではと考えた。

import numpy as np
import glob
import os
from PIL import Image
from gluoncv import model_zoo, data, utils
import mxnet as mx
from mxnet import gluon

ctx = mx.gpu()

net = model_zoo.get_model('mask_rcnn_fpn_resnet101_v1d_coco', pretrained=True)
net.collect_params().reset_ctx(ctx)

dataset = glob.glob('lfw-deepfunneled/*/*.jpg')

sampler = gluon.data.SequentialSampler(len(dataset))
data_loader = gluon.data.BatchSampler(sampler, batch_size=200)

count = 0

for batch in data_loader:

    img_list = [dataset[i] for i in batch]

    x, orig_img = data.transforms.presets.rcnn.load_test(img_list)

    for i, (img_array, img_data) in enumerate(zip(x,orig_img)):

        img_array = img_array.as_in_context(ctx)

        ids, scores, bboxes, masks = [xx[0].asnumpy() for xx in net(img_array)]

        person_index = np.where(ids[:,0]==0)

        if len(person_index[0]) != 0:
    
            index = person_index[0][0]
            width, height = img_data.shape[1], img_data.shape[0]
        
            #if gluoncv==1.5.0
            masks_ = utils.viz.expand_mask(masks[index:index+1], bboxes[index:index+1], (width, height), scores[index:index+1])
        
            ##if gluoncv==1.6 beta
            #masks_, _ = utils.viz.expand_mask(masks[index:index+1], bboxes[index:index+1], (width, height), scores[index:index+1])

            masks_ = masks_[0]
    
            index_non = np.where(masks_ == 0)
    
            img_data[:,:,0][index_non] = 255
            img_data[:,:,1][index_non] = 255
            img_data[:,:,2][index_non] = 255

            if len(index_non[0]) > (width*height*0.7):
                out_dir = 'faces_failure'
            else:
                out_dir = 'faces_success' 
            
            save_img = Image.fromarray(img_data).resize((250,250))
            save_img.save(os.path.join(out_dir,'%d.jpg'%count))
            count += 1

以下の一文で背景が多い画像を取り除いている。

if len(index_non[0]) > (width*height*0.7):

取り除かれた画像の一部。
f:id:touch-sp:20190818091124p:plain

WGANの結果

  • 元画像

f:id:touch-sp:20190818091230p:plain

  • 800 epoch

f:id:touch-sp:20190818091255p:plain

  • 1000 epoch

f:id:touch-sp:20190818091323p:plain