blanchon/Metric3D

1import numpy as np2import torch3from plyfile import PlyData, PlyElement4import cv25 6import trimesh7 8def get_pcd_base(H, W, u0, v0, fx, fy):9    x_row = np.arange(0, W)10    x = np.tile(x_row, (H, 1))11    x = x.astype(np.float32)12    u_m_u0 = x - u013 14    y_col = np.arange(0, H)  # y_col = np.arange(0, height)15    y = np.tile(y_col, (W, 1)).T16    y = y.astype(np.float32)17    v_m_v0 = y - v018 19    x = u_m_u0 / fx20    y = v_m_v0 / fy21    z = np.ones_like(x)22    pw = np.stack([x, y, z], axis=2)  # [h, w, c]23    return pw24 25 26def reconstruct_pcd(depth, fx, fy, u0, v0, pcd_base=None, mask=None):27    if type(depth) == torch.__name__:28        depth = depth.cpu().numpy().squeeze()29    depth = cv2.medianBlur(depth, 5)30    if pcd_base is None:31        H, W = depth.shape32        pcd_base = get_pcd_base(H, W, u0, v0, fx, fy)33    pcd = depth[:, :, None] * pcd_base34    if mask:35        pcd[mask] = 036    return pcd37 38 39def save_point_cloud(pcd, rgb, filename, binary=True):40    """Save an RGB point cloud as a PLY file.41    :paras42        @pcd: Nx3 matrix, the XYZ coordinates43        @rgb: Nx3 matrix, the rgb colors for each 3D point44    """45    assert pcd.shape[0] == rgb.shape[0]46 47    if rgb is None:48        gray_concat = np.tile(np.array([128], dtype=np.uint8),49                              (pcd.shape[0], 3))50        points_3d = np.hstack((pcd, gray_concat))51    else:52        points_3d = np.hstack((pcd, rgb))53    python_types = (float, float, float, int, int, int)54    npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'),55                 ('green', 'u1'), ('blue', 'u1')]56    if binary is True:57        # Format into Numpy structured array58        vertices = []59        for row_idx in range(points_3d.shape[0]):60            cur_point = points_3d[row_idx]61            vertices.append(62                tuple(63                    dtype(point)64                    for dtype, point in zip(python_types, cur_point)))65        vertices_array = np.array(vertices, dtype=npy_types)66        el = PlyElement.describe(vertices_array, 'vertex')67 68        # write69        PlyData([el]).write(filename)70    else:71        x = np.squeeze(points_3d[:, 0])72        y = np.squeeze(points_3d[:, 1])73        z = np.squeeze(points_3d[:, 2])74        r = np.squeeze(points_3d[:, 3])75        g = np.squeeze(points_3d[:, 4])76        b = np.squeeze(points_3d[:, 5])77 78        ply_head = 'ply\n' \79                    'format ascii 1.0\n' \80                    'element vertex %d\n' \81                    'property float x\n' \82                    'property float y\n' \83                    'property float z\n' \84                    'property uchar red\n' \85                    'property uchar green\n' \86                    'property uchar blue\n' \87                    'end_header' % r.shape[0]88        # ---- Save ply data to disk89        np.savetxt(filename, np.column_stack([x, y, z, r, g, b]), fmt='%f %f %f %d %d %d', header=ply_head, comments='')90 91def ply_to_obj(ply_file, obj_file):92    mesh = trimesh.load_mesh(ply_file)93    # T2 = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])94    # mesh.apply_transform(T2)95    mesh.export(obj_file)96 97 98# import numpy as np99 100# def save_point_cloud_to_obj(points, colors, file_name):101#     """102#     Save a numpy array of point cloud data with color to an OBJ file.103    104#     Args:105#     points (np.ndarray): A numpy array of shape (H, W, 3) where H is height, W is width.106#     colors (np.ndarray): A numpy array of color data, shape (H, W, 3), values should be in [0, 1].107#     file_name (str): The path to the output .obj file.108#     """109#     H, W, _ = points.shape110#     assert points.shape == colors.shape, "Points and colors must have the same shape"111    112#     with open(file_name, 'w') as file:113#         for i in range(H):114#             for j in range(W):115#                 x, y, z = points[i, j]116#                 r, g, b = colors[i, j]117#                 file.write(f"v {x} {y} {z} {r} {g} {b}\n")118