diff --git a/.gitignore b/.gitignore index 9a5df8f..b0a7a37 100644 --- a/.gitignore +++ b/.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 diff --git a/color_change/calib_io.py b/color_change/calib_io.py new file mode 100644 index 0000000..b55222c --- /dev/null +++ b/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, + ) \ No newline at end of file diff --git a/color_change/calibration.py b/color_change/calibration.py new file mode 100644 index 0000000..6977bc1 --- /dev/null +++ b/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) diff --git a/color_change/convert_txt_to_bin.py b/color_change/convert_txt_to_bin.py new file mode 100644 index 0000000..eae2b36 --- /dev/null +++ b/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 ") + sys.exit(1) + main() \ No newline at end of file diff --git a/color_change/run_calibration.py b/color_change/run_calibration.py new file mode 100644 index 0000000..4350bcd --- /dev/null +++ b/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() \ No newline at end of file