semantic segmentationで道路を青色に塗ってみました

画像をピクセル単位で処理する、セマンティックセグメンテーションのモデルも、表示に注意を払えば、十分高速な処理に適用可能だと思います
特集

SIN

2020.09.19

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

1 はじめに

CX事業本部の平内（SIN）です。

ディープラーニングによる画像認識は、基本的に、Object Detection、Image classification、semantic segmentationの3つとなりますが、今回は、セマンティックセグメンテーションのモデルを使用して、道路を識別する要領を試してみました。

最初に、試してみた様子です。

動画は、Pixels Videosに公開されているものを利用させて頂きました。

2 モデル

今回も、推論に使用したフレームワークは、Open VINO Toolkitです。

そして、モデルは、OpenVINOで利用可能なIR（中間表現フォーマット）へ変換されて公開されている、road-segmentation-adas-0001を利用させて頂きました。


road-segmentation-adas-0001

road-segmentation-adas-0001は、各ピクセルをバックグラウンド、道路、縁石、標識の4つのクラスに分類するセグメンテーションモデルです。

入出力は、以下のようになっています。

Inputs

A blob with a BGR image in the format: [B, C=4, H=512, W=896], where:

  • B – batch size
  • C – number of channels
  • H – image height
  • W – image width

Outputs

The output is a blob with the shape [B, C=4, H=512, W=896]. It can be treated as a four-channel feature map, where each channel is a probability of one of the classes: BG, road, curb, mark.

3 コード（静止画）

最初に、静止画を処理しているコードです。

RoadSegmentationは、モデルをラップするクラスです。

infer()が、推論を行うメソッドですが、戻り値は、各ピクセルを、4つのクラスに対する信頼度で表現されたTensolです。 しきい値を決めて、一定の値（ここでは、THRESHOLD=0.5としています）となるピクセルを色付けしています。

index.py

import numpy as np
import time
import random
import cv2
import glob
import os
import time
from model import Model
from mrcnn import visualize
from openvino.inference_engine import IEPlugin  

class RoadSegmentation(Model):

    def __init__(self, plugin, model_path, num_requests=2):
        super().__init__(plugin, model_path, num_requests, None)
        _, _, h, w = self.input_size
        self.__input_height = h
        self.__input_width = w

    def __prepare_frame(self, frame):
        initial_h, initial_w = frame.shape[:2]
        scale_h, scale_w = initial_h / float(self.__input_height), initial_w / float(self.__input_width)
        in_frame = cv2.resize(frame, (self.__input_width, self.__input_height))
        in_frame = in_frame.transpose((2, 0, 1))
        in_frame = in_frame.reshape(self.input_size)
        return in_frame, scale_h, scale_w

    def infer(self, frame):
        in_frame, _, _ = self.__prepare_frame(frame)
        result = super().infer(in_frame)
        # [1, 4, 512, 896] => [4, 512, 896]
        return result.squeeze()

# MacOS
device = "CPU"
plugin_dirs = "/opt/intel/openvino/deployment_tools/inference_engine/lib/intel64"
modelPath = "./FP32/"

plugin = IEPlugin(device=device, plugin_dirs = plugin_dirs)
segmentation = RoadSegmentation(plugin, modelPath + "road-segmentation-adas-0001")

input = "./input"
output = "./output"

THRESHOLD= 0.5
WIDTH = 896
HEIGHT = 512
color = (1.0, 0.0, 0.0) 

for file in glob.glob("{}/*.png".format(input)):
    baseName = os.path.basename(file)
    orgImg = cv2.imread(file)
    orgImg = cv2.resize(orgImg, (WIDTH, HEIGHT))
    segmentImg = orgImg.copy()

    start = time.time()

    result =  segmentation.infer(orgImg)
    # result[0]:BG [1]:road [2]:curb [3]:mark
    mask = result[1] > THRESHOLD
    segmentImg = visualize.apply_mask(segmentImg, mask, color)
    print ("processing time:{:.3f} sec {}".format(time.time() - start, baseName))

    img = cv2.vconcat([orgImg, segmentImg]) 
    path = "{}/{}".format(output, baseName)
    cv2.imwrite(path, img)

下記は、OpenVINOでコンピュータビジョン関連のモデルを使用する場合の、ベースとなるクラスです。

model.py

#
# 下記のコードを参考にさせて頂きました。
# https://github.com/openvinotoolkit/open_model_zoo/blob/master/demos/python_demos/asl_recognition_demo/asl_recognition_demo/common.py
#

from openvino.inference_engine import IENetwork

