Compare commits
96 Commits
15 changed files with 20005 additions and 8197 deletions
@ -0,0 +1,265 @@
@@ -0,0 +1,265 @@
|
||||
from __future__ import annotations |
||||
|
||||
import json |
||||
import os |
||||
from collections.abc import Callable, Mapping |
||||
from dataclasses import dataclass |
||||
from pathlib import Path, PurePosixPath, PureWindowsPath |
||||
from typing import TypeAlias |
||||
|
||||
import numpy as np |
||||
import pycolmap |
||||
from numpy.typing import NDArray |
||||
from PIL import Image |
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True) |
||||
class CalibrationView: |
||||
image_id: int |
||||
name: str |
||||
rgb_path: Path |
||||
mask_path: Path |
||||
camera: pycolmap.Camera |
||||
cam_from_world: "pycolmap.Rigid3d" |
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True) |
||||
class CalibrationDataset: |
||||
reconstruction: pycolmap.Reconstruction |
||||
views: tuple[CalibrationView, ...] |
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True) |
||||
class CalibrationResult: |
||||
corrected_images: Mapping[int, NDArray[np.uint8]] |
||||
report: Mapping[str, object] |
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True) |
||||
class CalibrationOutputs: |
||||
image_paths: tuple[Path, ...] |
||||
report_path: Path |
||||
|
||||
|
||||
Calibrator: TypeAlias = Callable[[CalibrationDataset], CalibrationResult] |
||||
|
||||
|
||||
def _relative_image_path(name: str) -> Path: |
||||
posix_path = PurePosixPath(name) |
||||
windows_path = PureWindowsPath(name) |
||||
if ( |
||||
not posix_path.parts |
||||
or posix_path.is_absolute() |
||||
or ".." in posix_path.parts |
||||
or windows_path.is_absolute() |
||||
or windows_path.drive |
||||
or ".." in windows_path.parts |
||||
): |
||||
raise ValueError(f"unsafe COLMAP image name: {name}") |
||||
return Path(*posix_path.parts) |
||||
|
||||
|
||||
def _validate_output_target( |
||||
output_root: Path, |
||||
resolved_output_root: Path, |
||||
relative: Path, |
||||
seen_targets: set[str], |
||||
seen_ancestors: set[str], |
||||
) -> Path: |
||||
target = output_root / relative |
||||
resolved_target = target.resolve() |
||||
if not resolved_target.is_relative_to(resolved_output_root): |
||||
raise ValueError(f"output path escapes output_dir: {target}") |
||||
ancestor = target.parent |
||||
while True: |
||||
if (ancestor.exists() or ancestor.is_symlink()) and not ancestor.is_dir(): |
||||
raise NotADirectoryError(f"output ancestor is not a directory: {ancestor}") |
||||
if ancestor == output_root: |
||||
break |
||||
ancestor = ancestor.parent |
||||
target_key = os.path.normcase(str(resolved_target)) |
||||
if target_key in seen_targets: |
||||
raise ValueError(f"duplicate output path: {target}") |
||||
if target_key in seen_ancestors or any( |
||||
os.path.normcase(str(parent)) in seen_targets |
||||
for parent in resolved_target.parents |
||||
): |
||||
raise ValueError(f"conflicting output path: {target}") |
||||
seen_targets.add(target_key) |
||||
seen_ancestors.update( |
||||
os.path.normcase(str(parent)) for parent in resolved_target.parents |
||||
) |
||||
if target.exists() or target.is_symlink(): |
||||
raise FileExistsError(f"output already exists: {target}") |
||||
return target |
||||
|
||||
|
||||
def load_dataset( |
||||
model_dir: str | os.PathLike[str], |
||||
rgb_dir: str | os.PathLike[str], |
||||
mask_dir: str | os.PathLike[str], |
||||
) -> CalibrationDataset: |
||||
# ✅ pycolmap 4.x: 直接用 Reconstruction 加载模型目录 |
||||
reconstruction = pycolmap.Reconstruction(str(Path(model_dir))) |
||||
|
||||
# 验证加载成功 |
||||
if reconstruction.num_images() == 0: |
||||
raise RuntimeError( |
||||
f"Failed to load COLMAP model from {model_dir}. " |
||||
f"Expected cameras.bin/images.bin/points3D.bin or .txt files." |
||||
) |
||||
|
||||
print(f"Loaded COLMAP model: {reconstruction.num_images()} images, " |
||||
f"{reconstruction.num_cameras()} cameras, " |
||||
f"{reconstruction.num_points3D()} 3D points") |
||||
|
||||
rgb_root = Path(rgb_dir) |
||||
mask_root = Path(mask_dir) |
||||
seen_names: set[str] = set() |
||||
views: list[CalibrationView] = [] |
||||
|
||||
# ✅ pycolmap 4.x: images 是一个 dict-like 对象 |
||||
for image_id, image in reconstruction.images.items(): |
||||
if not image.has_pose: |
||||
raise ValueError(f"COLMAP image has no pose: {image.name}") |
||||
|
||||
relative = _relative_image_path(image.name) |
||||
normalized_name = os.path.normcase(str(relative)) |
||||
if normalized_name in seen_names: |
||||
raise ValueError(f"duplicate COLMAP image name: {image.name}") |
||||
seen_names.add(normalized_name) |
||||
|
||||
rgb_path = rgb_root / relative |
||||
mask_path = mask_root / relative.with_suffix(".png") |
||||
if not rgb_path.is_file(): |
||||
raise FileNotFoundError(f"RGB image not found: {rgb_path}") |
||||
if not mask_path.is_file(): |
||||
raise FileNotFoundError(f"mask image not found: {mask_path}") |
||||
|
||||
camera = image.camera |
||||
if camera is None: |
||||
raise ValueError(f"COLMAP image has no camera: {image.name}") |
||||
|
||||
with Image.open(rgb_path) as rgb_image: |
||||
rgb_size = rgb_image.size |
||||
with Image.open(mask_path) as mask_image: |
||||
mask_size = mask_image.size |
||||
camera_size = (camera.width, camera.height) |
||||
if rgb_size != camera_size or mask_size != camera_size: |
||||
raise ValueError( |
||||
f"size mismatch for {image.name}: camera={camera_size}, " |
||||
f"rgb={rgb_size}, mask={mask_size}" |
||||
) |
||||
|
||||
views.append( |
||||
CalibrationView( |
||||
image_id=image_id, |
||||
name=image.name, |
||||
rgb_path=rgb_path, |
||||
mask_path=mask_path, |
||||
camera=camera, |
||||
# ✅ pycolmap 4.x: cam_from_world 是属性，不是方法 |
||||
cam_from_world=image.cam_from_world, |
||||
) |
||||
) |
||||
|
||||
return CalibrationDataset(reconstruction=reconstruction, views=tuple(views)) |
||||
|
||||
|
||||
def read_rgb(view: CalibrationView) -> NDArray[np.uint8]: |
||||
with Image.open(view.rgb_path) as image: |
||||
return np.array(image.convert("RGB"), dtype=np.uint8, copy=True) |
||||
|
||||
|
||||
def read_mask(view: CalibrationView) -> NDArray[np.bool_]: |
||||
with Image.open(view.mask_path) as image: |
||||
return np.array(image.convert("L"), dtype=np.uint8, copy=True) != 0 |
||||
|
||||
|
||||
def run_calibration( |
||||
model_dir: str | os.PathLike[str], |
||||
rgb_dir: str | os.PathLike[str], |
||||
mask_dir: str | os.PathLike[str], |
||||
output_dir: str | os.PathLike[str], |
||||
*, |
||||
calibrator: Calibrator, |
||||
) -> CalibrationOutputs: |
||||
dataset = load_dataset(model_dir, rgb_dir, mask_dir) |
||||
output_root = Path(output_dir) |
||||
resolved_output_root = output_root.resolve() |
||||
if resolved_output_root == Path(rgb_dir).resolve(): |
||||
raise ValueError("output_dir must differ from rgb_dir") |
||||
|
||||
result = calibrator(dataset) |
||||
if not isinstance(result, CalibrationResult): |
||||
raise ValueError("calibrator must return CalibrationResult") |
||||
if not isinstance(result.corrected_images, Mapping): |
||||
raise ValueError("calibrator corrected_images must be a mapping") |
||||
try: |
||||
report_json = json.dumps( |
||||
result.report, |
||||
ensure_ascii=False, |
||||
indent=2, |
||||
sort_keys=True, |
||||
allow_nan=False, |
||||
) |
||||
except (TypeError, ValueError) as error: |
||||
raise ValueError("calibration report is not valid JSON") from error |
||||
|
||||
expected_ids = {view.image_id for view in dataset.views} |
||||
returned_ids = set(result.corrected_images) |
||||
if returned_ids != expected_ids: |
||||
missing = sorted(expected_ids - returned_ids) |
||||
extra = sorted(returned_ids - expected_ids, key=repr) |
||||
raise ValueError( |
||||
f"calibrator returned wrong image ids: missing={missing} extra={extra}" |
||||
) |
||||
|
||||
Image.init() |
||||
output_formats = Image.registered_extensions() |
||||
seen_targets: set[str] = set() |
||||
seen_ancestors: set[str] = set() |
||||
outputs: list[tuple[Path, NDArray[np.uint8]]] = [] |
||||
for view in dataset.views: |
||||
rgb = result.corrected_images[view.image_id] |
||||
expected_shape = (view.camera.height, view.camera.width, 3) |
||||
if ( |
||||
not isinstance(rgb, np.ndarray) |
||||
or rgb.dtype != np.uint8 |
||||
or rgb.shape != expected_shape |
||||
): |
||||
raise ValueError( |
||||
f"invalid output for image {view.name}: expected shape " |
||||
f"{expected_shape} and dtype uint8, got shape " |
||||
f"{getattr(rgb, 'shape', None)} and dtype " |
||||
f"{getattr(rgb, 'dtype', None)}" |
||||
) |
||||
target = _validate_output_target( |
||||
output_root, |
||||
resolved_output_root, |
||||
_relative_image_path(view.name), |
||||
seen_targets, |
||||
seen_ancestors, |
||||
) |
||||
output_format = output_formats.get(target.suffix.lower()) |
||||
if output_format not in Image.SAVE: |
||||
raise ValueError(f"unsupported output image extension: {target.suffix}") |
||||
outputs.append((target, rgb)) |
||||
|
||||
report_path = _validate_output_target( |
||||
output_root, |
||||
resolved_output_root, |
||||
Path("calibration_report.json"), |
||||
seen_targets, |
||||
seen_ancestors, |
||||
) |
||||
for target, rgb in outputs: |
||||
target.parent.mkdir(parents=True, exist_ok=True) |
||||
Image.fromarray(rgb).save(target) |
||||
report_path.parent.mkdir(parents=True, exist_ok=True) |
||||
report_path.write_text(f"{report_json}\n", encoding="utf-8") |
||||
|
||||
return CalibrationOutputs( |
||||
image_paths=tuple(target for target, _ in outputs), |
||||
report_path=report_path, |
||||
) |
||||
@ -0,0 +1,467 @@
@@ -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) |
||||
@ -0,0 +1,154 @@
@@ -0,0 +1,154 @@
|
||||
#!/usr/bin/env python3 |
||||
"""把 COLMAP txt 格式转成 bin 格式（用 pycolmap 4.x 写入）""" |
||||
import sys |
||||
import struct |
||||
import numpy as np |
||||
from pathlib import Path |
||||
|
||||
|
||||
def write_cameras_bin(cameras, output_path): |
||||
"""写 cameras.bin""" |
||||
with open(output_path, 'wb') as f: |
||||
f.write(struct.pack('Q', len(cameras))) # uint64 num_cameras |
||||
for cam_id in sorted(cameras.keys()): |
||||
cam = cameras[cam_id] |
||||
f.write(struct.pack('I', cam_id)) # uint32 camera_id |
||||
# model 映射：SIMPLE_PINHOLE=0, PINHOLE=1, SIMPLE_RADIAL=2, etc. |
||||
model_map = { |
||||
'SIMPLE_PINHOLE': 0, |
||||
'PINHOLE': 1, |
||||
'SIMPLE_RADIAL': 2, |
||||
'RADIAL': 3, |
||||
'OPENCV': 4, |
||||
'OPENCV_FISHEYE': 5, |
||||
'FULL_OPENCV': 6, |
||||
'FOV': 7, |
||||
'THIN_PRISM_FISHEYE': 8, |
||||
} |
||||
model_id = model_map.get(cam['model'], 0) |
||||
f.write(struct.pack('I', model_id)) # uint32 model_id |
||||
f.write(struct.pack('Q', cam['width'])) # uint64 width |
||||
f.write(struct.pack('Q', cam['height'])) # uint64 height |
||||
# params |
||||
for p in cam['params']: |
||||
f.write(struct.pack('d', float(p))) |
||||
|
||||
|
||||
def write_images_bin(images, output_path): |
||||
"""写 images.bin""" |
||||
with open(output_path, 'wb') as f: |
||||
f.write(struct.pack('Q', len(images))) # uint64 num_images |
||||
for img_id in sorted(images.keys()): |
||||
img = images[img_id] |
||||
f.write(struct.pack('I', img_id)) # uint32 image_id |
||||
# qvec (4 doubles) |
||||
for q in img['q']: |
||||
f.write(struct.pack('d', float(q))) |
||||
# tvec (3 doubles) |
||||
for t in img['t']: |
||||
f.write(struct.pack('d', float(t))) |
||||
f.write(struct.pack('I', img['cam_id'])) # uint32 camera_id |
||||
# name (null-terminated string) |
||||
name_bytes = img['name'].encode('utf-8') |
||||
f.write(name_bytes) |
||||
f.write(b'\x00') |
||||
# num_points2D (uint64) |
||||
f.write(struct.pack('Q', 0)) # 简化：0 个 2D 点 |
||||
|
||||
|
||||
def write_points3D_bin(points3D, output_path): |
||||
"""写 points3D.bin""" |
||||
with open(output_path, 'wb') as f: |
||||
f.write(struct.pack('Q', len(points3D))) # uint64 num_points3D |
||||
for pt_id in sorted(points3D.keys()): |
||||
pt = points3D[pt_id] |
||||
f.write(struct.pack('Q', pt_id)) # uint64 point3D_id |
||||
# xyz (3 doubles) |
||||
for v in pt['xyz']: |
||||
f.write(struct.pack('d', float(v))) |
||||
# rgb (3 uint8) |
||||
f.write(struct.pack('B', 0)) |
||||
f.write(struct.pack('B', 0)) |
||||
f.write(struct.pack('B', 0)) |
||||
# error (double) |
||||
f.write(struct.pack('d', 0.0)) |
||||
# track_length (uint64) |
||||
f.write(struct.pack('Q', 0)) # 简化：空 track |
||||
|
||||
|
||||
def main(): |
||||
input_dir = Path(sys.argv[1]) |
||||
output_dir = Path(sys.argv[2]) |
||||
output_dir.mkdir(parents=True, exist_ok=True) |
||||
|
||||
print(f"Reading txt from {input_dir}") |
||||
|
||||
# 解析 cameras.txt |
||||
cameras = {} |
||||
with open(input_dir / "cameras.txt") as f: |
||||
for line in f: |
||||
if line.startswith('#'): |
||||
continue |
||||
parts = line.strip().split() |
||||
if len(parts) < 5: |
||||
continue |
||||
cam_id = int(parts[0]) |
||||
cameras[cam_id] = { |
||||
'model': parts[1], |
||||
'width': int(parts[2]), |
||||
'height': int(parts[3]), |
||||
'params': [float(x) for x in parts[4:]], |
||||
} |
||||
|
||||
# 解析 images.txt |
||||
images = {} |
||||
with open(input_dir / "images.txt") as f: |
||||
lines = f.readlines() |
||||
i = 0 |
||||
while i < len(lines): |
||||
line = lines[i].strip() |
||||
i += 1 |
||||
if not line or line.startswith('#'): |
||||
continue |
||||
parts = line.split() |
||||
if len(parts) < 10: |
||||
continue |
||||
img_id = int(parts[0]) |
||||
q = [float(x) for x in parts[1:5]] |
||||
t = [float(x) for x in parts[5:8]] |
||||
cam_id = int(parts[8]) |
||||
name = parts[9] |
||||
images[img_id] = {'name': name, 'cam_id': cam_id, 'q': q, 't': t} |
||||
# 跳过 2D 点行 |
||||
if i < len(lines): |
||||
i += 1 |
||||
|
||||
# 解析 points3D.txt |
||||
points3D = {} |
||||
with open(input_dir / "points3D.txt") as f: |
||||
for line in f: |
||||
if line.startswith('#'): |
||||
continue |
||||
parts = line.strip().split() |
||||
if len(parts) < 5: |
||||
continue |
||||
pt_id = int(parts[0]) |
||||
xyz = [float(x) for x in parts[1:4]] |
||||
points3D[pt_id] = {'xyz': xyz} |
||||
|
||||
print(f"Found: {len(cameras)} cameras, {len(images)} images, {len(points3D)} points") |
||||
|
||||
# 写 bin |
||||
write_cameras_bin(cameras, output_dir / "cameras.bin") |
||||
write_images_bin(images, output_dir / "images.bin") |
||||
write_points3D_bin(points3D, output_dir / "points3D.bin") |
||||
|
||||
print(f"Written bin files to {output_dir}") |
||||
print(f"Files: cameras.bin, images.bin, points3D.bin") |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
if len(sys.argv) != 3: |
||||
print("Usage: python convert_txt_to_bin.py <input_txt_dir> <output_bin_dir>") |
||||
sys.exit(1) |
||||
main() |
||||
@ -0,0 +1,118 @@
@@ -0,0 +1,118 @@
|
||||
#!/usr/bin/env python3 |
||||
"""光度标定入口脚本 |
||||
|
||||
用法: |
||||
python run_calibration.py \ |
||||
--model_dir /path/to/colmap_model \ |
||||
--rgb_dir /path/to/rgb \ |
||||
--mask_dir /path/to/masks \ |
||||
--output_dir /path/to/output |
||||
""" |
||||
|
||||
import argparse |
||||
import sys |
||||
from pathlib import Path |
||||
|
||||
from calib_io import run_calibration |
||||
from calibration import calibrate, CalibrationOptions |
||||
|
||||
|
||||
def parse_args(): |
||||
parser = argparse.ArgumentParser( |
||||
description="Photometric calibration using COLMAP reconstruction", |
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
||||
) |
||||
parser.add_argument( |
||||
"--model_dir", "-m", |
||||
type=str, required=True, |
||||
help="COLMAP sparse reconstruction directory (contains cameras.bin/images.bin/points3D.bin)", |
||||
) |
||||
parser.add_argument( |
||||
"--rgb_dir", "-r", |
||||
type=str, required=True, |
||||
help="Directory containing input RGB images", |
||||
) |
||||
parser.add_argument( |
||||
"--mask_dir", "-mk", |
||||
type=str, required=True, |
||||
help="Directory containing foreground masks (.png, same names as RGB)", |
||||
) |
||||
parser.add_argument( |
||||
"--output_dir", "-o", |
||||
type=str, required=True, |
||||
help="Output directory for corrected images and report", |
||||
) |
||||
|
||||
# ---- 可选参数 ---- |
||||
parser.add_argument("--mask_erosion_radius", type=int, default=2) |
||||
parser.add_argument("--min_shared_tracks", type=int, default=3) |
||||
parser.add_argument("--valid_color_min", type=int, default=1) |
||||
parser.add_argument("--valid_color_max", type=int, default=254) |
||||
parser.add_argument("--huber_k", type=float, default=1.345) |
||||
parser.add_argument("--max_iterations", type=int, default=20) |
||||
parser.add_argument("--tolerance", type=float, default=1e-6) |
||||
|
||||
return parser.parse_args() |
||||
|
||||
|
||||
def main(): |
||||
args = parse_args() |
||||
|
||||
# 路径校验 |
||||
model_dir = Path(args.model_dir) |
||||
rgb_dir = Path(args.rgb_dir) |
||||
mask_dir = Path(args.mask_dir) |
||||
output_dir = Path(args.output_dir) |
||||
|
||||
if not model_dir.exists(): |
||||
print(f"[ERROR] model_dir does not exist: {model_dir}", file=sys.stderr) |
||||
sys.exit(1) |
||||
if not rgb_dir.exists(): |
||||
print(f"[ERROR] rgb_dir does not exist: {rgb_dir}", file=sys.stderr) |
||||
sys.exit(1) |
||||
if not mask_dir.exists(): |
||||
print(f"[ERROR] mask_dir does not exist: {mask_dir}", file=sys.stderr) |
||||
sys.exit(1) |
||||
|
||||
# 参数 |
||||
options = CalibrationOptions( |
||||
mask_erosion_radius=args.mask_erosion_radius, |
||||
min_shared_tracks=args.min_shared_tracks, |
||||
valid_color_min=args.valid_color_min, |
||||
valid_color_max=args.valid_color_max, |
||||
huber_k=args.huber_k, |
||||
max_iterations=args.max_iterations, |
||||
tolerance=args.tolerance, |
||||
) |
||||
|
||||
print("=" * 50) |
||||
print("Photometric Calibration") |
||||
print("=" * 50) |
||||
print(f" model_dir : {model_dir}") |
||||
print(f" rgb_dir : {rgb_dir}") |
||||
print(f" mask_dir : {mask_dir}") |
||||
print(f" output_dir : {output_dir}") |
||||
print(f" min_shared_tracks: {options.min_shared_tracks}") |
||||
print(f" mask_erosion_radius: {options.mask_erosion_radius}") |
||||
print("-" * 50) |
||||
|
||||
# 执行标定 |
||||
print("Starting calibration...") |
||||
print(f"model_dir={model_dir}") |
||||
outputs = run_calibration( |
||||
model_dir=str(model_dir), |
||||
rgb_dir=str(rgb_dir), |
||||
mask_dir=str(mask_dir), |
||||
output_dir=str(output_dir), |
||||
calibrator=lambda dataset: calibrate(dataset, options=options), |
||||
) |
||||
|
||||
print("=" * 50) |
||||
print(f"[DONE] Corrected images: {len(outputs.image_paths)}") |
||||
print(f" Output dir : {output_dir}") |
||||
print(f" Report : {outputs.report_path}") |
||||
print("=" * 50) |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
main() |
||||
@ -0,0 +1,120 @@
@@ -0,0 +1,120 @@
|
||||
""" |
||||
calibrate.py —— 在跑 OpenMVS 之前，对输入照片做白平衡 + 曝光统一。 |
||||
这是消除"源视图色温差异"最彻底、最安全的办法：在像素进入纹理管线之前就对齐。 |
||||
用法: python calibrate.py <image_dir> <out_dir> |
||||
|
||||
做了两件事（都可选、都保守，不会破坏细节）： |
||||
1) Gray-World 白平衡：把每张图的通道均值拉向全局中值（去色温差） |
||||
2) 全局曝光归一化：把所有图的亮度中位数对齐到全局中位数（去曝光差） |
||||
|
||||
注意：OpenMVS 的 EstimateGlobalPhotometricCorrection 强依赖稠密点云， |
||||
如果没有点云它会被跳过（你之前日志里 "Photometric: no point cloud, skip"）。 |
||||
所以在这里预处理是最稳的兜底。 |
||||
""" |
||||
import os, sys, glob |
||||
import numpy as np |
||||
import cv2 |
||||
|
||||
|
||||
def gray_world(img_bgr, ref_bgr=None): |
||||
"""灰度世界白平衡。 |
||||
|
||||
若给定 ref_bgr（全局参考图的 BGR 均值），则把本图通道均值拉向 ref； |
||||
否则用本图三通道的均值作为中性灰参考（单图自校正）。 |
||||
不做任何“保守衰减”——必须完全对齐，否则等于没做。 |
||||
""" |
||||
b, g, r = [img_bgr[:, :, c].astype(np.float32) for c in range(3)] |
||||
mb, mg, mr = b.mean(), g.mean(), r.mean() |
||||
if ref_bgr is not None: |
||||
rb, rg, rr = ref_bgr |
||||
else: |
||||
# 单图自校正：以三通道共同均值作为中性灰目标 |
||||
m = (mb + mg + mr) / 3.0 |
||||
rb, rg, rr = m, m, m |
||||
kb = rb / mb if mb > 1e-6 else 1.0 |
||||
kg = rg / mg if mg > 1e-6 else 1.0 |
||||
kr = rr / mr if mr > 1e-6 else 1.0 |
||||
# 限制增益上限，防止极端单图把噪声放大 |
||||
def clamp_gain(k): |
||||
return min(max(k, 0.3), 3.0) |
||||
kb, kg, kr = clamp_gain(kb), clamp_gain(kg), clamp_gain(kr) |
||||
out = np.stack([ |
||||
np.clip(b * kb, 0, 255), |
||||
np.clip(g * kg, 0, 255), |
||||
np.clip(r * kr, 0, 255), |
||||
], axis=2).astype(np.uint8) |
||||
return out |
||||
|
||||
|
||||
def normalize_exposure(img_bgr, target_median): |
||||
"""曝光归一化：按亮度中位数线性缩放，对齐到全局目标中位数。 |
||||
只做温和限制（上限 2.0），保证不同视图曝光真正一致。""" |
||||
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) |
||||
med = np.median(gray) |
||||
if med < 1.0: |
||||
return img_bgr |
||||
scale = target_median / med |
||||
scale = min(scale, 2.0) # 防止极暗图放大噪声，但允许充分提亮 |
||||
out = np.clip(img_bgr.astype(np.float32) * scale, 0, 255).astype(np.uint8) |
||||
return out |
||||
|
||||
|
||||
def main(): |
||||
if len(sys.argv) < 3: |
||||
print("usage: python calibrate.py <image_dir> <out_dir>") |
||||
sys.exit(1) |
||||
in_dir, out_dir = sys.argv[1], sys.argv[2] |
||||
os.makedirs(out_dir, exist_ok=True) |
||||
|
||||
paths = sorted(glob.glob(os.path.join(in_dir, "*.*"))) |
||||
paths = [p for p in paths if p.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".bmp"))] |
||||
if not paths: |
||||
print("no images found in", in_dir) |
||||
sys.exit(1) |
||||
|
||||
# ---- 第一遍：统计所有图，确定全局参考（目标白平衡 + 目标亮度）---- |
||||
means = [] # (mb, mg, mr) |
||||
meds = [] |
||||
for p in paths: |
||||
img = cv2.imread(p) |
||||
if img is None: |
||||
continue |
||||
b, g, r = img[:, :, 0].astype(np.float32), img[:, :, 1].astype(np.float32), img[:, :, 2].astype(np.float32) |
||||
means.append((b.mean(), g.mean(), r.mean())) |
||||
gray = 0.299 * r + 0.587 * g + 0.114 * b |
||||
meds.append(np.median(gray)) |
||||
|
||||
means = np.array(means) |
||||
meds = np.array(meds) |
||||
|
||||
# 全局参考 = 所有图通道均值的均值（= 数据集整体真实色彩），亮度取中位数 |
||||
ref_b = means[:, 0].mean() |
||||
ref_g = means[:, 1].mean() |
||||
ref_r = means[:, 2].mean() |
||||
ref_bgr = (ref_b, ref_g, ref_r) |
||||
target_median = float(np.median(meds)) |
||||
|
||||
print(f"[calibrate] {len(paths)} images") |
||||
print(f"[calibrate] global reference white-balance (B,G,R) = ({ref_b:.1f}, {ref_g:.1f}, {ref_r:.1f})") |
||||
print(f"[calibrate] global reference luminance median = {target_median:.1f}") |
||||
|
||||
# ---- 第二遍：所有图统一校正到同一参考 ---- |
||||
for i, p in enumerate(paths): |
||||
img = cv2.imread(p) |
||||
if img is None: |
||||
continue |
||||
|
||||
# 1) 白平衡：把本图拉向全局参考 ref_bgr |
||||
out = gray_world(img, ref_bgr=ref_bgr) |
||||
|
||||
# 2) 曝光归一化：亮度中位数对齐到全局目标 |
||||
out = normalize_exposure(out, target_median) |
||||
|
||||
name = os.path.basename(p) |
||||
cv2.imwrite(os.path.join(out_dir, name), out) |
||||
|
||||
print(f"[calibrate] done -> {out_dir}") |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
main() |
||||
Loading…
Reference in new issue