セマンティックセグメンテーションで背景透過処理を行いたい

Question

前提・実現したいこと

前回『セマンティックセグメンテーション Deeplab v3+を用いて画像から人型を認識させたい。』の引き続きになります。

ローカルより人と認識するところへマスクする画像を取得することはできました。

手順
tensorflow/models
このコードを改変したものをanaconda仮想環境上で実行。

import os
from io import BytesIO
import tarfile
import tempfile
from six.moves import urllib

from matplotlib import gridspec
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image

import tensorflow as tf

import random

class DeepLabModel(object):
  """Class to load deeplab model and run inference."""

  INPUT_TENSOR_NAME = 'ImageTensor:0'
  OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
  INPUT_SIZE = 513
  FROZEN_GRAPH_NAME = 'frozen_inference_graph'

  def __init__(self, tarball_path):
    """Creates and loads pretrained deeplab model."""
    self.graph = tf.Graph()

    graph_def = None
    # Extract frozen graph from tar archive.
    tar_file = tarfile.open(tarball_path)
    for tar_info in tar_file.getmembers():
      if self.FROZEN_GRAPH_NAME in os.path.basename(tar_info.name):
        file_handle = tar_file.extractfile(tar_info)
        graph_def = tf.GraphDef.FromString(file_handle.read())
        break

    tar_file.close()

    if graph_def is None:
      raise RuntimeError('Cannot find inference graph in tar archive.')

    with self.graph.as_default():
      tf.import_graph_def(graph_def, name='')

    self.sess = tf.Session(graph=self.graph)

  def run(self, image):
    """Runs inference on a single image.
    Args:
      image: A PIL.Image object, raw input image.
    Returns:
      resized_image: RGB image resized from original input image.
      seg_map: Segmentation map of `resized_image`.
    """
    width, height = image.size
    resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
    target_size = (int(resize_ratio * width), int(resize_ratio * height))
    resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
    batch_seg_map = self.sess.run(
        self.OUTPUT_TENSOR_NAME,
        feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
    seg_map = batch_seg_map[0]
    return resized_image, seg_map


def create_pascal_label_colormap():
  """Creates a label colormap used in PASCAL VOC segmentation benchmark.
  Returns:
    A Colormap for visualizing segmentation results.
  """
  colormap = np.zeros((256, 3), dtype=int)
  ind = np.arange(256, dtype=int)

  for shift in reversed(range(8)):
    for channel in range(3):
      colormap[:, channel] |= ((ind >> channel) & 1) << shift
    ind >>= 3

  return colormap


def label_to_color_image(label):
  """Adds color defined by the dataset colormap to the label.
  Args:
    label: A 2D array with integer type, storing the segmentation label.
  Returns:
    result: A 2D array with floating type. The element of the array
      is the color indexed by the corresponding element in the input label
      to the PASCAL color map.
  Raises:
    ValueError: If label is not of rank 2 or its value is larger than color
      map maximum entry.
  """
  if label.ndim != 2:
    raise ValueError('Expect 2-D input label')

  colormap = create_pascal_label_colormap()

  if np.max(label) >= len(colormap):
    raise ValueError('label value too large.')

  return colormap[label]


def vis_segmentation(image, seg_map):
  """Visualizes input image, segmentation map and overlay view."""
  plt.figure(figsize=(15, 5))
  grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])

  plt.subplot(grid_spec[0])
  plt.imshow(image)
  plt.axis('off')
  plt.title('input image')

  plt.subplot(grid_spec[1])
  seg_image = label_to_color_image(seg_map).astype(np.uint8)
  plt.imshow(seg_image)
  plt.axis('off')
  plt.title('segmentation map')

  plt.subplot(grid_spec[2])
  plt.imshow(image)
  plt.imshow(seg_image, alpha=0.7)
  plt.axis('off')
  plt.title('segmentation overlay')

  unique_labels = np.unique(seg_map)
  ax = plt.subplot(grid_spec[3])
  plt.imshow(
      FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest')
  ax.yaxis.tick_right()
  plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
  plt.xticks([], [])
  ax.tick_params(width=0.0)
  plt.grid('off')
  plt.show()


LABEL_NAMES = np.asarray([
    'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
    'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
    'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'
])

FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)

