File size: 2,355 Bytes
a7aea10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import argparse
import os
import json
import numpy as np
import imageio
from create_input import depth_to_world_coords_points, camera_list
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--folder_path", type=str)
parser.add_argument("--video_path", type=str)
parser.add_argument("--frame_id", type=int, default=0)
parser.add_argument("--max_depth", type=float, default=25)
return parser.parse_args()
def save_ply(points: np.ndarray, colors: np.ndarray, out_path: str):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
n = points.shape[0]
colors = np.clip(colors, 0, 255).astype(np.uint8)
header = (
"ply\n"
"format ascii 1.0\n"
f"element vertex {n}\n"
"property float x\n"
"property float y\n"
"property float z\n"
"property uchar red\n"
"property uchar green\n"
"property uchar blue\n"
"end_header\n"
)
with open(out_path, "w") as f:
f.write(header)
for p, c in zip(points, colors):
f.write(f"{float(p[0])} {float(p[1])} {float(p[2])} {int(c[0])} {int(c[1])} {int(c[2])}\n")
if __name__ == "__main__":
args = parse_args()
folder_path = args.folder_path
video_path = args.video_path
frame_id = args.frame_id
max_depth = args.max_depth
reader = imageio.v2.get_reader(video_path)
for i, frame in enumerate(reader):
if i == frame_id:
frame = frame.astype(np.uint8)
break
with open(os.path.join(folder_path, "depth_range.json"), "r") as f:
depth_range = json.load(f)[frame_id]
rgb = frame[:512]
depth = frame[512:, :, 0] / 255.0
depth = depth * (depth_range[1] - depth_range[0]) + depth_range[0]
depth = 1 / (depth + 1e-6)
valid_mask = np.logical_and(depth > 0, depth < max_depth)
intrinsics, extrinsics = camera_list(
num_frames=1, type="forward", Width=512, Height=512, fx=256, fy=256
)
point_map = depth_to_world_coords_points(depth, extrinsics[0], intrinsics[0])
points = point_map[valid_mask].reshape(-1, 3)
colors = rgb[valid_mask].reshape(-1, 3)
out_ply = os.path.join(folder_path, f"frame_{frame_id:06d}.ply")
save_ply(points, colors, out_ply)
print(f"Saved point cloud: {out_ply}, number of points: {points.shape[0]}")
|