OpenVINO™ を使用した 3D 点群パーツのセグメント化#

import platform 

%pip install -q "openvino>=2023.1.0" "tqdm" 

if platform.system() != "Windows":
     %pip install -q "matplotlib>=3.4" 
else:
     %pip install -q "matplotlib>=3.4,<3.7"

from pathlib import Path 
from typing import Union 
import numpy as np 
import matplotlib.pyplot as plt 
import openvino as ov 

# `notebook_utils` モジュールを取得 
import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py", 
) 
open("notebook_utils.py", "w").write(r.text) 

from notebook_utils import download_file

# データとモデルのディレクトリー、モデルソースの URL、モデルファイル名を設定 
MODEL_DIR = Path("model") 
MODEL_DIR.mkdir(exist_ok=True) 
download_file( 
    "https://storage.googleapis.com/ailia-models/pointnet_pytorch/chair_100.onnx", 
    directory=Path(MODEL_DIR), 
    show_progress=False, 
) 
onnx_model_path = MODEL_DIR / "chair_100.onnx"

ir_model_xml = onnx_model_path.with_suffix(".xml") 

core = ov.Core() 

if not ir_model_xml.exists():
    # モデルを OpenVINO モデルに変換 
    model = ov.convert_model(onnx_model_path) 
    # モデルを OpenVINO IR 形式 xml + bin でシリアル化 
    ov.save_model(model, ir_model_xml) 
else:
    # モデルをリード 
    model = core.read_model(model=ir_model_xml)

def load_data(point_file: Union[str, Path]): 
    """
    Load the point cloud data and convert it to ndarray 

    Parameters: 
        point_file: string, path of .pts data 
    Returns: 
        point_set: point clound represented in np.array format 
    """

    point_set = np.loadtxt(point_file).astype(np.float32) 

    # 正規化 
    point_set = point_set - np.expand_dims(np.mean(point_set, axis=0), 0) # center 
    dist = np.max(np.sqrt(np.sum(point_set**2, axis=1)), 0) 
    point_set = point_set / dist # scale 

    return point_set 

def visualize(point_set: np.ndarray): 
    """
    Create a 3D view for data visualization 

    Parameters: 
        point_set: np.ndarray, the coordinate data in X Y Z format 
    """

    fig = plt.figure(dpi=192, figsize=(4, 4)) 
    ax = fig.add_subplot(111, projection="3d") 
    X = point_set[:, 0] 
    Y = point_set[:, 2] 
    Z = point_set[:, 1] 

    # 各軸のビューをスケーリングして、座標データ分布に適応 
    max_range = np.array([X.max() - X.min(), Y.max() - Y.min(), Z.max() - Z.min()]).max() * 0.5 
    mid_x = (X.max() + X.min()) * 0.5 
    mid_y = (Y.max() + Y.min()) * 0.5 
    mid_z = (Z.max() + Z.min()) * 0.5 
    ax.set_xlim(mid_x - max_range, mid_x + max_range) 
    ax.set_ylim(mid_y - max_range, mid_y + max_range) 
    ax.set_zlim(mid_z - max_range, mid_z + max_range) 

    plt.tick_params(labelsize=5) 
    ax.set_xlabel("X", fontsize=10) 
    ax.set_ylabel("Y", fontsize=10) 
    ax.set_zlabel("Z", fontsize=10) 

    return ax

# Download data from the openvino_notebooks storage 
point_data = download_file( 

"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/pts/chair.pts", 
    directory="data", 
) 

points = load_data(str(point_data)) 
X = points[:, 0] 
Y = points[:, 2] 
Z = points[:, 1] 
ax = visualize(points) 
ax.scatter3D(X, Y, Z, s=5, cmap="jet", marker="o", label="chair") 
ax.set_title("3D Visualization") 
plt.legend(loc="upper right", fontsize=8) 
plt.show()

# 椅子の部品 
classes = ["back", "seat", "leg", "arm"] 

# 入力データを前処理 
point = points.transpose(1, 0) 
point = np.expand_dims(point, axis=0) 

# モデルの入力と出力の形状に関する情報をプリント 
print(f"input shape: {model.input(0).partial_shape}") 
print(f"output shape: {model.output(0).partial_shape}")