| import multiprocessing |
| import os |
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
| import torch |
| from tqdm import tqdm |
|
|
|
|
| def _quantize_coords_batch(coords, coord_scale, coord_offset): |
| coords = np.asarray(coords, dtype=np.float64) |
| coord_offset = np.asarray(coord_offset, dtype=np.float64) |
| return np.rint((coords - coord_offset) / coord_scale).astype(np.int32) |
|
|
|
|
| def _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset=None): |
| if tile_offset is None: |
| tile_offset = np.zeros(3, dtype=np.float64) |
| tile_offset = np.asarray(tile_offset, dtype=np.float64) |
| coord_offset = np.asarray(coord_offset, dtype=np.float64) |
| coords_global = np.asarray(coords, dtype=np.float64) + tile_offset |
| coords_offset_global = coord_offset + tile_offset |
| batch = { |
| "visibility": torch.tensor(np.array(vis), dtype=torch.float32), |
| "logits": torch.tensor(np.array(logits), dtype=torch.float32), |
| "mask": torch.tensor(np.array(masks), dtype=torch.bool), |
| "target": torch.tensor(np.array(targets), dtype=torch.long), |
| "coords_int": torch.tensor(_quantize_coords_batch(coords_global, coord_scale, coords_offset_global), dtype=torch.int32), |
| "coords_scale": torch.tensor(coord_scale, dtype=torch.float64), |
| "coords_offset": torch.tensor(coords_offset_global, dtype=torch.float64), |
| "coords_tile_offset": torch.tensor(tile_offset, dtype=torch.float64), |
| } |
| batch["visibility"] = normalize_visibility(batch["visibility"], vmin, vmax) |
| save_path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(batch, save_path, pickle_protocol=4) |
|
|
|
|
| def decode_coordinates(dataset, handle): |
| coords = np.asarray(dataset) |
| if np.issubdtype(coords.dtype, np.integer) and "coords_scale" in handle.attrs: |
| scale = float(handle.attrs["coords_scale"]) |
| offset = np.asarray(handle.attrs.get("coords_offset", np.zeros(3, dtype=np.float32)), dtype=np.float32) |
| return coords.astype(np.float32) * scale + offset |
| return coords.astype(np.float32) |
|
|
|
|
| def decode_visibility(dataset, handle): |
| visibility = np.asarray(dataset) |
| if np.issubdtype(visibility.dtype, np.integer) and "visibility_quant_max" in handle.attrs: |
| quant_max = float(handle.attrs["visibility_quant_max"]) |
| vmin = np.asarray(handle.attrs["visibility_vmin"], dtype=np.float32) |
| vmax = np.asarray(handle.attrs["visibility_vmax"], dtype=np.float32) |
| normalized = visibility.astype(np.float32) / max(quant_max, 1.0) |
| return normalized * (vmax - vmin) + vmin |
| return visibility.astype(np.float32) |
|
|
|
|
| def decode_logits(dataset, handle): |
| logits = np.asarray(dataset) |
| if np.issubdtype(logits.dtype, np.integer) and "logits_quant_max" in handle.attrs: |
| quant_max = float(handle.attrs["logits_quant_max"]) |
| return logits.astype(np.float32) / max(quant_max, 1.0) |
| return logits.astype(np.float32) |
|
|
|
|
| def decode_dense_visibility(visibility, observations): |
| quant_max = float(observations["visibility_quant_max"]) |
| vmin = np.asarray(observations["visibility_vmin"], dtype=np.float32) |
| vmax = np.asarray(observations["visibility_vmax"], dtype=np.float32) |
| normalized = visibility.astype(np.float32) / max(quant_max, 1.0) |
| return normalized * (vmax - vmin) + vmin |
|
|
|
|
| def decode_dense_logits(logits, observations): |
| quant_max = float(observations["logits_quant_max"]) |
| return logits.astype(np.float32) / max(quant_max, 1.0) |
|
|
|
|
| def normalize_visibility(vis, vmin, vmax, eps=1e-8): |
| mins = torch.tensor(vmin, dtype=torch.float32).view(1, 1, -1) |
| maxs = torch.tensor(vmax, dtype=torch.float32).view(1, 1, -1) |
| vis_clamped = torch.max(torch.min(vis, maxs), mins) |
| return (vis_clamped - mins) / (maxs - mins + eps) |
|
|
|
|
| def _pad_views(array, max_views, pad_value=0, dtype=np.float32): |
| current_views = array.shape[0] |
| if current_views >= max_views: |
| return array[:max_views] |
| pad_shape = (max_views - current_views,) + array.shape[1:] |
| padding = np.full(pad_shape, pad_value, dtype=dtype) |
| return np.concatenate([array, padding], axis=0) |
|
|
|
|
| def process_h5_file(args): |
| path, output_dir, vmin, vmax, batch_size, max_views, coord_scale, train, val, test = args |
| batch_idx = 0 |
| filename_prefix = Path(path).stem |
|
|
| with h5py.File(path, "r") as handle: |
| points = handle["points"] |
| keys = list(points.keys()) |
| for start in range(0, len(keys), batch_size): |
| batch_keys = keys[start : start + batch_size] |
| vis = [] |
| logits = [] |
| masks = [] |
| targets = [] |
| coords = [] |
|
|
| for key in batch_keys: |
| point = points[key] |
| visibility = decode_visibility(point["visibility"], handle) |
| logit_vectors = decode_logits(point["logit_vectors"], handle) |
| num_views = min(len(visibility), max_views) |
|
|
| vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32)) |
| logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32)) |
|
|
| mask = np.zeros(max_views, dtype=bool) |
| mask[:num_views] = True |
| masks.append(mask) |
|
|
| targets.append(point["ground_truth"][()]) |
| coords.append(decode_coordinates(point["coordinates"], handle)) |
|
|
| if path in train: |
| save_path = Path(output_dir) / "train" / f"{filename_prefix}_batch_{batch_idx:05d}.pt" |
| elif path in test: |
| save_path = Path(output_dir) / "test" / f"{filename_prefix}_batch_{batch_idx:05d}.pt" |
| else: |
| save_path = Path(output_dir) / "val" / f"{filename_prefix}_batch_{batch_idx:05d}.pt" |
|
|
| coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0) |
| _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, np.zeros(3, dtype=np.float64)) |
| batch_idx += 1 |
|
|
| return f"{path} done" |
|
|
|
|
| def save_tile_observations_to_pt( |
| cfg, |
| tile_name, |
| split_name, |
| coords_array, |
| tile_offset, |
| all_observations, |
| output_dir, |
| batch_size=1024, |
| ): |
| output_dir = Path(output_dir) / split_name |
| output_dir.mkdir(parents=True, exist_ok=True) |
| max_views = int(cfg["selection"]["max_views"]) |
| min_views = int(cfg["selection"]["min_views"]) |
| vmin = cfg["data"]["vmin"] |
| vmax = cfg["data"]["vmax"] |
| batch_idx = 0 |
| coord_scale = float(cfg.get("storage", {}).get("coord_scale", 0.001)) |
|
|
| vis = [] |
| logits = [] |
| masks = [] |
| targets = [] |
| coords = [] |
| saved_paths = [] |
|
|
| if isinstance(all_observations, dict) and all_observations.get("mode") == "dense": |
| counts = all_observations["counts"] |
| ground_truth = all_observations["ground_truth"] |
| for point_idx in range(len(counts)): |
| num_views = int(counts[point_idx]) |
| if num_views < min_views: |
| continue |
| gt = int(ground_truth[point_idx]) |
| if gt < 0: |
| continue |
|
|
| visibility = decode_dense_visibility(all_observations["visibility"][point_idx, :num_views], all_observations) |
| logit_vectors = decode_dense_logits(all_observations["logit_vectors"][point_idx, :num_views], all_observations) |
|
|
| vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32)) |
| logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32)) |
|
|
| mask = np.zeros(max_views, dtype=bool) |
| mask[:num_views] = True |
| masks.append(mask) |
| targets.append(gt) |
| coords.append(np.asarray(coords_array[point_idx], dtype=np.float32)) |
|
|
| if len(vis) == batch_size: |
| save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt" |
| coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0) |
| _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset) |
| saved_paths.append(str(save_path)) |
| vis, logits, masks, targets, coords = [], [], [], [], [] |
| batch_idx += 1 |
| else: |
| for point_idx, observations in all_observations.items(): |
| if len(observations["camera"]) < min_views: |
| continue |
| gt = observations["ground_truth"] |
| if gt is None: |
| continue |
|
|
| visibility = np.asarray(observations["visibility"], dtype=np.float32) |
| logit_vectors = np.asarray(observations["logit_vectors"], dtype=np.float32) |
| num_views = min(len(visibility), max_views) |
|
|
| vis.append(_pad_views(visibility, max_views, pad_value=0.0, dtype=np.float32)) |
| logits.append(_pad_views(logit_vectors, max_views, pad_value=0.0, dtype=np.float32)) |
|
|
| mask = np.zeros(max_views, dtype=bool) |
| mask[:num_views] = True |
| masks.append(mask) |
| targets.append(gt) |
| coords.append(np.asarray(coords_array[point_idx], dtype=np.float32)) |
|
|
| if len(vis) == batch_size: |
| save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt" |
| coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0) |
| _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset) |
| saved_paths.append(str(save_path)) |
| vis, logits, masks, targets, coords = [], [], [], [], [] |
| batch_idx += 1 |
|
|
| if vis: |
| save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt" |
| coord_offset = np.min(np.asarray(coords, dtype=np.float32), axis=0) |
| _save_batch(save_path, vis, logits, masks, targets, coords, vmin, vmax, coord_scale, coord_offset, tile_offset) |
| saved_paths.append(str(save_path)) |
|
|
| return saved_paths |
|
|
|
|
| def save_compact_payloads_to_pt( |
| cfg, |
| tile_name, |
| split_name, |
| coords_array, |
| tile_offset, |
| payloads, |
| output_dir, |
| batch_size=1024, |
| ): |
| output_dir = Path(output_dir) / split_name |
| output_dir.mkdir(parents=True, exist_ok=True) |
| max_views = int(cfg["selection"]["max_views"]) |
| min_views = int(cfg["selection"]["min_views"]) |
| vmin = cfg["data"]["vmin"] |
| vmax = cfg["data"]["vmax"] |
|
|
| payloads = [payload for payload in payloads if payload is not None and payload.get("point_indices") is not None] |
| if not payloads: |
| return [] |
|
|
| point_indices = np.concatenate([payload["point_indices"] for payload in payloads], axis=0) |
| if point_indices.size == 0: |
| return [] |
| visibility = np.concatenate([payload["visibility"] for payload in payloads], axis=0) |
| logits = np.concatenate([payload["logit_vectors"] for payload in payloads], axis=0) |
| ground_truth = np.concatenate([payload["ground_truth"] for payload in payloads], axis=0) |
|
|
| order = np.argsort(point_indices, kind="mergesort") |
| point_indices = point_indices[order] |
| visibility = visibility[order] |
| logits = logits[order] |
| ground_truth = ground_truth[order] |
|
|
| quant_max_vis = float(payloads[0]["visibility_quant_max"]) |
| quant_max_logits = float(payloads[0]["logits_quant_max"]) |
| payload_vmin = np.asarray(payloads[0]["visibility_vmin"], dtype=np.float32) |
| payload_vmax = np.asarray(payloads[0]["visibility_vmax"], dtype=np.float32) |
|
|
| saved_paths = [] |
| vis_batch = [] |
| logits_batch = [] |
| masks_batch = [] |
| targets_batch = [] |
| coords_batch = [] |
| batch_idx = 0 |
| coord_scale = float(cfg.get("storage", {}).get("coord_scale", 0.001)) |
|
|
| 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]) |
| num_obs = end - start |
| if num_obs < min_views: |
| continue |
|
|
| gt_slice = ground_truth[start:end] |
| valid_gt = gt_slice[gt_slice >= 0] |
| if valid_gt.size == 0: |
| continue |
| gt = int(valid_gt[0]) |
|
|
| local_visibility = visibility[start:end] |
| if num_obs > max_views: |
| keep = np.lexsort((local_visibility[:, 0], local_visibility[:, 1]))[:max_views] |
| local_visibility = local_visibility[keep] |
| local_logits = logits[start:end][keep] |
| num_views = max_views |
| else: |
| local_logits = logits[start:end] |
| num_views = num_obs |
|
|
| vis_float = (local_visibility.astype(np.float32) / max(quant_max_vis, 1.0)) * (payload_vmax - payload_vmin) + payload_vmin |
| logits_float = local_logits.astype(np.float32) / max(quant_max_logits, 1.0) |
|
|
| vis_batch.append(_pad_views(vis_float, max_views, pad_value=0.0, dtype=np.float32)) |
| logits_batch.append(_pad_views(logits_float, max_views, pad_value=0.0, dtype=np.float32)) |
| mask = np.zeros(max_views, dtype=bool) |
| mask[:num_views] = True |
| masks_batch.append(mask) |
| targets_batch.append(gt) |
| coords_batch.append(np.asarray(coords_array[point_idx], dtype=np.float32)) |
|
|
| if len(vis_batch) == batch_size: |
| save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt" |
| coord_offset = np.min(np.asarray(coords_batch, dtype=np.float32), axis=0) |
| _save_batch( |
| save_path, |
| vis_batch, |
| logits_batch, |
| masks_batch, |
| targets_batch, |
| coords_batch, |
| vmin, |
| vmax, |
| coord_scale, |
| coord_offset, |
| tile_offset, |
| ) |
| saved_paths.append(str(save_path)) |
| vis_batch, logits_batch, masks_batch, targets_batch, coords_batch = [], [], [], [], [] |
| batch_idx += 1 |
|
|
| if vis_batch: |
| save_path = output_dir / f"{tile_name}_batch_{batch_idx:05d}.pt" |
| coord_offset = np.min(np.asarray(coords_batch, dtype=np.float32), axis=0) |
| _save_batch( |
| save_path, |
| vis_batch, |
| logits_batch, |
| masks_batch, |
| targets_batch, |
| coords_batch, |
| vmin, |
| vmax, |
| coord_scale, |
| coord_offset, |
| tile_offset, |
| ) |
| saved_paths.append(str(save_path)) |
|
|
| return saved_paths |
|
|
|
|
| def save_dataset_to_pt_parallel(cfg, train, val, test, output_dir, batch_size=1024, num_workers=4): |
| output_dir = Path(output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| max_views = cfg["selection"]["max_views"] |
| vmin = cfg["data"]["vmin"] |
| vmax = cfg["data"]["vmax"] |
| h5_paths = train + val + test |
| args_list = [ |
| (path, str(output_dir), vmin, vmax, batch_size, max_views, float(cfg.get("storage", {}).get("coord_scale", 0.001)), train, val, test) |
| for path in h5_paths |
| ] |
|
|
| with multiprocessing.Pool(num_workers) as pool: |
| for result in pool.imap_unordered(process_h5_file, args_list): |
| print(result) |
|
|
