画像から深度マップを取得したい

特別なカメラやセンサーやステレオ画像なしで、シンプルに一枚の画像から深度マップを得たい。

Monocular Depth Estimation　で取得できる。

Monocular Depth Estimation機械学習モデルMiDaSで、単一画像から深度を取得できる。

モデルの初期化。

import torch
model_type = "MiDaS_small"
midas = torch.hub.load("intel-isl/MiDaS", model_type)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()

推論。

img = cv2.imread("pexels-pixabay-164634.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

input_batch = transform(img).to(device)
with torch.no_grad():
    prediction = midas(input_batch)

    prediction = torch.nn.functional.interpolate(
        prediction.unsqueeze(1),
        size=img.shape[:2],
        mode="bicubic",
        align_corners=False,
    ).squeeze()

depth = prediction.cpu().numpy()

depth_min = depth.min()
depth_max = depth.max()

max_val = 255

if depth_max - depth_min > np.finfo("float").eps:
    out = max_val * (depth - depth_min) / (depth_max - depth_min)
else:
    out = np.zeros(depth.shape, dtype=depth.type)

cv2.imwrite(path + ".png", out.astype("uint8"))

🐣

フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium
GitHub