import argparse import json import logging import os import subprocess import sys import time from pathlib import Path import cv2 as cv import h5py import jax import jax.numpy as jnp import laspy import numpy as np import torch import yaml from tqdm import tqdm from dataset_generation.fast_proj import ( compute_depth_map, f_frame_agi, parse_calibration_xml, read_camera_file, ) from dataset_generation.batch_serialization import ( save_compact_payloads_to_pt, save_dataset_to_pt_parallel, save_tile_observations_to_pt, ) from dataset_generation.gridnet_hd_manifest import load_or_build_manifest, resolve_logits_dir from utils.utilities import load_normals_from_h5, map_main_class_indices from utils.visibility_criteria import ( compute_criteria_maps, compute_visibility_criteria_numpy, radian_to_degree, round_to_three_digits, ) logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") _PROJECTION_FN_CACHE = {} _DEPTH_FN_CACHE = {} def morton2D(x, y, bits=21): z = 0 for i in range(bits): z |= ((x >> i) & 1) << (2 * i) | ((y >> i) & 1) << (2 * i + 1) return z def load_calibration(calibration_file): try: return parse_calibration_xml(calibration_file) except Exception as exc: logging.error("Failed to load calibration %s: %s", calibration_file, exc) sys.exit(1) def load_camera_parameters(camera_file, offset): try: return read_camera_file(camera_file, offset) except Exception as exc: logging.error("Failed to load camera poses %s: %s", camera_file, exc) sys.exit(1) def _extract_labels(las_data): if hasattr(las_data, "ground_truth"): return np.asarray(las_data.ground_truth) return None def load_las_points(las_file_path): try: las_data = laspy.read(las_file_path) offset = np.array([np.min(las_data.x), np.min(las_data.y), np.min(las_data.z)]) labels = _extract_labels(las_data) points = { "X": np.array(las_data.x - offset[0], dtype=np.float32), "Y": np.array(las_data.y - offset[1], dtype=np.float32), "Z": np.array(las_data.z - offset[2], dtype=np.float32), "GT": labels, "GT_AVAILABLE": labels is not None, } codes = morton2D(points["X"].astype(int), points["Y"].astype(int)) perm = np.argsort(codes) for key in ("X", "Y", "Z"): points[key] = points[key][perm] if points["GT"] is not None: points["GT"] = points["GT"][perm] points_for_normals = np.vstack((points["X"], points["Y"], points["Z"])).T.astype(np.float32) return points, offset, las_data, points_for_normals except Exception as exc: logging.error("Failed to load LAS %s: %s", las_file_path, exc) sys.exit(1) def _get_projection_fn(param_keys): cache_key = tuple(param_keys) if cache_key not in _PROJECTION_FN_CACHE: in_axes_dict = {key: None for key in param_keys} for key in ("X", "Y", "Z"): in_axes_dict[key] = 0 _PROJECTION_FN_CACHE[cache_key] = jax.jit(jax.vmap(f_frame_agi, in_axes=(in_axes_dict,))) return _PROJECTION_FN_CACHE[cache_key] def _get_depth_fn(buffer_size, threshold): cache_key = (int(buffer_size), float(threshold)) if cache_key not in _DEPTH_FN_CACHE: _DEPTH_FN_CACHE[cache_key] = jax.jit( lambda px, py, pz, dmap: compute_depth_map( px, py, pz, dmap, buffer_size=buffer_size, threshold=threshold, ) ) return _DEPTH_FN_CACHE[cache_key] def empty_compact_observations(cam_key, logits_dim): return { "camera": cam_key, "point_indices": np.empty((0,), dtype=np.int64), "pixel_coords": np.empty((0, 2), dtype=np.int32), "visibility": np.empty((0, 6), dtype=np.float32), "ground_truth": np.empty((0,), dtype=np.int16), "logit_vectors": np.empty((0, logits_dim), dtype=np.float32), } def filter_compact_observations_to_valid_points(observations, valid_point_mask): if not observations or "point_indices" not in observations: return observations if np.any(valid_point_mask): return { "camera": observations["camera"], "point_indices": observations["point_indices"][valid_point_mask], "pixel_coords": observations["pixel_coords"][valid_point_mask], "visibility": observations["visibility"][valid_point_mask], "ground_truth": observations["ground_truth"][valid_point_mask], "logit_vectors": observations["logit_vectors"][valid_point_mask], } return empty_compact_observations(observations["camera"], observations["logit_vectors"].shape[1]) def project_points_raw(cam_params, global_params, subscale_factor, target_device=None): params = global_params.copy() params.update(cam_params) f_proj = _get_projection_fn(params.keys()) if target_device is None: x, y, z, in_bounds = f_proj(params) else: with jax.default_device(target_device): x, y, z, in_bounds = f_proj(params) full_i = x.astype(int) full_j = y.astype(int) i = (x / subscale_factor).astype(int) j = (y / subscale_factor).astype(int) return x, y, z, in_bounds, full_i, full_j, i, j def project_points_for_camera(cam_params, global_params, subscale_factor, buffer_size, threshold, target_device=None): _, _, z, in_bounds, full_i, full_j, i, j = project_points_raw( cam_params, global_params, subscale_factor, target_device=target_device, ) if jnp.sum(in_bounds) == 0: return None, None, None, None, None, None depth_map = jnp.full((global_params["height"] // subscale_factor, global_params["width"] // subscale_factor), jnp.inf) compute_depth_map_jit = _get_depth_fn(buffer_size, threshold) if target_device is None: _, visible_in_bounds = compute_depth_map_jit(i[in_bounds], j[in_bounds], z[in_bounds], depth_map) else: with jax.default_device(target_device): _, visible_in_bounds = compute_depth_map_jit(i[in_bounds], j[in_bounds], z[in_bounds], depth_map) return in_bounds, visible_in_bounds, full_i, full_j, i, j def rot_x_np(omega): return np.array( [[1.0, 0.0, 0.0], [0.0, np.cos(omega), np.sin(omega)], [0.0, -np.sin(omega), np.cos(omega)]], dtype=np.float32, ) def rot_y_np(phi): return np.array( [[np.cos(phi), 0.0, -np.sin(phi)], [0.0, 1.0, 0.0], [np.sin(phi), 0.0, np.cos(phi)]], dtype=np.float32, ) def rot_z_np(kappa): return np.array( [[np.cos(kappa), np.sin(kappa), 0.0], [-np.sin(kappa), np.cos(kappa), 0.0], [0.0, 0.0, 1.0]], dtype=np.float32, ) def rot_zyx_np(omega, phi, kappa): return rot_z_np(kappa) @ rot_y_np(phi) @ rot_x_np(omega) def prefilter_points_for_camera(points_xyz, cam_params, calibration_params, cfg): if not cfg["load"].get("prefilter_enabled", True): return np.ones(points_xyz.shape[0], dtype=bool) camera_position = np.array([cam_params["Xs"], cam_params["Ys"], cam_params["Zs"]], dtype=np.float32) rotation = rot_zyx_np(cam_params["omega"], cam_params["phi"], cam_params["kappa"]) relative = points_xyz - camera_position[None, :] rms = relative @ rotation.T z_cam = -rms[:, 2] mask = z_cam > 0.0 if not np.any(mask): return mask max_distance = cfg["load"].get("max_distance") if max_distance is not None: mask &= np.linalg.norm(relative, axis=1) <= float(max_distance) if not np.any(mask): return mask margin = float(cfg["load"].get("prefilter_margin", 0.15)) width = float(calibration_params["width"]) height = float(calibration_params["height"]) focal = float(calibration_params["f"]) w_half = (width / focal / 2.0) * (1.0 + margin) h_half = (height / focal / 2.0) * (1.0 + margin) x_norm = np.zeros(points_xyz.shape[0], dtype=np.float32) y_norm = np.zeros(points_xyz.shape[0], dtype=np.float32) valid = mask x_norm[valid] = -rms[valid, 0] / rms[valid, 2] y_norm[valid] = -rms[valid, 1] / rms[valid, 2] return valid & (x_norm >= -w_half) & (x_norm < w_half) & (y_norm >= -h_half) & (y_norm < h_half) def load_logits_image(logits_path, expected_shape=None): logits_img = np.load(logits_path) if logits_img.ndim != 3: raise ValueError(f"Expected logits image of shape [H, W, C], got {logits_img.shape}") if expected_shape is not None and logits_img.shape[:2] != expected_shape: raise ValueError(f"Logits/image size mismatch for {logits_path}: {logits_img.shape} vs {expected_shape}") return logits_img.astype(np.float32) def build_compact_observations( cam_key, points_for_normals, normals, points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cam_params, in_bounds, visible_in_bounds, full_i, full_j, logits_i, logits_j, prefilter_mask, allow_missing_ground_truth=False, global_indices=None, remap_ground_truth=True, ): logits_dim = logits_img.shape[2] in_bounds_indices = np.flatnonzero(in_bounds) visible_indices = in_bounds_indices[np.asarray(visible_in_bounds, dtype=bool)] if visible_indices.size == 0: return empty_compact_observations(cam_key, logits_dim), 0, 0 filtered_points = points_for_normals[prefilter_mask] filtered_normals = normals[prefilter_mask] filtered_local_indices = np.nonzero(prefilter_mask)[0] crit = compute_visibility_criteria_numpy(filtered_points[visible_indices], filtered_normals[visible_indices], cam_params) angles = np.degrees(crit[:, 0]).astype(np.float32) distances = np.round(crit[:, 1], 3).astype(np.float32) local_chunk_indices = filtered_local_indices[visible_indices] original_point_indices = ( local_chunk_indices.astype(np.int64) if global_indices is None else global_indices[local_chunk_indices].astype(np.int64) ) u_full = full_i[visible_indices].astype(np.int32) v_full = full_j[visible_indices].astype(np.int32) u_logit = logits_i[visible_indices].astype(np.int32) v_logit = logits_j[visible_indices].astype(np.int32) contrast_values = contrast_map[v_full, u_full].astype(np.float32) blur_values = blur_map[v_full, u_full].astype(np.float32) snr_values = snr_map[v_full, u_full].astype(np.float32) saturation_values = saturation_map[v_full, u_full].astype(np.float32) logits_values = logits_img[v_logit, u_logit, :].astype(np.float32) if points["GT_AVAILABLE"]: if remap_ground_truth: class_ids = map_main_class_indices(points["GT"][original_point_indices]) else: class_ids = np.asarray(points["GT"][original_point_indices], dtype=np.int16) valid_mask = class_ids != 255 else: class_ids = np.full(original_point_indices.shape[0], -1, dtype=np.int16) valid_mask = np.ones(original_point_indices.shape[0], dtype=bool) if allow_missing_ground_truth else np.zeros( original_point_indices.shape[0], dtype=bool ) valid_indices = np.nonzero(valid_mask)[0] if valid_indices.size == 0: return empty_compact_observations(cam_key, logits_dim), int(visible_indices.size), 0 point_indices = original_point_indices[valid_indices] order = np.argsort(point_indices, kind="mergesort") sorted_indices = valid_indices[order] compact_observations = { "camera": cam_key, "point_indices": original_point_indices[sorted_indices].astype(np.int64), "pixel_coords": np.stack((u_full[sorted_indices], v_full[sorted_indices]), axis=1).astype(np.int32), "visibility": np.stack( ( angles[sorted_indices], distances[sorted_indices], contrast_values[sorted_indices], blur_values[sorted_indices], snr_values[sorted_indices], saturation_values[sorted_indices], ), axis=1, ).astype(np.float32), "ground_truth": class_ids[sorted_indices].astype(np.int16), "logit_vectors": logits_values[sorted_indices].astype(np.float32), } selected_points = int(np.unique(compact_observations["point_indices"]).size) return compact_observations, int(visible_indices.size), selected_points def build_camera_observations( cam_key, cam_params, calibration_params, points_for_normals, normals, points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, allow_missing_ground_truth=False, global_indices=None, target_device=None, remap_ground_truth=True, ): step_start = time.perf_counter() prefilter_mask = prefilter_points_for_camera(points_for_normals, cam_params, calibration_params, cfg) if not np.any(prefilter_mask): return {}, { "prefilter_s": time.perf_counter() - step_start, "projection_s": 0.0, "criteria_s": 0.0, "visible_points": 0, "selected_points": 0, "prefilter_kept": 0, } filtered_params = calibration_params.copy() filtered_params["X"] = calibration_params["X"][prefilter_mask] filtered_params["Y"] = calibration_params["Y"][prefilter_mask] filtered_params["Z"] = calibration_params["Z"][prefilter_mask] prefilter_s = time.perf_counter() - step_start projection_start = time.perf_counter() in_bounds_jax, visible_in_bounds_jax, full_i_jax, full_j_jax, logits_i_jax, logits_j_jax = project_points_for_camera( cam_params, filtered_params, cfg["load"]["subscale"], cfg["load"]["buffer_size"], cfg["load"]["threshold"], target_device=target_device, ) if in_bounds_jax is None: return {}, { "prefilter_s": prefilter_s, "projection_s": time.perf_counter() - projection_start, "criteria_s": 0.0, "visible_points": 0, "selected_points": 0, "prefilter_kept": int(prefilter_mask.sum()), } in_bounds, visible_in_bounds, full_i, full_j, logits_i, logits_j = jax.device_get( (in_bounds_jax, visible_in_bounds_jax, full_i_jax, full_j_jax, logits_i_jax, logits_j_jax) ) if np.count_nonzero(np.asarray(visible_in_bounds, dtype=bool)) == 0: return {}, { "prefilter_s": prefilter_s, "projection_s": time.perf_counter() - projection_start, "criteria_s": 0.0, "visible_points": 0, "selected_points": 0, "prefilter_kept": int(prefilter_mask.sum()), } projection_s = time.perf_counter() - projection_start criteria_start = time.perf_counter() compact_observations, visible_points, selected_points = build_compact_observations( cam_key, points_for_normals, normals, points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cam_params, in_bounds, visible_in_bounds, full_i, full_j, logits_i, logits_j, prefilter_mask, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices, remap_ground_truth=remap_ground_truth, ) return compact_observations, { "prefilter_s": prefilter_s, "projection_s": projection_s, "criteria_s": time.perf_counter() - criteria_start, "visible_points": visible_points, "selected_points": selected_points, "prefilter_kept": int(prefilter_mask.sum()), } def process_camera_chunk( start_idx, end_idx, calibration_params, points, coords, normals, points_gt, logits_img, contrast_map, blur_map, snr_map, saturation_map, cam_key, cam_params, cfg, allow_missing_ground_truth, global_indices, target_device, ): chunk_params = calibration_params.copy() chunk_params["X"] = points["X"][start_idx:end_idx] chunk_params["Y"] = points["Y"][start_idx:end_idx] chunk_params["Z"] = points["Z"][start_idx:end_idx] chunk_points = { "GT": points_gt, "GT_AVAILABLE": points_gt is not None, } observations, stats = build_camera_observations( cam_key, cam_params, chunk_params, coords[start_idx:end_idx], normals[start_idx:end_idx], chunk_points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices[start_idx:end_idx], target_device=target_device, ) return observations, stats def compute_camera_global_depth_map( cam_key, cam_params, calibration_params, points_xyz, cfg, chunk_size, target_device=None, ): subscale_factor = int(cfg["load"]["subscale"]) buffer_size = int(cfg["load"]["buffer_size"]) threshold = float(cfg["load"]["threshold"]) depth_map = jnp.full( ( int(calibration_params["height"]) // subscale_factor, int(calibration_params["width"]) // subscale_factor, ), jnp.inf, ) compute_depth_map_jit = _get_depth_fn(buffer_size, threshold) total_prefilter_s = 0.0 total_projection_s = 0.0 total_prefilter_kept = 0 for start_idx in range(0, points_xyz.shape[0], chunk_size): end_idx = start_idx + chunk_size step_start = time.perf_counter() chunk_points_xyz = points_xyz[start_idx:end_idx] prefilter_mask = prefilter_points_for_camera(chunk_points_xyz, cam_params, calibration_params, cfg) total_prefilter_s += time.perf_counter() - step_start if not np.any(prefilter_mask): continue filtered_params = calibration_params.copy() filtered_params["X"] = calibration_params["X"][start_idx:end_idx][prefilter_mask] filtered_params["Y"] = calibration_params["Y"][start_idx:end_idx][prefilter_mask] filtered_params["Z"] = calibration_params["Z"][start_idx:end_idx][prefilter_mask] total_prefilter_kept += int(prefilter_mask.sum()) projection_start = time.perf_counter() _, _, z, in_bounds, _, _, i, j = project_points_raw( cam_params, filtered_params, subscale_factor, target_device=target_device, ) if jnp.sum(in_bounds) != 0: if target_device is None: depth_map, _ = compute_depth_map_jit(i[in_bounds], j[in_bounds], z[in_bounds], depth_map) else: with jax.default_device(target_device): depth_map, _ = compute_depth_map_jit(i[in_bounds], j[in_bounds], z[in_bounds], depth_map) total_projection_s += time.perf_counter() - projection_start logging.debug("Built global depth map for camera %s", cam_key) return depth_map, total_prefilter_s, total_projection_s, total_prefilter_kept def build_camera_observations_global_depth( cam_key, cam_params, calibration_params, points_for_normals, normals, points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, global_depth_map, allow_missing_ground_truth=False, global_indices=None, target_device=None, ): step_start = time.perf_counter() prefilter_mask = prefilter_points_for_camera(points_for_normals, cam_params, calibration_params, cfg) if not np.any(prefilter_mask): return {}, { "prefilter_s": time.perf_counter() - step_start, "projection_s": 0.0, "criteria_s": 0.0, "visible_points": 0, "selected_points": 0, "prefilter_kept": 0, } filtered_params = calibration_params.copy() filtered_params["X"] = calibration_params["X"][prefilter_mask] filtered_params["Y"] = calibration_params["Y"][prefilter_mask] filtered_params["Z"] = calibration_params["Z"][prefilter_mask] prefilter_s = time.perf_counter() - step_start projection_start = time.perf_counter() _, _, z_jax, in_bounds_jax, full_i_jax, full_j_jax, logits_i_jax, logits_j_jax = project_points_raw( cam_params, filtered_params, cfg["load"]["subscale"], target_device=target_device, ) if jnp.sum(in_bounds_jax) == 0: return {}, { "prefilter_s": prefilter_s, "projection_s": time.perf_counter() - projection_start, "criteria_s": 0.0, "visible_points": 0, "selected_points": 0, "prefilter_kept": int(prefilter_mask.sum()), } compute_depth_map_jit = _get_depth_fn(cfg["load"]["buffer_size"], cfg["load"]["threshold"]) if target_device is None: _, visible_in_bounds_jax = compute_depth_map_jit( logits_i_jax[in_bounds_jax], logits_j_jax[in_bounds_jax], z_jax[in_bounds_jax], global_depth_map, ) else: with jax.default_device(target_device): _, visible_in_bounds_jax = compute_depth_map_jit( logits_i_jax[in_bounds_jax], logits_j_jax[in_bounds_jax], z_jax[in_bounds_jax], global_depth_map, ) in_bounds, visible_in_bounds, full_i, full_j, logits_i, logits_j = jax.device_get( (in_bounds_jax, visible_in_bounds_jax, full_i_jax, full_j_jax, logits_i_jax, logits_j_jax) ) projection_s = time.perf_counter() - projection_start criteria_start = time.perf_counter() compact_observations, visible_points, selected_points = build_compact_observations( cam_key, points_for_normals, normals, points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cam_params, in_bounds, visible_in_bounds, full_i, full_j, logits_i, logits_j, prefilter_mask, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices, ) return compact_observations, { "prefilter_s": prefilter_s, "projection_s": projection_s, "criteria_s": time.perf_counter() - criteria_start, "visible_points": visible_points, "selected_points": selected_points, "prefilter_kept": int(prefilter_mask.sum()), } def process_camera_items_for_tile_global_depth( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ): output_mode = cfg.get("runtime", {}).get("tile_output_format", "h5_then_pt") use_dense = output_mode in {"pt_direct", "pt_sharded_direct"} all_observations = None if use_dense else {} extension = cfg["logits"].get("extension", ".npy") max_views = int(cfg["selection"]["max_views"]) selection_strategy = cfg["selection"]["strategy"] chunk_size = int(cfg["load"].get("point_chunk_size", len(points["X"]))) total_points = len(points["X"]) global_indices = np.arange(total_points, dtype=np.int64) pad_count = (-total_points) % chunk_size padded_points = {key: value.copy() if isinstance(value, np.ndarray) else value for key, value in points.items()} padded_coords = coords.copy() padded_normals = normals.copy() if pad_count: padded_points["X"] = np.pad(padded_points["X"], (0, pad_count), mode="constant", constant_values=0) padded_points["Y"] = np.pad(padded_points["Y"], (0, pad_count), mode="constant", constant_values=0) padded_points["Z"] = np.pad(padded_points["Z"], (0, pad_count), mode="constant", constant_values=0) padded_coords = np.pad(padded_coords, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) padded_normals = np.pad(padded_normals, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) global_indices = np.pad(global_indices, (0, pad_count), mode="constant", constant_values=-1) padded_total_points = len(padded_points["X"]) logging.info( "Using JAX backend=%s for tile %s with chunk_size=%s max_distance=%s visibility_mode=two_pass_global_depth", jax.default_backend(), tile_info["tile"], chunk_size, cfg["load"].get("max_distance"), ) for cam_key, cam_params in tqdm(camera_items, desc=f"{tile_info['tile']} cameras", unit="camera", leave=False): camera_start = time.perf_counter() logits_path = logits_dir / f"{cam_key}{extension}" image_path = Path(tile_info["image_dir"]) / f"{cam_key}.JPG" if not logits_path.exists(): logging.warning("Missing logits for %s in tile %s", cam_key, tile_info["tile"]) continue if not image_path.exists(): logging.warning("Missing image for %s in tile %s", cam_key, tile_info["tile"]) continue io_start = time.perf_counter() image = cv.imread(str(image_path), cv.IMREAD_COLOR) if image is None: raise FileNotFoundError(f"Unable to read image {image_path}") expected_logits_shape = (image.shape[0] // cfg["load"]["subscale"], image.shape[1] // cfg["load"]["subscale"]) logits_img = load_logits_image(logits_path, expected_shape=expected_logits_shape) if use_dense and all_observations is None: all_observations = create_dense_observations( total_points, max_views, logits_img.shape[2], cfg["data"]["vmin"], cfg["data"]["vmax"], get_storage_config(cfg)["visibility_quant_max"], get_storage_config(cfg)["logits_quant_max"], ) io_s = time.perf_counter() - io_start maps_start = time.perf_counter() contrast_map, blur_map, snr_map, saturation_map = compute_criteria_maps(image) maps_s = time.perf_counter() - maps_start depth_map, pass1_prefilter_s, pass1_projection_s, pass1_prefilter_kept = compute_camera_global_depth_map( cam_key, cam_params, calibration_params, padded_coords, cfg, chunk_size, ) chunk_prefilter_s = 0.0 chunk_projection_s = 0.0 chunk_criteria_s = 0.0 chunk_merge_s = 0.0 prefilter_kept = 0 visible_points = 0 selected_points = 0 num_chunks = padded_total_points // chunk_size for start_idx in range(0, padded_total_points, chunk_size): end_idx = start_idx + chunk_size chunk_params = calibration_params.copy() chunk_params["X"] = padded_points["X"][start_idx:end_idx] chunk_params["Y"] = padded_points["Y"][start_idx:end_idx] chunk_params["Z"] = padded_points["Z"][start_idx:end_idx] observations, stats = build_camera_observations_global_depth( cam_key, cam_params, chunk_params, padded_coords[start_idx:end_idx], padded_normals[start_idx:end_idx], padded_points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, depth_map, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices[start_idx:end_idx], ) if pad_count and end_idx > total_points and observations and "point_indices" in observations: observations = filter_compact_observations_to_valid_points(observations, observations["point_indices"] >= 0) chunk_prefilter_s += stats["prefilter_s"] chunk_projection_s += stats["projection_s"] chunk_criteria_s += stats["criteria_s"] prefilter_kept += stats["prefilter_kept"] visible_points += stats["visible_points"] selected_points += stats["selected_points"] merge_start = time.perf_counter() if use_dense: merge_compact_observations_dense(all_observations, observations, max_views, selection_strategy) else: merge_compact_observations_bounded(all_observations, observations, max_views, selection_strategy) chunk_merge_s += time.perf_counter() - merge_start logging.info( "Tile %s camera %s done in %.2fs | io=%.2fs maps=%.2fs pass1_prefilter=%.2fs pass1_projection=%.2fs " "prefilter=%.2fs projection=%.2fs criteria=%.2fs merge=%.2fs prefilter_kept=%s pass1_prefilter_kept=%s visible=%s " "selected=%s chunks=%s", tile_info["tile"], cam_key, time.perf_counter() - camera_start, io_s, maps_s, pass1_prefilter_s, pass1_projection_s, chunk_prefilter_s, chunk_projection_s, chunk_criteria_s, chunk_merge_s, prefilter_kept, pass1_prefilter_kept, visible_points, selected_points, num_chunks, ) if use_dense and all_observations is None: all_observations = create_dense_observations( total_points, max_views, 0, cfg["data"]["vmin"], cfg["data"]["vmax"], get_storage_config(cfg)["visibility_quant_max"], get_storage_config(cfg)["logits_quant_max"], ) return all_observations def select_observations(observations, max_views, strategy): if strategy == "topk_distance": ranked = sorted(observations, key=lambda obs: (obs["visibility"][1], obs["visibility"][0])) else: ranked = observations return ranked[:max_views] def trim_bucket(bucket, max_views, strategy): if len(bucket["camera"]) <= max_views: return if strategy == "topk_distance": order = sorted( range(len(bucket["camera"])), key=lambda idx: (bucket["visibility"][idx][1], bucket["visibility"][idx][0]), )[:max_views] else: order = list(range(max_views)) bucket["camera"] = [bucket["camera"][idx] for idx in order] bucket["pixel_coords"] = [bucket["pixel_coords"][idx] for idx in order] bucket["visibility"] = [bucket["visibility"][idx] for idx in order] bucket["logit_vectors"] = [bucket["logit_vectors"][idx] for idx in order] def merge_compact_observations_bounded(all_observations, compact_observations, max_views, strategy): if not compact_observations or "point_indices" not in compact_observations: return point_indices = compact_observations["point_indices"] if point_indices.size == 0: return group_starts = np.r_[0, np.flatnonzero(np.diff(point_indices)) + 1] group_ends = np.r_[group_starts[1:], point_indices.size] camera_key = compact_observations["camera"] for start, end in zip(group_starts, group_ends): point_idx = int(point_indices[start]) bucket = all_observations.setdefault( point_idx, { "ground_truth": None, "camera": [], "pixel_coords": [], "visibility": [], "logit_vectors": [], }, ) gt_slice = compact_observations["ground_truth"][start:end] valid_gt = gt_slice[gt_slice >= 0] if bucket["ground_truth"] is None and valid_gt.size > 0: bucket["ground_truth"] = int(valid_gt[0]) count = end - start bucket["camera"].extend([camera_key] * count) bucket["pixel_coords"].extend(compact_observations["pixel_coords"][start:end].tolist()) bucket["visibility"].extend(compact_observations["visibility"][start:end].tolist()) bucket["logit_vectors"].extend([row for row in compact_observations["logit_vectors"][start:end]]) trim_bucket(bucket, max_views, strategy) def merge_worker_observations_bounded(all_observations, worker_observations, max_views, strategy): for point_idx, worker_bucket in worker_observations.items(): bucket = all_observations.setdefault( point_idx, { "ground_truth": None, "camera": [], "pixel_coords": [], "visibility": [], "logit_vectors": [], }, ) if bucket["ground_truth"] is None and worker_bucket.get("ground_truth") is not None: bucket["ground_truth"] = int(worker_bucket["ground_truth"]) bucket["camera"].extend(worker_bucket["camera"]) bucket["pixel_coords"].extend(worker_bucket["pixel_coords"]) bucket["visibility"].extend(worker_bucket["visibility"]) bucket["logit_vectors"].extend(worker_bucket["logit_vectors"]) trim_bucket(bucket, max_views, strategy) def create_dense_observations(num_points, max_views, logits_dim, visibility_vmin, visibility_vmax, visibility_quant_max, logits_quant_max): return { "mode": "dense", "counts": np.zeros(num_points, dtype=np.uint8), "ground_truth": np.full(num_points, -1, dtype=np.int16), "pixel_coords": np.zeros((num_points, max_views, 2), dtype=np.uint16), "visibility": np.zeros((num_points, max_views, 6), dtype=np.uint16), "logit_vectors": np.zeros((num_points, max_views, logits_dim), dtype=np.uint8), "visibility_vmin": np.asarray(visibility_vmin, dtype=np.float32), "visibility_vmax": np.asarray(visibility_vmax, dtype=np.float32), "visibility_quant_max": int(visibility_quant_max), "logits_quant_max": int(logits_quant_max), } def count_nonempty_observations(observations): if isinstance(observations, dict) and observations.get("mode") == "dense": return int(np.count_nonzero(observations["counts"])) return len(observations) def slice_point_data(points, coords, normals, start_idx, end_idx): sliced_points = {"GT_AVAILABLE": points["GT_AVAILABLE"]} for key in ("X", "Y", "Z"): sliced_points[key] = points[key][start_idx:end_idx] sliced_points["GT"] = None if points["GT"] is None else points["GT"][start_idx:end_idx] return sliced_points, coords[start_idx:end_idx], normals[start_idx:end_idx] def _sort_dense_point_slots(dense_observations, point_idx): count = int(dense_observations["counts"][point_idx]) if count <= 1: return order = np.lexsort( ( dense_observations["visibility"][point_idx, :count, 0], dense_observations["visibility"][point_idx, :count, 1], ) ) dense_observations["pixel_coords"][point_idx, :count] = dense_observations["pixel_coords"][point_idx, :count][order] dense_observations["visibility"][point_idx, :count] = dense_observations["visibility"][point_idx, :count][order] dense_observations["logit_vectors"][point_idx, :count] = dense_observations["logit_vectors"][point_idx, :count][order] def _select_topk_indices(visibility, max_views, strategy): if visibility.shape[0] <= max_views: order = np.lexsort((visibility[:, 0], visibility[:, 1])) return order if strategy == "topk_distance": order = np.lexsort((visibility[:, 0], visibility[:, 1])) return order[:max_views] return np.arange(max_views, dtype=np.int64) def _merge_dense_point_arrays(dense_observations, point_idx, new_pixel_coords, new_visibility, new_logit_vectors, max_views, strategy): current_count = int(dense_observations["counts"][point_idx]) if current_count > 0: merged_pixel_coords = np.concatenate( (dense_observations["pixel_coords"][point_idx, :current_count], new_pixel_coords), axis=0, ) merged_visibility = np.concatenate( (dense_observations["visibility"][point_idx, :current_count], new_visibility), axis=0, ) merged_logit_vectors = np.concatenate( (dense_observations["logit_vectors"][point_idx, :current_count], new_logit_vectors), axis=0, ) else: merged_pixel_coords = new_pixel_coords merged_visibility = new_visibility merged_logit_vectors = new_logit_vectors keep = _select_topk_indices(merged_visibility, max_views, strategy) keep_count = len(keep) dense_observations["pixel_coords"][point_idx, :keep_count] = merged_pixel_coords[keep] dense_observations["visibility"][point_idx, :keep_count] = merged_visibility[keep] dense_observations["logit_vectors"][point_idx, :keep_count] = merged_logit_vectors[keep] dense_observations["counts"][point_idx] = keep_count def merge_compact_observations_dense(dense_observations, compact_observations, max_views, strategy): if not compact_observations or "point_indices" not in compact_observations: return point_indices = compact_observations["point_indices"] if point_indices.size == 0: return quantized_visibility = quantize_visibility( compact_observations["visibility"], dense_observations["visibility_vmin"], dense_observations["visibility_vmax"], dense_observations["visibility_quant_max"], ) quantized_logits = quantize_logits( compact_observations["logit_vectors"], dense_observations["logits_quant_max"], ) group_starts = np.r_[0, np.flatnonzero(np.diff(point_indices)) + 1] group_ends = np.r_[group_starts[1:], point_indices.size] for start, end in zip(group_starts, group_ends): point_idx = int(point_indices[start]) gt_slice = compact_observations["ground_truth"][start:end] valid_gt = gt_slice[gt_slice >= 0] if valid_gt.size > 0 and dense_observations["ground_truth"][point_idx] < 0: dense_observations["ground_truth"][point_idx] = int(valid_gt[0]) _merge_dense_point_arrays( dense_observations, point_idx, compact_observations["pixel_coords"][start:end].astype(np.uint16, copy=False), quantized_visibility[start:end], quantized_logits[start:end], max_views, strategy, ) def merge_dense_observations(target, source, max_views, strategy): source_counts = source["counts"] for point_idx in np.flatnonzero(source_counts): point_idx = int(point_idx) src_count = int(source_counts[point_idx]) if src_count == 0: continue if target["ground_truth"][point_idx] < 0 and source["ground_truth"][point_idx] >= 0: target["ground_truth"][point_idx] = source["ground_truth"][point_idx] _merge_dense_point_arrays( target, point_idx, source["pixel_coords"][point_idx, :src_count], source["visibility"][point_idx, :src_count], source["logit_vectors"][point_idx, :src_count], max_views, strategy, ) def create_compact_payload_template(cfg, logits_dim): storage_cfg = get_storage_config(cfg) return { "point_indices": [], "visibility": [], "logit_vectors": [], "ground_truth": [], "visibility_vmin": np.asarray(cfg["data"]["vmin"], dtype=np.float32), "visibility_vmax": np.asarray(cfg["data"]["vmax"], dtype=np.float32), "visibility_quant_max": int(storage_cfg["visibility_quant_max"]), "logits_quant_max": int(storage_cfg["logits_quant_max"]), "logits_dim": int(logits_dim), } def append_compact_payload(payload, compact_observations): if not compact_observations or "point_indices" not in compact_observations: return if compact_observations["point_indices"].size == 0: return payload["point_indices"].append(compact_observations["point_indices"].astype(np.uint32, copy=False)) payload["visibility"].append( quantize_visibility( compact_observations["visibility"], payload["visibility_vmin"], payload["visibility_vmax"], payload["visibility_quant_max"], ) ) payload["logit_vectors"].append( quantize_logits(compact_observations["logit_vectors"], payload["logits_quant_max"]) ) payload["ground_truth"].append(compact_observations["ground_truth"].astype(np.int16, copy=False)) def finalize_compact_payload(payload): if payload is None: return None logits_dim = payload["logits_dim"] if payload["point_indices"]: return { "point_indices": np.concatenate(payload["point_indices"], axis=0), "visibility": np.concatenate(payload["visibility"], axis=0), "logit_vectors": np.concatenate(payload["logit_vectors"], axis=0), "ground_truth": np.concatenate(payload["ground_truth"], axis=0), "visibility_vmin": payload["visibility_vmin"], "visibility_vmax": payload["visibility_vmax"], "visibility_quant_max": payload["visibility_quant_max"], "logits_quant_max": payload["logits_quant_max"], } return { "point_indices": np.empty((0,), dtype=np.uint32), "visibility": np.empty((0, 6), dtype=np.uint16), "logit_vectors": np.empty((0, logits_dim), dtype=np.uint8), "ground_truth": np.empty((0,), dtype=np.int16), "visibility_vmin": payload["visibility_vmin"], "visibility_vmax": payload["visibility_vmax"], "visibility_quant_max": payload["visibility_quant_max"], "logits_quant_max": payload["logits_quant_max"], } def point_ground_truth(selected): labels = [obs["ground_truth"] for obs in selected if obs["ground_truth"] is not None] if not labels: return None values, counts = np.unique(np.asarray(labels, dtype=np.int64), return_counts=True) return int(values[np.argmax(counts)]) def get_storage_config(cfg): storage_cfg = cfg.get("storage", {}) return { "coord_scale": float(storage_cfg.get("coord_scale", 0.001)), "pixel_coords_dtype": storage_cfg.get("pixel_coords_dtype", "uint16"), "visibility_quant_max": int(storage_cfg.get("visibility_quant_max", 65535)), "logits_quant_max": int(storage_cfg.get("logits_quant_max", 255)), "compression": storage_cfg.get("h5_compression", "lzf"), } def quantize_coordinates(coords_row, coord_scale): return np.rint(coords_row / coord_scale).astype(np.int32) def quantize_visibility(visibility, vmin, vmax, quant_max): vmin_arr = np.asarray(vmin, dtype=np.float32) vmax_arr = np.asarray(vmax, dtype=np.float32) clipped = np.clip(visibility.astype(np.float32), vmin_arr, vmax_arr) normalized = (clipped - vmin_arr) / np.maximum(vmax_arr - vmin_arr, 1e-8) return np.rint(normalized * quant_max).astype(np.uint16) def quantize_logits(logits, quant_max): logits = np.asarray(logits) if np.issubdtype(logits.dtype, np.integer) or float(np.nanmax(logits)) > 1.0: return np.clip(logits, 0, quant_max).astype(np.uint8) clipped = np.clip(logits.astype(np.float32), 0.0, 1.0) return np.rint(clipped * quant_max).astype(np.uint8) def write_tile_h5(output_h5_file, tile_name, split_name, coords, all_observations, cfg): output_h5_file.parent.mkdir(parents=True, exist_ok=True) min_views = cfg["selection"]["min_views"] max_views = cfg["selection"]["max_views"] strategy = cfg["selection"]["strategy"] total_points = len(all_observations) storage_cfg = get_storage_config(cfg) compression = storage_cfg["compression"] vmin = cfg["data"]["vmin"] vmax = cfg["data"]["vmax"] logging.info( "Writing tile %s to %s with %s aggregated point buckets", tile_name, output_h5_file, total_points, ) with h5py.File(output_h5_file, "w") as handle: handle.attrs["tile"] = tile_name handle.attrs["split"] = split_name handle.attrs["feature_names"] = np.array(cfg["data"]["visibility_feature_names"], dtype="S") handle.attrs["coords_scale"] = storage_cfg["coord_scale"] handle.attrs["coords_offset"] = np.zeros(3, dtype=np.float32) handle.attrs["visibility_vmin"] = np.asarray(vmin, dtype=np.float32) handle.attrs["visibility_vmax"] = np.asarray(vmax, dtype=np.float32) handle.attrs["visibility_quant_max"] = storage_cfg["visibility_quant_max"] handle.attrs["logits_quant_max"] = storage_cfg["logits_quant_max"] points_group = handle.create_group("points") kept = 0 skipped = 0 for point_idx, observations in all_observations.items(): if len(observations["camera"]) < min_views: skipped += 1 continue gt = observations["ground_truth"] if gt is None: skipped += 1 continue group = points_group.create_group(str(point_idx)) group.create_dataset( "coordinates", data=quantize_coordinates(coords[point_idx], storage_cfg["coord_scale"]), dtype="int32", compression=compression, ) group.create_dataset( "image_ids", data=np.array(observations["camera"], dtype=h5py.string_dtype("utf-8")), compression=compression, ) group.create_dataset( "pixel_coords", data=np.array(observations["pixel_coords"], dtype=np.uint16), dtype=storage_cfg["pixel_coords_dtype"], compression=compression, ) group.create_dataset( "visibility", data=quantize_visibility( np.array(observations["visibility"], dtype=np.float32), vmin, vmax, storage_cfg["visibility_quant_max"], ), dtype="uint16", compression=compression, ) group.create_dataset( "logit_vectors", data=quantize_logits( np.array(observations["logit_vectors"], dtype=np.float32), storage_cfg["logits_quant_max"], ), dtype="uint8", compression=compression, ) group.create_dataset("ground_truth", data=np.uint8(gt), dtype="uint8") kept += 1 handle.attrs["num_points_kept"] = kept handle.attrs["num_points_skipped"] = skipped logging.info( "Finished writing tile %s: kept=%s skipped=%s", tile_name, kept, skipped, ) def process_camera_items_for_tile( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ): output_mode = cfg.get("runtime", {}).get("tile_output_format", "h5_then_pt") use_dense = output_mode in {"pt_direct", "pt_sharded_direct"} all_observations = None if use_dense else {} extension = cfg["logits"].get("extension", ".npy") max_views = int(cfg["selection"]["max_views"]) selection_strategy = cfg["selection"]["strategy"] chunk_size = int(cfg["load"].get("point_chunk_size", len(points["X"]))) total_points = len(points["X"]) global_indices = np.arange(total_points, dtype=np.int64) pad_count = (-total_points) % chunk_size padded_points = {key: value.copy() if isinstance(value, np.ndarray) else value for key, value in points.items()} padded_coords = coords.copy() padded_normals = normals.copy() if pad_count: padded_points["X"] = np.pad(padded_points["X"], (0, pad_count), mode="constant", constant_values=0) padded_points["Y"] = np.pad(padded_points["Y"], (0, pad_count), mode="constant", constant_values=0) padded_points["Z"] = np.pad(padded_points["Z"], (0, pad_count), mode="constant", constant_values=0) padded_coords = np.pad(padded_coords, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) padded_normals = np.pad(padded_normals, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) global_indices = np.pad(global_indices, (0, pad_count), mode="constant", constant_values=-1) padded_total_points = len(padded_points["X"]) logging.info( "Using JAX backend=%s for tile %s with chunk_size=%s max_distance=%s visibility_mode=chunk_local_depth", jax.default_backend(), tile_info["tile"], chunk_size, cfg["load"].get("max_distance"), ) for cam_key, cam_params in tqdm(camera_items, desc=f"{tile_info['tile']} cameras", unit="camera", leave=False): camera_start = time.perf_counter() logits_path = logits_dir / f"{cam_key}{extension}" image_path = Path(tile_info["image_dir"]) / f"{cam_key}.JPG" if not logits_path.exists(): logging.warning("Missing logits for %s in tile %s", cam_key, tile_info["tile"]) continue if not image_path.exists(): logging.warning("Missing image for %s in tile %s", cam_key, tile_info["tile"]) continue io_start = time.perf_counter() image = cv.imread(str(image_path), cv.IMREAD_COLOR) if image is None: raise FileNotFoundError(f"Unable to read image {image_path}") expected_logits_shape = (image.shape[0] // cfg["load"]["subscale"], image.shape[1] // cfg["load"]["subscale"]) logits_img = load_logits_image(logits_path, expected_shape=expected_logits_shape) if use_dense and all_observations is None: all_observations = create_dense_observations( total_points, max_views, logits_img.shape[2], cfg["data"]["vmin"], cfg["data"]["vmax"], get_storage_config(cfg)["visibility_quant_max"], get_storage_config(cfg)["logits_quant_max"], ) io_s = time.perf_counter() - io_start maps_start = time.perf_counter() contrast_map, blur_map, snr_map, saturation_map = compute_criteria_maps(image) maps_s = time.perf_counter() - maps_start chunk_prefilter_s = 0.0 chunk_projection_s = 0.0 chunk_criteria_s = 0.0 chunk_merge_s = 0.0 prefilter_kept = 0 visible_points = 0 selected_points = 0 num_chunks = padded_total_points // chunk_size for start_idx in range(0, padded_total_points, chunk_size): end_idx = start_idx + chunk_size chunk_params = calibration_params.copy() chunk_params["X"] = padded_points["X"][start_idx:end_idx] chunk_params["Y"] = padded_points["Y"][start_idx:end_idx] chunk_params["Z"] = padded_points["Z"][start_idx:end_idx] observations, stats = build_camera_observations( cam_key, cam_params, chunk_params, padded_coords[start_idx:end_idx], padded_normals[start_idx:end_idx], padded_points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices[start_idx:end_idx], ) if pad_count and end_idx > total_points and observations and "point_indices" in observations: observations = filter_compact_observations_to_valid_points(observations, observations["point_indices"] >= 0) chunk_prefilter_s += stats["prefilter_s"] chunk_projection_s += stats["projection_s"] chunk_criteria_s += stats["criteria_s"] prefilter_kept += stats["prefilter_kept"] visible_points += stats["visible_points"] selected_points += stats["selected_points"] merge_start = time.perf_counter() if use_dense: merge_compact_observations_dense(all_observations, observations, max_views, selection_strategy) else: merge_compact_observations_bounded(all_observations, observations, max_views, selection_strategy) chunk_merge_s += time.perf_counter() - merge_start logging.info( "Tile %s camera %s done in %.2fs | io=%.2fs maps=%.2fs prefilter=%.2fs projection=%.2fs criteria=%.2fs merge=%.2fs prefilter_kept=%s visible=%s selected=%s chunks=%s", tile_info["tile"], cam_key, time.perf_counter() - camera_start, io_s, maps_s, chunk_prefilter_s, chunk_projection_s, chunk_criteria_s, chunk_merge_s, prefilter_kept, visible_points, selected_points, num_chunks, ) if use_dense and all_observations is None: all_observations = create_dense_observations( total_points, max_views, 0, cfg["data"]["vmin"], cfg["data"]["vmax"], get_storage_config(cfg)["visibility_quant_max"], get_storage_config(cfg)["logits_quant_max"], ) return all_observations def process_camera_items_for_tile_compact( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ): extension = cfg["logits"].get("extension", ".npy") chunk_size = int(cfg["load"].get("point_chunk_size", len(points["X"]))) total_points = len(points["X"]) global_indices = np.arange(total_points, dtype=np.int64) pad_count = (-total_points) % chunk_size padded_points = {key: value.copy() if isinstance(value, np.ndarray) else value for key, value in points.items()} padded_coords = coords.copy() padded_normals = normals.copy() if pad_count: padded_points["X"] = np.pad(padded_points["X"], (0, pad_count), mode="constant", constant_values=0) padded_points["Y"] = np.pad(padded_points["Y"], (0, pad_count), mode="constant", constant_values=0) padded_points["Z"] = np.pad(padded_points["Z"], (0, pad_count), mode="constant", constant_values=0) padded_coords = np.pad(padded_coords, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) padded_normals = np.pad(padded_normals, ((0, pad_count), (0, 0)), mode="constant", constant_values=0) global_indices = np.pad(global_indices, (0, pad_count), mode="constant", constant_values=-1) padded_total_points = len(padded_points["X"]) compact_payload = None logging.info( "Using JAX backend=%s for tile %s with chunk_size=%s max_distance=%s visibility_mode=%s compact_reduce=true", jax.default_backend(), tile_info["tile"], chunk_size, cfg["load"].get("max_distance"), get_visibility_mode(cfg), ) for cam_key, cam_params in tqdm(camera_items, desc=f"{tile_info['tile']} cameras", unit="camera", leave=False): camera_start = time.perf_counter() logits_path = logits_dir / f"{cam_key}{extension}" image_path = Path(tile_info["image_dir"]) / f"{cam_key}.JPG" if not logits_path.exists(): logging.warning("Missing logits for %s in tile %s", cam_key, tile_info["tile"]) continue if not image_path.exists(): logging.warning("Missing image for %s in tile %s", cam_key, tile_info["tile"]) continue io_start = time.perf_counter() image = cv.imread(str(image_path), cv.IMREAD_COLOR) if image is None: raise FileNotFoundError(f"Unable to read image {image_path}") expected_logits_shape = (image.shape[0] // cfg["load"]["subscale"], image.shape[1] // cfg["load"]["subscale"]) logits_img = load_logits_image(logits_path, expected_shape=expected_logits_shape) if compact_payload is None: compact_payload = create_compact_payload_template(cfg, logits_img.shape[2]) io_s = time.perf_counter() - io_start maps_start = time.perf_counter() contrast_map, blur_map, snr_map, saturation_map = compute_criteria_maps(image) maps_s = time.perf_counter() - maps_start chunk_prefilter_s = 0.0 chunk_projection_s = 0.0 chunk_criteria_s = 0.0 chunk_append_s = 0.0 prefilter_kept = 0 visible_points = 0 selected_points = 0 num_chunks = padded_total_points // chunk_size for start_idx in range(0, padded_total_points, chunk_size): end_idx = start_idx + chunk_size chunk_params = calibration_params.copy() chunk_params["X"] = padded_points["X"][start_idx:end_idx] chunk_params["Y"] = padded_points["Y"][start_idx:end_idx] chunk_params["Z"] = padded_points["Z"][start_idx:end_idx] observations, stats = build_camera_observations( cam_key, cam_params, chunk_params, padded_coords[start_idx:end_idx], padded_normals[start_idx:end_idx], padded_points, logits_img, contrast_map, blur_map, snr_map, saturation_map, cfg, allow_missing_ground_truth=allow_missing_ground_truth, global_indices=global_indices[start_idx:end_idx], ) if pad_count and end_idx > total_points and observations and "point_indices" in observations: observations = filter_compact_observations_to_valid_points(observations, observations["point_indices"] >= 0) chunk_prefilter_s += stats["prefilter_s"] chunk_projection_s += stats["projection_s"] chunk_criteria_s += stats["criteria_s"] prefilter_kept += stats["prefilter_kept"] visible_points += stats["visible_points"] selected_points += stats["selected_points"] append_start = time.perf_counter() append_compact_payload(compact_payload, observations) chunk_append_s += time.perf_counter() - append_start logging.info( "Tile %s camera %s done in %.2fs | io=%.2fs maps=%.2fs prefilter=%.2fs projection=%.2fs criteria=%.2fs append=%.2fs prefilter_kept=%s visible=%s selected=%s chunks=%s", tile_info["tile"], cam_key, time.perf_counter() - camera_start, io_s, maps_s, chunk_prefilter_s, chunk_projection_s, chunk_criteria_s, chunk_append_s, prefilter_kept, visible_points, selected_points, num_chunks, ) if compact_payload is None: compact_payload = create_compact_payload_template(cfg, 0) return finalize_compact_payload(compact_payload) def get_visibility_mode(cfg): return cfg["load"].get("visibility_mode", "chunk_local_depth") def get_worker_output_mode(cfg): return cfg.get("runtime", {}).get("worker_output_mode", "inherit") def get_point_shard_size(cfg, total_points): return int(cfg.get("runtime", {}).get("point_shard_size", total_points)) def _launch_camera_workers( cfg, tile_info, camera_items, temp_dir, worker_count, shard_start=None, shard_end=None, ): worker_output_mode = get_worker_output_mode(cfg) camera_groups = [[] for _ in range(worker_count)] for idx, camera_item in enumerate(camera_items): camera_groups[idx % worker_count].append(camera_item[0]) processes = [] output_files = [] for worker_idx, camera_group in enumerate(camera_groups): if not camera_group: continue camera_list_file = temp_dir / f"worker_{worker_idx}_cameras.json" output_file = temp_dir / f"worker_{worker_idx}_observations.pt" stdout_file = temp_dir / f"worker_{worker_idx}.stdout.log" stderr_file = temp_dir / f"worker_{worker_idx}.stderr.log" camera_list_file.write_text(json.dumps(camera_group), encoding="utf-8") env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = str(worker_idx) env.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") cmd = [ sys.executable, str(Path(__file__).resolve().parents[1] / "scripts" / "run_batch_generation.py"), "--config", cfg["_config_path"], "--worker-mode", "camera-batch", "--tile", tile_info["tile"], "--camera-list-file", str(camera_list_file), "--worker-output", str(output_file), ] if shard_start is not None and shard_end is not None: cmd.extend(["--shard-start", str(shard_start), "--shard-end", str(shard_end)]) logging.info( "Launching worker %s for tile %s%s on GPU %s with %s cameras", worker_idx, tile_info["tile"], "" if shard_start is None else f" shard {shard_start}:{shard_end}", worker_idx, len(camera_group), ) stdout_handle = None stderr_handle = None if worker_output_mode == "files": stdout_handle = open(stdout_file, "w", encoding="utf-8") stderr_handle = open(stderr_file, "w", encoding="utf-8") process = subprocess.Popen( cmd, cwd=str(Path(__file__).resolve().parents[1]), env=env, stdout=stdout_handle, stderr=stderr_handle, ) elif worker_output_mode == "inherit": process = subprocess.Popen( cmd, cwd=str(Path(__file__).resolve().parents[1]), env=env, ) else: raise ValueError(f"Unsupported worker_output_mode: {worker_output_mode}") processes.append((worker_idx, process, stdout_handle, stderr_handle, stdout_file, stderr_file)) output_files.append(output_file) for worker_idx, process, stdout_handle, stderr_handle, stdout_file, stderr_file in processes: return_code = process.wait() if stdout_handle is not None: stdout_handle.close() if stderr_handle is not None: stderr_handle.close() if return_code != 0: log_hint = ( f" Worker logs: stdout={stdout_file} stderr={stderr_file}" if worker_output_mode == "files" else "" ) if return_code == -9: raise RuntimeError( f"Camera worker {worker_idx} was killed with code -9 for tile {tile_info['tile']}. " f"This usually indicates an out-of-memory kill (GPU or system RAM).{log_hint}" ) raise RuntimeError( f"Camera worker {worker_idx} failed for tile {tile_info['tile']} with code {return_code}.{log_hint}" ) return output_files def process_tile_sharded_to_pt( tile_info, cfg, calibration_params, points, coords, offset, normals, camera_items, logits_dir, allow_missing_ground_truth, ): max_views = int(cfg["selection"]["max_views"]) selection_strategy = cfg["selection"]["strategy"] shard_size = get_point_shard_size(cfg, len(points["X"])) total_points = len(points["X"]) parallel_camera_workers = int(cfg.get("runtime", {}).get("parallel_camera_workers", 1)) gpu_count = len(jax.devices("gpu")) if jax.default_backend() == "gpu" else 0 visibility_mode = get_visibility_mode(cfg) saved_paths = [] temp_root = Path(cfg["data"]["batches_dir"]) / ".camera_workers" / tile_info["tile"] for shard_idx, shard_start in enumerate(range(0, total_points, shard_size)): shard_end = min(shard_start + shard_size, total_points) shard_points, shard_coords, shard_normals = slice_point_data(points, coords, normals, shard_start, shard_end) shard_label = f"{tile_info['tile']}_shard_{shard_idx:04d}" logging.info( "Processing tile %s shard %s (%s:%s, %s points)", tile_info["tile"], shard_idx, shard_start, shard_end, shard_end - shard_start, ) use_parallel_workers = ( parallel_camera_workers > 1 and gpu_count > 0 and len(camera_items) > 1 and cfg.get("_worker_mode") is None ) if use_parallel_workers: worker_count = min(parallel_camera_workers, gpu_count, len(camera_items)) temp_dir = temp_root / f"shard_{shard_idx:04d}" temp_dir.mkdir(parents=True, exist_ok=True) output_files = _launch_camera_workers( cfg, tile_info, camera_items, temp_dir, worker_count, shard_start=shard_start, shard_end=shard_end, ) shard_observations = None for output_file in output_files: worker_payload = torch.load(output_file, map_location="cpu", weights_only=False) if shard_observations is None: shard_observations = worker_payload["observations"] else: merge_dense_observations(shard_observations, worker_payload["observations"], max_views, selection_strategy) if shard_observations is None: shard_observations = create_dense_observations( shard_end - shard_start, max_views, 0, cfg["data"]["vmin"], cfg["data"]["vmax"], get_storage_config(cfg)["visibility_quant_max"], get_storage_config(cfg)["logits_quant_max"], ) else: if visibility_mode == "two_pass_global_depth": shard_observations = process_camera_items_for_tile_global_depth( tile_info, camera_items, cfg, calibration_params, shard_points, shard_coords, shard_normals, logits_dir, allow_missing_ground_truth, ) elif visibility_mode == "chunk_local_depth": shard_observations = process_camera_items_for_tile( tile_info, camera_items, cfg, calibration_params, shard_points, shard_coords, shard_normals, logits_dir, allow_missing_ground_truth, ) else: raise ValueError(f"Unsupported visibility_mode: {visibility_mode}") shard_saved_paths = save_tile_observations_to_pt( cfg, shard_label, tile_info["split"], shard_coords, offset, shard_observations, cfg["data"]["batches_dir"], batch_size=cfg["data"]["dataloader_batch_size"], ) saved_paths.extend(shard_saved_paths) logging.info( "Finished tile %s shard %s: %s non-empty point buckets -> %s pt batches", tile_info["tile"], shard_idx, count_nonempty_observations(shard_observations), len(shard_saved_paths), ) return saved_paths def build_tile_dataset(tile_info, cfg, overwrite=False): prepared_root = Path(cfg["data"]["prepared_h5_dir"]) / tile_info["split"] output_h5_file = prepared_root / f"{tile_info['tile']}.h5" output_mode = cfg.get("runtime", {}).get("tile_output_format", "h5_then_pt") if output_mode == "pt_direct": pt_output_dir = Path(cfg["data"]["batches_dir"]) / tile_info["split"] existing_pt_batches = sorted(pt_output_dir.glob(f"{tile_info['tile']}_batch_*.pt")) if existing_pt_batches and not overwrite: logging.info("Skipping prepared tile %s", tile_info["tile"]) return [str(path) for path in existing_pt_batches] elif output_mode == "pt_compact_reduce": pt_output_dir = Path(cfg["data"]["batches_dir"]) / tile_info["split"] existing_pt_batches = sorted(pt_output_dir.glob(f"{tile_info['tile']}_batch_*.pt")) if existing_pt_batches and not overwrite: logging.info("Skipping prepared tile %s", tile_info["tile"]) return [str(path) for path in existing_pt_batches] elif output_mode == "pt_sharded_direct": pt_output_dir = Path(cfg["data"]["batches_dir"]) / tile_info["split"] existing_pt_batches = sorted(pt_output_dir.glob(f"{tile_info['tile']}_shard_*_batch_*.pt")) if existing_pt_batches and not overwrite: logging.info("Skipping prepared tile %s", tile_info["tile"]) return [str(path) for path in existing_pt_batches] elif output_h5_file.exists() and not overwrite: logging.info("Skipping prepared tile %s", tile_info["tile"]) return output_h5_file normals_path = Path(cfg["data"]["normals_dir"]) / tile_info["split"] / f"{tile_info['tile']}.h5" if not normals_path.exists(): raise FileNotFoundError(f"Normals not found for tile {tile_info['tile']}: {normals_path}") logits_dir = resolve_logits_dir(cfg, tile_info) if not logits_dir.exists(): raise FileNotFoundError(f"Logits directory not found for tile {tile_info['tile']}: {logits_dir}") calibration_params = load_calibration(tile_info["calibration_file"]) points, offset, _, coords = load_las_points(tile_info["lidar_path"]) camera_dict = load_camera_parameters(tile_info["camera_file"], offset) normals = load_normals_from_h5(normals_path) allow_missing_ground_truth = tile_info["split"] == "test" if not points["GT_AVAILABLE"] and not allow_missing_ground_truth: raise ValueError( f"Tile {tile_info['tile']} in split {tile_info['split']} has no ground_truth field in LAS." ) camera_items = list(camera_dict.items()) max_cameras = cfg.get("debug", {}).get("max_cameras_per_tile") if max_cameras is not None: camera_items = camera_items[: int(max_cameras)] if output_mode == "pt_sharded_direct": saved_paths = process_tile_sharded_to_pt( tile_info, cfg, calibration_params, points, coords, offset, normals, camera_items, logits_dir, allow_missing_ground_truth, ) logging.info("Prepared tile %s directly to %s pt shard batches", tile_info["tile"], len(saved_paths)) return saved_paths parallel_camera_workers = int(cfg.get("runtime", {}).get("parallel_camera_workers", 1)) gpu_count = len(jax.devices("gpu")) if jax.default_backend() == "gpu" else 0 visibility_mode = get_visibility_mode(cfg) worker_output_mode = get_worker_output_mode(cfg) if parallel_camera_workers > 1 and gpu_count > 0 and len(camera_items) > 1 and cfg.get("_worker_mode") is None: worker_count = min(parallel_camera_workers, gpu_count, len(camera_items)) temp_dir = prepared_root / ".camera_workers" / tile_info["tile"] temp_dir.mkdir(parents=True, exist_ok=True) camera_groups = [[] for _ in range(worker_count)] for idx, camera_item in enumerate(camera_items): camera_groups[idx % worker_count].append(camera_item[0]) processes = [] output_files = [] for worker_idx, camera_group in enumerate(camera_groups): if not camera_group: continue camera_list_file = temp_dir / f"worker_{worker_idx}_cameras.json" output_file = temp_dir / f"worker_{worker_idx}_observations.pt" stdout_file = temp_dir / f"worker_{worker_idx}.stdout.log" stderr_file = temp_dir / f"worker_{worker_idx}.stderr.log" camera_list_file.write_text(json.dumps(camera_group), encoding="utf-8") env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = str(worker_idx) env.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") cmd = [ sys.executable, str(Path(__file__).resolve().parents[1] / "scripts" / "run_batch_generation.py"), "--config", cfg["_config_path"], "--worker-mode", "camera-batch", "--tile", tile_info["tile"], "--camera-list-file", str(camera_list_file), "--worker-output", str(output_file), ] logging.info("Launching worker %s for tile %s on GPU %s with %s cameras", worker_idx, tile_info["tile"], worker_idx, len(camera_group)) stdout_handle = None stderr_handle = None if worker_output_mode == "files": stdout_handle = open(stdout_file, "w", encoding="utf-8") stderr_handle = open(stderr_file, "w", encoding="utf-8") process = subprocess.Popen( cmd, cwd=str(Path(__file__).resolve().parents[1]), env=env, stdout=stdout_handle, stderr=stderr_handle, ) elif worker_output_mode == "inherit": process = subprocess.Popen( cmd, cwd=str(Path(__file__).resolve().parents[1]), env=env, ) else: raise ValueError(f"Unsupported worker_output_mode: {worker_output_mode}") processes.append((worker_idx, process, stdout_handle, stderr_handle, stdout_file, stderr_file)) output_files.append(output_file) for worker_idx, process, stdout_handle, stderr_handle, stdout_file, stderr_file in processes: return_code = process.wait() if stdout_handle is not None: stdout_handle.close() if stderr_handle is not None: stderr_handle.close() if return_code != 0: log_hint = ( f" Worker logs: stdout={stdout_file} stderr={stderr_file}" if worker_output_mode == "files" else "" ) if return_code == -9: raise RuntimeError( f"Camera worker {worker_idx} was killed with code -9 for tile {tile_info['tile']}. " f"This usually indicates an out-of-memory kill (GPU or system RAM). " f"{log_hint}" ) raise RuntimeError( f"Camera worker {worker_idx} failed for tile {tile_info['tile']} with code {return_code}. " f"{log_hint}" ) if output_mode == "pt_compact_reduce": compact_payloads = [] for output_file in output_files: worker_payload = torch.load(output_file, map_location="cpu", weights_only=False) compact_payloads.append(worker_payload["observations"]) saved_paths = save_compact_payloads_to_pt( cfg, tile_info["tile"], tile_info["split"], coords, offset, compact_payloads, cfg["data"]["batches_dir"], batch_size=cfg["data"]["dataloader_batch_size"], ) logging.info("Prepared tile %s with compact_reduce to %s pt batches", tile_info["tile"], len(saved_paths)) return saved_paths all_observations = {} max_views = int(cfg["selection"]["max_views"]) selection_strategy = cfg["selection"]["strategy"] for output_file in output_files: worker_payload = torch.load(output_file, map_location="cpu", weights_only=False) merge_worker_observations_bounded( all_observations, worker_payload["observations"], max_views, selection_strategy, ) else: if output_mode == "pt_compact_reduce": compact_payload = process_camera_items_for_tile_compact( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) saved_paths = save_compact_payloads_to_pt( cfg, tile_info["tile"], tile_info["split"], coords, offset, [compact_payload], cfg["data"]["batches_dir"], batch_size=cfg["data"]["dataloader_batch_size"], ) logging.info("Prepared tile %s with compact_reduce to %s pt batches", tile_info["tile"], len(saved_paths)) return saved_paths if visibility_mode == "two_pass_global_depth": all_observations = process_camera_items_for_tile_global_depth( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) elif visibility_mode == "chunk_local_depth": all_observations = process_camera_items_for_tile( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) else: raise ValueError(f"Unsupported visibility_mode: {visibility_mode}") logging.info( "Tile %s finished camera processing; aggregated observations for %s point buckets", tile_info["tile"], count_nonempty_observations(all_observations), ) if output_mode == "pt_direct": saved_paths = save_tile_observations_to_pt( cfg, tile_info["tile"], tile_info["split"], coords, offset, all_observations, cfg["data"]["batches_dir"], batch_size=cfg["data"]["dataloader_batch_size"], ) logging.info("Prepared tile %s directly to %s pt batches", tile_info["tile"], len(saved_paths)) return saved_paths write_tile_h5(output_h5_file, tile_info["tile"], tile_info["split"], coords, all_observations, cfg) logging.info("Prepared tile %s -> %s", tile_info["tile"], output_h5_file) return output_h5_file def run_camera_batch_worker(cfg, tile_name, camera_list_file, worker_output, shard_start=None, shard_end=None): manifest = load_or_build_manifest(cfg) tile_info = next(item for item in manifest if item["tile"] == tile_name) logits_dir = resolve_logits_dir(cfg, tile_info) calibration_params = load_calibration(tile_info["calibration_file"]) points, offset, _, coords = load_las_points(tile_info["lidar_path"]) camera_dict = load_camera_parameters(tile_info["camera_file"], offset) normals_path = Path(cfg["data"]["normals_dir"]) / tile_info["split"] / f"{tile_info['tile']}.h5" normals = load_normals_from_h5(normals_path) allow_missing_ground_truth = tile_info["split"] == "test" if not points["GT_AVAILABLE"] and not allow_missing_ground_truth: raise ValueError( f"Tile {tile_info['tile']} in split {tile_info['split']} has no ground_truth field in LAS." ) selected_camera_names = json.loads(Path(camera_list_file).read_text(encoding="utf-8")) camera_items = [(name, camera_dict[name]) for name in selected_camera_names if name in camera_dict] if shard_start is not None and shard_end is not None: points, coords, normals = slice_point_data(points, coords, normals, int(shard_start), int(shard_end)) output_mode = cfg.get("runtime", {}).get("tile_output_format", "h5_then_pt") if output_mode == "pt_compact_reduce": observations = process_camera_items_for_tile_compact( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) Path(worker_output).parent.mkdir(parents=True, exist_ok=True) torch.save({"observations": observations}, worker_output, pickle_protocol=4) return visibility_mode = get_visibility_mode(cfg) if visibility_mode == "two_pass_global_depth": observations = process_camera_items_for_tile_global_depth( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) elif visibility_mode == "chunk_local_depth": observations = process_camera_items_for_tile( tile_info, camera_items, cfg, calibration_params, points, coords, normals, logits_dir, allow_missing_ground_truth, ) else: raise ValueError(f"Unsupported visibility_mode: {visibility_mode}") Path(worker_output).parent.mkdir(parents=True, exist_ok=True) torch.save({"observations": observations}, worker_output, pickle_protocol=4) def batch_generation_main(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default="configs/config_gridnet_hd_dataset_generation.yaml") parser.add_argument("--worker-mode", choices=["camera-batch"], default=None) parser.add_argument("--tile", default=None) parser.add_argument("--camera-list-file", default=None) parser.add_argument("--worker-output", default=None) parser.add_argument("--shard-start", type=int, default=None) parser.add_argument("--shard-end", type=int, default=None) args = parser.parse_args() cfg = yaml.safe_load(args.config.read_text()) cfg["_config_path"] = str(args.config.resolve()) cfg["_worker_mode"] = args.worker_mode if args.worker_mode == "camera-batch": run_camera_batch_worker( cfg, args.tile, args.camera_list_file, args.worker_output, shard_start=args.shard_start, shard_end=args.shard_end, ) return manifest = load_or_build_manifest(cfg) debug_tiles = cfg.get("debug", {}).get("tiles") if debug_tiles: wanted_tiles = set(debug_tiles) manifest = [tile_info for tile_info in manifest if tile_info["tile"] in wanted_tiles] overwrite = cfg.get("runtime", {}).get("overwrite_existing", False) output_mode = cfg.get("runtime", {}).get("tile_output_format", "h5_then_pt") h5_paths_by_split = {"train": [], "val": [], "test": []} for tile_info in tqdm(manifest, desc="Tiles", unit="tile"): output_path = build_tile_dataset(tile_info, cfg, overwrite=overwrite) if output_mode == "h5_then_pt": h5_paths_by_split[tile_info["split"]].append(str(output_path)) if output_mode == "h5_then_pt": save_dataset_to_pt_parallel( cfg, h5_paths_by_split["train"], h5_paths_by_split["val"], h5_paths_by_split["test"], cfg["data"]["batches_dir"], batch_size=cfg["data"]["dataloader_batch_size"], num_workers=cfg["data"]["max_workers"], )