in scale_calibration.py [0:0]
def calibrate_scale(video, out_dir, frame_range, args):
# COLMAP reconstruction.
print_banner("COLMAP reconstruction")
colmap_dir = pjoin(video.path, 'colmap_dense')
src_meta_file = pjoin(colmap_dir, "metadata.npz")
colmap = COLMAPProcessor(args.colmap_bin_path)
dense_dir = colmap.dense_dir(colmap_dir, 0)
if os.path.isfile(src_meta_file):
print("Checked metadata file exists.")
else:
color_dir = prepare_colmap_color(video)
if not colmap.check_dense(
dense_dir, color_dir, valid_ratio=args.dense_frame_ratio
):
path_args = [color_dir, colmap_dir]
mask_path = pjoin(video.path, 'colmap_mask')
if os.path.isdir(mask_path):
path_args.extend(['--mask_path', mask_path])
colmap_args = COLMAPParams().parse_args(
args=path_args + ['--dense_max_size', str(args.size)],
namespace=args
)
colmap.process(colmap_args)
intrinsics, extrinsics = make_camera_params_from_colmap(
video.path, colmap.sparse_dir(colmap_dir, 0)
)
np.savez(src_meta_file, intrinsics=intrinsics, extrinsics=extrinsics)
# Convert COLMAP dense depth maps to .raw file format.
print_banner("Convert COLMAP depth maps")
converted_depth_fmt = pjoin(
video.path, "depth_colmap_dense", "depth", "frame_{:06d}.raw"
)
# convert colmap dense depths to .raw
converted_depth_dir = os.path.dirname(converted_depth_fmt)
dense_depth_dir = pjoin(dense_dir, "stereo", "depth_maps")
frames = frame_range.frames()
if not check_frames(
dense_depth_dir, colmap.dense_depth_suffix(), converted_depth_dir, "",
frame_names={f"frame_{i:06d}.png" for i in frames},
):
os.makedirs(converted_depth_dir, exist_ok=True)
colmap_depth_fmt = pjoin(
dense_depth_dir, "frame_{:06d}.png" + colmap.dense_depth_suffix()
)
for i in frames:
colmap_depth_fn = colmap_depth_fmt.format(i)
if not os.path.isfile(colmap_depth_fn):
logging.warning(
"[SCALE CALIBRATION] %s does not exist.",
colmap_depth_fn
)
continue
cmp_depth = load_colmap.read_array(colmap_depth_fn)
inv_cmp_depth = 1.0 / cmp_depth
ix = np.isinf(inv_cmp_depth) | (inv_cmp_depth < 0)
inv_cmp_depth[ix] = float("nan")
image_io.save_raw_float32_image(
converted_depth_fmt.format(i), inv_cmp_depth
)
with SuppressedStdout():
visualization.visualize_depth_dir(
converted_depth_dir, converted_depth_dir,
force=True, min_percentile=0, max_percentile=99,
)
# Compute scaled depth maps
print_banner("Compute per-frame scales")
scaled_depth_dir = pjoin(out_dir, "depth_scaled_by_colmap_dense", "depth")
scaled_depth_fmt = pjoin(scaled_depth_dir, "frame_{:06d}.raw")
scales_file = pjoin(out_dir, "scales.csv")
src_depth_fmt = pjoin(
video.path, f"depth_{args.model_type}", "depth", "frame_{:06d}.raw"
)
frames = frame_range.frames()
if (
check_frames(
converted_depth_dir, ".png",
os.path.dirname(scaled_depth_fmt), ".raw"
)
and os.path.isfile(scales_file)
):
src_to_colmap_scales = np.loadtxt(scales_file, delimiter=',')
assert src_to_colmap_scales.shape[0] >= len(frames) * args.dense_frame_ratio \
and src_to_colmap_scales.shape[1] == 2, \
(f"scales shape is {src_to_colmap_scales.shape} does not match "
+ f"({len(frames)}, 2) with threshold {args.dense_frame_ratio}")
print("Existing scales file loaded.")
else:
# Scale depth maps
os.makedirs(scaled_depth_dir, exist_ok=True)
src_to_colmap_scales_map = {}
for i in frames:
converted_depth_fn = converted_depth_fmt.format(i)
if not os.path.isfile(converted_depth_fn):
logging.warning("[SCALE CALIBRATION] %s does not exist",
converted_depth_fn)
continue
# convert colmap_depth to raw
inv_cmp_depth = image_io.load_raw_float32_image(converted_depth_fn)
# compute scale for init depths
inv_src_depth = image_io.load_raw_float32_image(src_depth_fmt.format(i))
# src_depth * scale = (1/inv_src_depth) * scale == cmp_depth
inv_cmp_depth = cv2.resize(
inv_cmp_depth, inv_src_depth.shape[:2][::-1],
interpolation=cv2.INTER_NEAREST
)
ix = np.isfinite(inv_cmp_depth)
if np.sum(ix) / ix.size < args.dense_pixel_ratio:
# not enough pixels are valid and hence the frame is invalid.
continue
scales = (inv_src_depth / inv_cmp_depth)[ix]
scale = np.median(scales)
print(f"Scale[{i}]: median={scale}, std={np.std(scales)}")
# scale = np.median(inv_depth) * np.median(cmp_depth)
src_to_colmap_scales_map[i] = float(scale)
scaled_inv_src_depth = inv_src_depth / scale
image_io.save_raw_float32_image(
scaled_depth_fmt.format(i), scaled_inv_src_depth
)
with SuppressedStdout():
visualization.visualize_depth_dir(
scaled_depth_dir, scaled_depth_dir, force=True
)
# Write scales.csv
xs = sorted(src_to_colmap_scales_map.keys())
ys = [src_to_colmap_scales_map[x] for x in xs]
src_to_colmap_scales = np.stack((np.array(xs), np.array(ys)), axis=-1)
np.savetxt(scales_file, src_to_colmap_scales, delimiter=",")
valid_frames = {int(s) for s in src_to_colmap_scales[:, 0]}
# Scale the extrinsics' translations
scaled_meta_file = pjoin(out_dir, "metadata_scaled.npz")
if os.path.isfile(scaled_meta_file):
print("Scaled metadata file exists.")
else:
scales = src_to_colmap_scales[:, 1]
mean_scale = scales.mean()
print(f"[scales] mean={mean_scale}, std={np.std(scales)}")
with np.load(src_meta_file) as meta_colmap:
intrinsics = meta_colmap["intrinsics"]
extrinsics = meta_colmap["extrinsics"]
extrinsics[..., -1] /= mean_scale
np.savez(
scaled_meta_file,
intrinsics=intrinsics,
extrinsics=extrinsics,
scales=src_to_colmap_scales,
)
color_fmt = pjoin(video.path, "color_down", "frame_{:06d}.raw")
vis_dir = pjoin(out_dir, "vis_calibration_dense")
visualize_all_calibration(
extrinsics, intrinsics, scaled_depth_fmt,
color_fmt, frame_range, vis_dir,
)
return valid_frames