[Python] 圖像處理 – Depthmap 002 – 深度圖轉 Point Cloud

import open3d as o3d
import cv2
import torch
import numpy as np

# Load the BGR image
image = cv2.imread('demo.png')
# Turn it into RGB
imageRGB = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Load the MiDaS Model and decide if use GPU or CPU
midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
device = torch.device(
    "cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()

# Transform the image for input to MiDaS
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
transform = midas_transforms.dpt_transform if hasattr(
    midas_transforms, 'dpt_transform') else midas_transforms.default_transform
input_batch = transform(imageRGB).to(device)

# Predict depth map
with torch.no_grad():
    prediction = midas(input_batch)
    prediction = torch.nn.functional.interpolate(prediction.unsqueeze(
        1), size=imageRGB.shape[:2], mode="bicubic", align_corners=False).squeeze()

depth_map = prediction.cpu().numpy()

# Convert depth map to 3D point cloud
# -- Focal Length
fx = 500
fy = 500
# -- Principal Point
cx = imageRGB.shape[1] / 2
cy = imageRGB.shape[0] / 2

# Generate 3D point cloud from depth map
meshX, meshY = np.meshgrid(
    np.arange(depth_map.shape[1]), np.arange(depth_map.shape[0]))
x = ((meshX - cx) * depth_map) / fx
y = ((meshY - cy) * depth_map) / fy
z = depth_map

# Create Open3D point cloud with colors from original RGB image
pointCloud = o3d.geometry.PointCloud()
pointCloud.points = o3d.utility.Vector3dVector(
    np.dstack((y, x, z)).reshape(-1, 3))
pointCloud.colors = o3d.utility.Vector3dVector(imageRGB.reshape(-1, 3) / 255.0)

# [IF Needed]
# Rotate the point cloud
# * A pi(π) represents 180 degrees
rotation = pointCloud.get_rotation_matrix_from_xyz(
    (0, 0, -(np.pi/2)))
pointCloud.rotate(rotation, center=(0, 0, 0))

# Visualize point cloud
o3d.visualization.draw_geometries([pointCloud])