#@title Select and download models {display-mode: "form"}

MODEL_NAME = 'mobilenetv2_coco_voctrainaug'  # @param ['mobilenetv2_coco_voctrainaug', 'mobilenetv2_coco_voctrainval', 'xception_coco_voctrainaug', 'xception_coco_voctrainval']

_DOWNLOAD_URL_PREFIX = 'http://download.tensorflow.org/models/'
_MODEL_URLS = {
    'mobilenetv2_coco_voctrainaug':
        'deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz',
    'mobilenetv2_coco_voctrainval':
        'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz',
    'xception_coco_voctrainaug':
        'deeplabv3_pascal_train_aug_2018_01_04.tar.gz',
    'xception_coco_voctrainval':
        'deeplabv3_pascal_trainval_2018_01_04.tar.gz',
}
_TARBALL_NAME = 'deeplab_model.tar.gz'

model_dir = tempfile.mkdtemp()
tf.gfile.MakeDirs(model_dir)

download_path = os.path.join(model_dir, _TARBALL_NAME)
print('downloading model, this might take a while...')
urllib.request.urlretrieve(_DOWNLOAD_URL_PREFIX + _MODEL_URLS[MODEL_NAME],
                   download_path)
print('download completed! loading DeepLab model...')

MODEL = DeepLabModel(download_path)
print('model loaded successfully!')

#@title Run on sample images {display-mode: "form"}

'''
SAMPLE_IMAGE = 'image1'  # @param ['image1', 'image2', 'image3']
IMAGE_URL = ''  #@param {type:"string"}
_SAMPLE_URL = ('https://github.com/tensorflow/models/blob/master/research/'
               'deeplab/g3doc/img/%s.jpg?raw=true')
'''


def run_visualization(url):
  """Inferences DeepLab model and visualizes result."""
  try:
    # f = urllib.request.urlopen(url)
    # jpeg_str = f.read()
    # original_im = Image.open(BytesIO(jpeg_str))
    original_im = Image.open(url)
  except IOError:
    print('Cannot retrieve image. Please check url: ' + url)
    return

  print('running deeplab on image %s...' % url)
  resized_im, seg_map = MODEL.run(original_im)

  vis_segmentation(resized_im, seg_map)
  import csv
  f = open('./test_image/seg_map.csv', 'w')
  writer = csv.writer(f, lineterminator='\n')
  writer.writerow(seg_map)
  f.close()

image_path = './test_image/3.jpg'
run_visualization(image_path)

test_imageフォルダを同階層に準備し目当ての画像（3.jpgなど)置いて実行。

例）（再使用が許可された画像を使用しています。）

↓

ここから、人物以外を透過処理させたいと考えています。
特にこちらの方がphpで実践していましたが、pythonとしてローカル上で実行させたいと考えています。

もしこれを利用するとなりますと、xamppや他のサーバなどの準備が必要になると考えています。
最終的にはpythonのexeにまとめたいため、phpのexecファイルが絡むと他サーバなどの準備が絡むので不可能と判断しています。そのためpythonとして統一して作りたいと思いました。

私の環境

OS Windows10
GPU RTX2070
GPUのドライバー 26.21.14.4166
python Python 3.7 :: Anaconda, Inc.
CUDA CUDA Toolkit 9.0
cuDNN release 9.0, V9.0.176
tensorflow_gpu v1.13.1
keras 2.3.4

行ったこと

seg_mapの中身を確認しますと、

import csv
f = open('./test_image/seg_map2.csv', 'w')
writer = csv.writer(f, lineterminator='\n')
writer.writerow(seg_map)
f.close()

seg_map.csv

ここから視認性を高くするために
①\r\nを空白置換
②","を",\r\n"へ置換

置換後seg_map.csv

semantic segmentationはピクセル単位で認識するため、人を15と表示しています。
これは姉妹兄弟？のお姉ちゃんの頭の部分でしょうか。

座標情報から0と表示されているところのみ、以下コードを利用しようと考えています。
OpenCV python 座標情報から画像を切り抜きたい

import cv2
import numpy as np

