源图均衡

.gitignore

@@ -74,3 +74,6 @@ libs/MVS/texture_output/texture_33.png
 libs/MVS/texture_output/texture_34.png
 libs/MVS/texture_output/texture_35.png
 texture_output/texture_0.png
+color_change/__pycache__/calib_io.cpython-310.pyc
+.gitignore
+.gitignore

color_change/calib_io.py

@@ -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,
+    )

color_change/calibration.py

@@ -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)

color_change/convert_txt_to_bin.py

@@ -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()

color_change/run_calibration.py

@@ -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()