from __future__ import annotations from collections import defaultdict from dataclasses import dataclass from itertools import combinations import numpy as np from numpy.typing import NDArray from calib_io import CalibrationDataset, CalibrationResult, read_mask, read_rgb _PATCH_RADIUS = 2 @dataclass(frozen=True, slots=True) class CalibrationOptions: mask_erosion_radius: int = 2 min_shared_tracks: int = 3 valid_color_min: int = 1 valid_color_max: int = 254 huber_k: float = 1.345 max_iterations: int = 20 tolerance: float = 1e-6 def _validate_options(options: CalibrationOptions) -> None: if options.mask_erosion_radius < 0: raise ValueError("mask_erosion_radius must be non-negative") if options.min_shared_tracks < 1: raise ValueError("min_shared_tracks must be positive") if not 1 <= options.valid_color_min < options.valid_color_max <= 255: raise ValueError("valid color range must satisfy 1 <= min < max <= 255") if not np.isfinite(options.huber_k) or options.huber_k <= 0: raise ValueError("huber_k must be finite and positive") if options.max_iterations < 1: raise ValueError("max_iterations must be positive") if not np.isfinite(options.tolerance) or options.tolerance <= 0: raise ValueError("tolerance must be finite and positive") def _erode_mask(mask: NDArray[np.bool_], radius: int) -> NDArray[np.bool_]: source = np.asarray(mask, dtype=np.bool_) if radius == 0: return source.copy() result = np.zeros_like(source) size = 2 * radius + 1 if source.shape[0] < size or source.shape[1] < size: return result windows = np.lib.stride_tricks.sliding_window_view(source, (size, size)) result[radius:-radius, radius:-radius] = windows.all(axis=(-2, -1)) return result def _sample_tracks( dataset: CalibrationDataset, options: CalibrationOptions, ) -> tuple[dict[int, tuple[tuple[int, NDArray[np.float64]], ...]], int, int]: observations: dict[int, list[tuple[int, NDArray[np.float64]]]] = defaultdict( list ) candidate_observations = 0 valid_observations = 0 duplicate_tracks: set[int] = set() for view in dataset.views: rgb = read_rgb(view) eroded_mask = _erode_mask( read_mask(view), options.mask_erosion_radius, ) # ✅ pycolmap 4.x: 从 images 字典取 image = dataset.reconstruction.images[view.image_id] seen_point3d_ids: set[int] = set() for point2D_idx in range(image.num_points2D()): point = image.point2D(point2D_idx) if not point.has_point3D(): continue candidate_observations += 1 point3d_id = int(point.point3D_id) if point3d_id in seen_point3d_ids: duplicate_tracks.add(point3d_id) else: seen_point3d_ids.add(point3d_id) x, y = np.asarray(point.xy, dtype=np.float64) if not np.isfinite(x) or not np.isfinite(y): continue center_x = int(np.floor(x + 0.5)) center_y = int(np.floor(y + 0.5)) if not ( _PATCH_RADIUS <= center_x < rgb.shape[1] - _PATCH_RADIUS and _PATCH_RADIUS <= center_y < rgb.shape[0] - _PATCH_RADIUS and eroded_mask[center_y, center_x] ): continue median = np.median( rgb[ center_y - _PATCH_RADIUS : center_y + _PATCH_RADIUS + 1, center_x - _PATCH_RADIUS : center_x + _PATCH_RADIUS + 1, ].reshape(-1, 3), axis=0, ) if np.any(median < options.valid_color_min) or np.any( median > options.valid_color_max ): continue observations[point3d_id].append( (view.image_id, median.astype(np.float64, copy=False)) ) valid_observations += 1 del rgb, eroded_mask, seen_point3d_ids sampled: dict[int, tuple[tuple[int, NDArray[np.float64]], ...]] = {} for point3d_id in sorted(observations): if point3d_id in duplicate_tracks: continue track = sorted(observations[point3d_id], key=lambda item: item[0]) image_ids = [image_id for image_id, _ in track] if len(track) < 2 or len(set(image_ids)) != len(image_ids): continue sampled[point3d_id] = tuple(track) return sampled, candidate_observations, valid_observations def _build_constraints( tracks: dict[int, tuple[tuple[int, NDArray[np.float64]], ...]], min_shared_tracks: int, ) -> tuple[ dict[tuple[int, int], int], NDArray[np.int64], NDArray[np.int64], NDArray[np.float64], ]: edge_counts: dict[tuple[int, int], int] = defaultdict(int) for point3d_id in sorted(tracks): for (image_i, _), (image_j, _) in combinations(tracks[point3d_id], 2): edge_counts[(image_i, image_j)] += 1 edges = { edge: edge_counts[edge] for edge in sorted(edge_counts) if edge_counts[edge] >= min_shared_tracks } del edge_counts constraint_count = sum(edges.values()) left = np.empty(constraint_count, dtype=np.int64) right = np.empty(constraint_count, dtype=np.int64) deltas = np.empty((constraint_count, 3), dtype=np.float64) index = 0 for point3d_id in sorted(tracks): for (image_i, color_i), (image_j, color_j) in combinations( tracks[point3d_id], 2 ): if (image_i, image_j) not in edges: continue left[index] = image_i right[index] = image_j deltas[index] = np.log(color_j / 255.0) - np.log(color_i / 255.0) index += 1 return edges, left, right, deltas def _connected_components( image_ids: tuple[int, ...], edges: dict[tuple[int, int], int], ) -> tuple[tuple[int, ...], ...]: neighbors = {image_id: set() for image_id in image_ids} for image_i, image_j in edges: neighbors[image_i].add(image_j) neighbors[image_j].add(image_i) components: list[tuple[int, ...]] = [] remaining = set(image_ids) while remaining: pending = [min(remaining)] component: set[int] = set() while pending: image_id = pending.pop() if image_id in component: continue component.add(image_id) pending.extend(sorted(neighbors[image_id] - component, reverse=True)) remaining -= component components.append(tuple(sorted(component))) return tuple(sorted(components, key=lambda component: component[0])) _ComponentSystem = tuple[ tuple[int, ...], NDArray[np.int64], NDArray[np.int64], NDArray[np.int64], NDArray[np.int64], ] def _prepare_component_systems( image_ids: tuple[int, ...], components: tuple[tuple[int, ...], ...], left_ids: NDArray[np.int64], right_ids: NDArray[np.int64], ) -> tuple[ NDArray[np.int64], NDArray[np.int64], tuple[_ComponentSystem, ...], ]: id_to_index = {image_id: index for index, image_id in enumerate(image_ids)} left_indices = np.fromiter( (id_to_index[int(value)] for value in left_ids), dtype=np.int64, count=len(left_ids), ) right_indices = np.fromiter( (id_to_index[int(value)] for value in right_ids), dtype=np.int64, count=len(right_ids), ) component_by_index = np.empty(len(image_ids), dtype=np.int64) component_indices: list[NDArray[np.int64]] = [] for component_index, component in enumerate(components): indices = np.fromiter( (id_to_index[image_id] for image_id in component), dtype=np.int64, count=len(component), ) component_indices.append(indices) component_by_index[indices] = component_index if len(components) == 1: rows_by_component = (np.arange(len(left_indices), dtype=np.int64),) else: constraint_components = component_by_index[left_indices] order = np.argsort(constraint_components, kind="stable") offsets = np.concatenate( ( np.zeros(1, dtype=np.int64), np.cumsum( np.bincount( constraint_components, minlength=len(components), ) ), ) ) rows_by_component = tuple( order[offsets[index] : offsets[index + 1]] for index in range(len(components)) ) systems: list[_ComponentSystem] = [] global_to_local = np.empty(len(image_ids), dtype=np.int64) for component_index, component in enumerate(components): indices = component_indices[component_index] global_to_local[indices] = np.arange(len(indices), dtype=np.int64) rows = rows_by_component[component_index] systems.append( ( component, indices, rows, global_to_local[left_indices[rows]].copy(), global_to_local[right_indices[rows]].copy(), ) ) return left_indices, right_indices, tuple(systems) def _solve_weighted_channel( image_count: int, systems: tuple[_ComponentSystem, ...], deltas: NDArray[np.float64], weights: NDArray[np.float64], ) -> NDArray[np.float64]: solved = np.zeros(image_count, dtype=np.float64) for component, indices, rows, local_left, local_right in systems: if len(component) == 1: continue row_weights = weights[rows] row_deltas = deltas[rows] size = len(component) laplacian = np.zeros((size, size), dtype=np.float64) rhs = np.zeros(size, dtype=np.float64) np.add.at(laplacian, (local_left, local_left), row_weights) np.add.at(laplacian, (local_right, local_right), row_weights) np.add.at(laplacian, (local_left, local_right), -row_weights) np.add.at(laplacian, (local_right, local_left), -row_weights) np.add.at(rhs, local_left, row_weights * row_deltas) np.add.at(rhs, local_right, -row_weights * row_deltas) component_solution = np.zeros(size, dtype=np.float64) try: component_solution[1:] = np.linalg.solve( laplacian[1:, 1:], rhs[1:] ) except np.linalg.LinAlgError as error: raise ValueError( f"gain solve failed for component {list(component)}" ) from error component_solution -= component_solution.mean() solved[indices] = component_solution return solved def _solve_log_gains( image_ids: tuple[int, ...], components: tuple[tuple[int, ...], ...], left_ids: NDArray[np.int64], right_ids: NDArray[np.int64], delta_log_rgb: NDArray[np.float64], options: CalibrationOptions, ) -> tuple[NDArray[np.float64], list[int], list[bool], list[float]]: log_gains = np.zeros((len(image_ids), 3), dtype=np.float64) if len(left_ids) == 0: return log_gains, [0, 0, 0], [True, True, True], [0.0, 0.0, 0.0] left_indices, right_indices, systems = _prepare_component_systems( image_ids, components, left_ids, right_ids, ) iterations_rgb: list[int] = [] converged_rgb: list[bool] = [] residual_rgb: list[float] = [] for channel in range(3): weights = np.ones(len(left_indices), dtype=np.float64) previous: NDArray[np.float64] | None = None converged = False solution = np.zeros(len(image_ids), dtype=np.float64) iterations = 0 for iterations in range(1, options.max_iterations + 1): solution = _solve_weighted_channel( len(image_ids), systems, delta_log_rgb[:, channel], weights, ) if previous is not None and np.max(np.abs(solution - previous)) <= options.tolerance: converged = True break residual = ( solution[left_indices] - solution[right_indices] - delta_log_rgb[:, channel] ) scale = max(1.4826 * float(np.median(np.abs(residual))), 1e-6) threshold = options.huber_k * scale absolute_residual = np.abs(residual) weights = np.where( absolute_residual <= threshold, 1.0, threshold / np.maximum(absolute_residual, np.finfo(np.float64).tiny), ) previous = solution final_residual = ( solution[left_indices] - solution[right_indices] - delta_log_rgb[:, channel] ) log_gains[:, channel] = solution iterations_rgb.append(iterations) converged_rgb.append(converged) residual_rgb.append(float(np.median(np.abs(final_residual)))) if not np.all(np.isfinite(log_gains)): raise ValueError("gain solve produced non-finite values") return log_gains, iterations_rgb, converged_rgb, residual_rgb def calibrate( dataset: CalibrationDataset, *, options: CalibrationOptions = CalibrationOptions(), ) -> CalibrationResult: _validate_options(options) tracks, candidate_observations, valid_observations = _sample_tracks( dataset, options ) edges, left_ids, right_ids, delta_log_rgb = _build_constraints( tracks, options.min_shared_tracks ) constraint_count = len(left_ids) sampled_track_count = len(tracks) del tracks image_ids = tuple(sorted(view.image_id for view in dataset.views)) components = _connected_components(image_ids, edges) log_gains, iterations, converged, residuals = _solve_log_gains( image_ids, components, left_ids, right_ids, delta_log_rgb, options, ) del left_ids, right_ids, delta_log_rgb gains = np.exp(log_gains) if not np.all(np.isfinite(gains)): raise ValueError("gain solve produced non-finite values") id_to_index = {image_id: index for index, image_id in enumerate(image_ids)} corrected_images: dict[int, NDArray[np.uint8]] = {} for view in dataset.views: corrected = read_rgb(view).astype(np.float32) corrected *= gains[id_to_index[view.image_id]].astype(np.float32) np.rint(corrected, out=corrected) np.clip(corrected, 0, 255, out=corrected) corrected_images[view.image_id] = corrected.astype(np.uint8) del corrected component_by_image = { image_id: component_index for component_index, component in enumerate(components) for image_id in component } report = { "schema_version": 1, "settings": { "patch_size": 2 * _PATCH_RADIUS + 1, "mask_erosion_radius": options.mask_erosion_radius, "min_shared_tracks": options.min_shared_tracks, "valid_color_range": [ options.valid_color_min, options.valid_color_max, ], "huber_k": options.huber_k, "max_iterations": options.max_iterations, "tolerance": options.tolerance, }, "summary": { "images": len(dataset.views), "total_tracks": dataset.reconstruction.num_points3D(), "candidate_observations": candidate_observations, "valid_observations": valid_observations, "sampled_tracks": sampled_track_count, "retained_edges": len(edges), "constraints": constraint_count, "connected_components": len(components), }, "images": [ { "image_id": view.image_id, "name": view.name, "component": component_by_image[view.image_id], "log_gain_rgb": log_gains[ id_to_index[view.image_id] ].tolist(), "gain_rgb": gains[id_to_index[view.image_id]].tolist(), } for view in dataset.views ], "edges": [ { "image_id_a": image_i, "image_id_b": image_j, "shared_tracks": shared_tracks, } for (image_i, image_j), shared_tracks in edges.items() ], "solver": { "iterations_rgb": iterations, "converged_rgb": converged, "median_abs_residual_rgb": residuals, }, } return CalibrationResult(corrected_images=corrected_images, report=report)