「最も人らしい人」(=「人である確率が最も高い物体」)を抽出する。
コード
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):
取り除かれた画像の一部。
WGANの結果
- 元画像
- 800 epoch
- 1000 epoch