openMVS/color_change/calibration.py

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)