源图均衡
This commit is contained in:
hesuicong committed
1 parent
56f3c93216
commit
adedcfd4f1
5 files changed
+1007
No files matched your search
@@ -0,0 +1,467 @@
|
||||
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)
|
||||
Reference in new issue
Block a user