Compare commits
96 Commits
15 changed files with 16951 additions and 1555 deletions
@ -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 @@ |
|||||||
|
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 @@ |
|||||||
|
#!/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 @@ |
|||||||
|
#!/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 @@ |
|||||||
|
""" |
||||||
|
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