OpenVINO™ による車両の検出と認識#

# de -> 検出 
# re -> 認識 
# 検出モデルの初期化 
input_key_de, output_keys_de, compiled_model_de = model_init(detection_model_path) 
# 認識モデルの初期化 
input_key_re, output_keys_re, compiled_model_re = model_init(recognition_model_path) 

# 入力サイズを取得 - 検出 
height_de, width_de = list(input_key_de.shape)[2:]
# 入力サイズを取得 - 認識 
height_re, width_re = list(input_key_re.shape)[2:]

def plt_show(raw_image):     """ 
    Use matplot to show image inline 
    raw_image: input image 

    :param: raw_image:image array 
    """ 
    plt.figure(figsize=(10, 6)) 
    plt.axis("off") 
    plt.imshow(raw_image)

# 画像をロード 
url = "https://storage.openvinotoolkit.org/data/test_data/images/person-bicycle-car-detection.bmp" 
filename = "cars.jpg" 
directory = "data" 
image_file = utils.download_file( 
    url, 
    filename=filename, 
    directory=directory, 
    show_progress=False, 
    silent=True, 
    timeout=30, 
) 
assert Path(image_file).exists() 

# 画像をリード 
image_de = cv2.imread("data/cars.jpg") 
# [3, 256, 256] にリサイズ 
resized_image_de = cv2.resize(image_de, (width_de, height_de)) 
# バッチチャネルを [1, 3, 256, 256] に拡張 
input_image_de = np.expand_dims(resized_image_de.transpose(2, 0, 1), 0) 
# 画像を表示 
plt_show(cv2.cvtColor(image_de, cv2.COLOR_BGR2RGB))

def crop_images(bgr_image, resized_image, boxes, threshold=0.6) -> np.ndarray:     """ 
    Use bounding boxes from detection model to find the absolute car position 

    :param: bgr_image: raw image 
    :param: resized_image: resized image 
    :param: boxes: detection model returns rectangle position 
    :param: threshold: confidence threshold 
    :returns: car_position: car's absolute position 
    """ 
    # 画像の形状を取得して比率を計算 
    (real_y, real_x), (resized_y, resized_x) = ( 
        bgr_image.shape[:2], 
        resized_image.shape[:2], 
    ) 
    ratio_x, ratio_y = real_x / resized_x, real_y / resized_y 

    # ボックスの比率を検出 
    boxes = boxes[:, 2:]
    # 車両の位置を保存 
    car_position = [] 
    # ゼロ以外のボックスを反復処理 
    for box in boxes:
        # 配列の最後の位置から信頼度係数を選択 
        conf = box[0] 
        if conf > threshold:
                # float を int に変換し、各ボックスのコーナーの位置を x と y の比率で乗算 
                # 境界ボックスが画像の上部にある場合は、 
                # 画像上で見えるように上部のボックスバーを少し下に配置 
                (x_min, y_min, x_max, y_max) = [ 
                        (int(max(corner_position * ratio_y * resized_y, 10)) if idx % 2 else int(corner_position * ratio_x * resized_x)) 
                        for idx, corner_position in enumerate(box[1:])                 ] 

                car_position.append([x_min, y_min, x_max, y_max]) 

    return car_position

# 車の位置を検出 
car_position = crop_images(image_de, resized_image_de, boxes)

# 認識する車両を選択 
pos = car_position[0] 
# [y_min:y_max, x_min:x_max] で画像をトリミング 
test_car = image_de[pos[1] : pos[3], pos[0] : pos[2]] 
# 画像のサイズを input_size に変更 
resized_image_re = cv2.resize(test_car, (width_re, height_re)) 
input_image_re = np.expand_dims(resized_image_re.transpose(2, 0, 1), 0) 
plt_show(cv2.cvtColor(resized_image_re, cv2.COLOR_BGR2RGB))

def convert_result_to_image(compiled_model_re, bgr_image, resized_image, boxes, threshold=0.6):     """ 
    Use Detection model boxes to draw rectangles and plot the result 
    :param: compiled_model_re: recognition net 
    :param: input_key_re: recognition input key 
    :param: bgr_image: raw image 
    :param: resized_image: resized image 
    :param: boxes: detection model returns rectangle position 
    :param: threshold: confidence threshold 
    :returns: rgb_image: processed image 
    """ 
    # ボックスと説明の色を定義 
    colors = {"red": (255, 0, 0), "green": (0, 255, 0)} 

    # ベースイメージを BGR 形式から RGB 形式に変換 
    rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB) 

    # 車の位置を検出 
    car_position = crop_images(image_de, resized_image, boxes) 

    for x_min, y_min, x_max, y_max in car_position:
        # 車両認識推論を実行 
        attr_color, attr_type = vehicle_recognition(compiled_model_re, (72, 72), image_de[y_min:y_max, x_min:x_max]) 

        # 車でウィンドウを閉じる 
        plt.close() 

        # 位置に基づいて境界ボックスを描画
        # `rectangle` 関数のパラメーターは、image、start_point、end_point、color、thickness 
        rgb_image = cv2.rectangle(rgb_image, (x_min, y_min), (x_max, y_max), colors["red"], 2) 

        # 車両の属性をプリント
        # `putText` 関数のパラメーターは、img、text、org、fontFace、fontScale、color、thickness、lineType 
        rgb_image = cv2.putText( 
            rgb_image, 
            f"{attr_color} {attr_type}", 
            (x_min, y_min - 10),
             cv2.FONT_HERSHEY_SIMPLEX, 
            2, 
            colors["green"], 
            10, 
            cv2.LINE_AA, 
        ) 
    return rgb_image

plt_show(convert_result_to_image(compiled_model_re, image_de, resized_image_de, boxes))