# img1: 元の画像, 形状が (Width, Height, 3) の numpy 配列
# contour: 輪郭の頂点情報, 形状が (NumPoints, 2) の numpy 配列

# マスク画像を作成する。
# 前景の画素は (255, 255, 255)、背景の画素は (0, 0, 0)
mask = np.zeros_like(img1)
cv2.fillConvexPoly(mask, contour, color=(255, 255, 255))

# 背景画像
bg_color = (50, 200, 0)
img2 = np.full_like(img, bg_color)

# np.where() はマスクの値が (255, 255, 255) の要素は前景画像 img1 の値、
# マスクの値が (0, 0, 0) の要素は背景画像 img2 の値を返す。
result = np.where(mask==255, img1, img2)

しかし、このcontourを取得する方法が必要でした。
contourは、もともと多角形を記述を行うために

contours = np.array( [ [100,100], [100,230], [230, 250], [150,70] ] )

と座標を並べ行くイメージで、この人物と背景の堺の座標が取得できれば上記codeで可能とは考えているのですが、そもそもこの座標情報から背景を透過させるライブラリなどあるものなのでしょうか。人へのマスキングができているのなら、それ以外を透過させるのはすぐにできると考えていましたが、これはどのような方向性でいったほうがよろしいのでしょうか。知見ある方アドバイス頂けないでしょうか。よろしくお願い致します。

20200705 13：50追記

seg_map.csvと対象画像になります。
結果画像は保存を試したのですがmatplotlibで表示されるのみでした。

fig.figure()
fig.savefig(image.jpg)

こちらのコードを途中入れたのですが保存ができませんでした。

20200705 14:05 追記

追加コマンド	結果
print(type(seg_map))	<class 'numpy.ndarray'>
print(seg_map.shape, seg_map.dtype)	(341, 513) int64
np.save("seg_map.npy", seg_map)	（GoogleDriveに追加します）
cv2.imwrite("seg_map.png", seg_map.astype(np.uint8))	（GoogleDriveに追加します）

Comments

[tiitoi]

入力画像とセグメンテーション結果のファイルをどこかにアップしていただくことはできますか？
変数 seg_map も実態が ndarray ならば、 CSV 形式ではなく、png 形式で保存すれば、画像になるので、teratail の質問欄にアップできると思います。

[dendenmushi]

ありがとうございます。結果などアップします。お待ちください。

[dendenmushi]

Google 共有ドライブへアップしました。結果画像は保存が難しいようです。。

[tiitoi]

segmap の型はどうなっていますか？
print(type(segmap)) の結果を教えて下さい
その結果が ndarray 配列であれば、
print(segmap.shape, segmap.dtype) の出力結果も記載していただけますか
もし ndarray 配列であれば、
np.save("segmap.npy", segmap) で npy ファイルで保存するか
または
cv2.imwrite("segmap.png", segmap.astype(np.uint8))
で画像として保存できるのではないかと思います。

[dendenmushi]

ありがとうございます。npyで保存もできましたし画像も保存できました。結果を追記しました。

Answer 1

セグメンテーション結果を表す形状が (Height, Width) の int64 型の配列は、各要素が対応する画素のラベルを表しています。
クラス:人 のラベルは15のようなので、

1 元画像にアルファチャンネルを追加する。
cv2.cvtColor() でアルファチャンネルを追加できます。

dst = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

2 セグメンテーション結果の値が15以外の位置のアルファチャンネルの値を0 (透過) にする。

dst[segmap != 15, 3] = 0

以上で、人以外の画素を透過できます。

サンプルコード

import cv2
import numpy as np

# 画像を読み込む。
img = cv2.imread("test.jpg")

# セグメンテーション結果を読み込む。
segmap = np.load("seg_map.npy")

# セグメンテーション結果を元の画像サイズに合わせる。
segmap = cv2.resize(
    segmap, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST
)
print(img.shape, segmap.shape)

# アルファチャンネル追加する。
dst = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

# 人以外のアルファチャンネルを透過する。
dst[segmap != 15, 3] = 0

# 保存する。
cv2.imwrite("output.png", dst)

Comments

[dendenmushi]

非常に勉強になりました。アルファチャンネル追加じたい知りませんでした。本当にありがとうございました。