265 lines
8.9 KiB
Python
265 lines
8.9 KiB
Python
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,
|
|
) |