468 lines
16 KiB
Python
468 lines
16 KiB
Python
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)
|