class Model:
    def __init__(self, plugin, model_path, num_requests, output_shape=None):

        if model_path.endswith((".xml", ".bin")):
            model_path = model_path[:-4]

        model = IENetwork(model_path + ".xml", model_path + ".bin")    
        self.net = plugin.load(network=model)
        assert len(self.net.input_info) == 1, "One input is expected"
        self.input_name = next(iter(self.net.input_info))
        if len(self.net.outputs) > 1:
            if output_shape is not None:
                candidates = []
                for candidate_name in self.net.outputs:
                    candidate_shape = self.exec_net.requests[0].output_blobs[candidate_name].buffer.shape
                    if len(candidate_shape) != len(output_shape):
                        continue

                    matches = [src == trg or trg < 0
                               for src, trg in zip(candidate_shape, output_shape)]
                    if all(matches):
                        candidates.append(candidate_name)

                if len(candidates) != 1:
                    raise Exception("One output is expected")

                self.output_name = candidates[0]
            else:
                raise Exception("One output is expected")
        else:
            self.output_name = next(iter(self.net.outputs))

        self.input_size = self.net.input_info[self.input_name].input_data.shape
        self.output_size = self.net.requests[0].output_blobs[self.output_name].buffer.shape
        self.num_requests = num_requests

    def infer(self, data):
        input_data = {self.input_name: data}
        infer_result = self.net.infer(input_data)
        return infer_result[self.output_name]

4 出力

上記のコードを実行すると、inputフォルダに置かれた画像を処理し、上下に連結した画像がoutputに出力されます。

そして、コンソールへの出力は、以下のようになっており、各画像ごとの推論及び、着色処理は、0.05sec程度となっていることが分かります。

processing time:0.059 sec img-006.png
processing time:0.049 sec img-007.png
processing time:0.049 sec img-005.png
processing time:0.046 sec img-004.png
processing time:0.042 sec img-001.png
processing time:0.043 sec img-003.png
processing time:0.036 sec img-002.png
processing time:0.038 sec img-009.png
processing time:0.043 sec img-008.png

推論の出力から、画像に色付けする処理ですが、ピクセルごとループして処理することはもちろん可能です。

しかし、ピクセル単位の画像処理は、非常に負荷が高いため高速化はできません。

今回は、しきい値を超えているかどうかのマスク作成をnumpyの演算、そして、マスクで当該ピクセルに色付けを行う処理は、visualize.apply_mask()を利用しました。

5 コード（動画）

1画像の処理が、0.05sec程度であれば、充分動画にも適用できるということで書いたコードが、以下です。

冒頭の動画は、このコードによるものです。

再生ループの負荷が少しでも下がるように、予め動画のサイズを、モデルの入力サイズに合せました。アスペクト比が維持されていない可能性があるので、ここは、割り切りとさせて下さい。

$ ffmpeg -i input.mp4 -vf scale=896:512 output.mp4

index2.py

import numpy as np
import time
import random
import cv2
import glob
import os
import time
from model import Model
from mrcnn import visualize
from openvino.inference_engine import IEPlugin  

class RoadSegmentation(Model):

    def __init__(self, plugin, model_path, num_requests=2):
        super().__init__(plugin, model_path, num_requests, None)
        _, _, h, w = self.input_size
        self.__input_height = h
        self.__input_width = w

    def __prepare_frame(self, frame):
        initial_h, initial_w = frame.shape[:2]
        scale_h, scale_w = initial_h / float(self.__input_height), initial_w / float(self.__input_width)
        in_frame = cv2.resize(frame, (self.__input_width, self.__input_height))
        in_frame = in_frame.transpose((2, 0, 1))
        in_frame = in_frame.reshape(self.input_size)
        return in_frame, scale_h, scale_w

    def infer(self, frame):
        in_frame, _, _ = self.__prepare_frame(frame)

        result = super().infer(in_frame)
        # [1, 4, 512, 896] => [4, 512, 896]
        return result.squeeze()

# MacOS
device = "CPU"
plugin_dirs = "/opt/intel/openvino/deployment_tools/inference_engine/lib/intel64"
modelPath = "./FP32/"

plugin = IEPlugin(device=device, plugin_dirs = plugin_dirs)
segmentation = RoadSegmentation(plugin, modelPath + "road-segmentation-adas-0001")

VIDEO = "video/output.mp4"
cap = cv2.VideoCapture (VIDEO)

THRESHOLD= 0.5
color = (1.0, 0.0, 0.0) 

while(True):

    grabbed, frame = cap.read()
    if not grabbed: # ループ再生
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        continue

    start = time.time()

    result =  segmentation.infer(frame)
    # result[0]:BG [1]:road [2]:curb [3]:mark
    mask = result[1] > THRESHOLD
    frame = visualize.apply_mask(frame, mask, color)
    print ("processing time:{:.3f} sec".format(time.time() - start))

    cv2.imshow('frame', frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

model.pyは同じです。

6 最後に

今回は、セマンティックセグメンテーションのモデルを使用してみました。モデルのサイズにもよりますが、そこまで処理時間はかかりませんでした。

表示の部分に気をつければ、充分、高速な処理に利用できるのでは、と感じました。

*

関連記事

Intel OpenVinoの開発環境を準備してデモを動かしてみた

水島 達哉

2023.08.31