Update scripts/infer_unisharp.py
Browse files- scripts/infer_unisharp.py +288 -103
scripts/infer_unisharp.py
CHANGED
|
@@ -33,6 +33,30 @@ IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".PNG", ".JPG", ".JPEG", ".W
|
|
| 33 |
CameraKind = Literal["perspective", "fisheye", "panorama"]
|
| 34 |
FACE_NAMES = ["up", "back", "left", "front", "right", "down"]
|
| 35 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 36 |
|
| 37 |
def _configure_torchhub_cache() -> None:
|
| 38 |
torchhub_dir = REPO_ROOT / "checkpoints" / "torchhub"
|
|
@@ -97,6 +121,52 @@ def _collect_image_paths(args: argparse.Namespace) -> list[Path]:
|
|
| 97 |
return paths[: int(args.max_images)] if int(args.max_images) > 0 else paths
|
| 98 |
|
| 99 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 100 |
def _load_rgb_u8(image_path: Path, max_long_edge: int) -> torch.Tensor:
|
| 101 |
with Image.open(image_path) as raw:
|
| 102 |
image = ImageOps.exif_transpose(raw).convert("RGB")
|
|
@@ -206,15 +276,15 @@ def _classify_camera(stats: dict[str, float], args: argparse.Namespace) -> Camer
|
|
| 206 |
v_fov = float(stats["vertical_fov_deg"])
|
| 207 |
diag_fov = float(stats["diagonal_fov_deg"])
|
| 208 |
if (
|
| 209 |
-
|
| 210 |
-
and h_fov >=
|
| 211 |
-
and v_fov >=
|
| 212 |
):
|
| 213 |
return "panorama"
|
| 214 |
-
fishlike_aspect = aspect <=
|
| 215 |
fishlike_fov = (
|
| 216 |
-
max(h_fov, v_fov) >=
|
| 217 |
-
or (diag_fov >=
|
| 218 |
)
|
| 219 |
if fishlike_aspect and fishlike_fov:
|
| 220 |
return "fisheye"
|
|
@@ -321,6 +391,13 @@ def _camera_name_from_json(entry: dict[str, Any] | None) -> str | None:
|
|
| 321 |
return str(value).strip().lower() if value is not None and str(value).strip() else None
|
| 322 |
|
| 323 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 324 |
@torch.no_grad()
|
| 325 |
def _predict_unik3d_rays(
|
| 326 |
model: UnisharpFeatureModel,
|
|
@@ -371,6 +448,72 @@ def _build_rotate_poses(num_views: int, radius_m: float, device: torch.device) -
|
|
| 371 |
return poses
|
| 372 |
|
| 373 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 374 |
def _render_pinhole_frame(
|
| 375 |
renderer: GSplatRenderer,
|
| 376 |
gaussians: Any,
|
|
@@ -515,85 +658,114 @@ def _process_one(
|
|
| 515 |
step: int,
|
| 516 |
args: argparse.Namespace,
|
| 517 |
) -> None:
|
| 518 |
-
|
| 519 |
-
_, h, w = rgb_u8.shape
|
| 520 |
-
if h < 4 or w < 4:
|
| 521 |
-
raise ValueError(f"Invalid image size for {image_path}: {tuple(rgb_u8.shape)}")
|
| 522 |
-
|
| 523 |
-
device = next(model.parameters()).device
|
| 524 |
-
image_u8 = rgb_u8.unsqueeze(0).to(device=device)
|
| 525 |
-
image = image_u8.to(torch.float32) / 255.0
|
| 526 |
-
|
| 527 |
camera_json_entry = _camera_json_for_image(getattr(args, "_camera_json_data", None), image_path)
|
| 528 |
-
|
| 529 |
-
|
| 530 |
-
|
| 531 |
-
|
| 532 |
-
|
| 533 |
-
|
| 534 |
-
|
| 535 |
-
|
| 536 |
-
|
| 537 |
-
|
| 538 |
-
|
| 539 |
-
|
| 540 |
-
|
| 541 |
-
|
| 542 |
-
|
| 543 |
-
|
| 544 |
-
|
| 545 |
-
|
| 546 |
-
|
| 547 |
-
|
| 548 |
-
|
| 549 |
-
|
| 550 |
-
|
| 551 |
-
|
| 552 |
-
|
| 553 |
-
|
| 554 |
-
|
| 555 |
-
|
| 556 |
-
render_camera_params =
|
| 557 |
-
|
| 558 |
-
|
| 559 |
-
|
| 560 |
-
|
| 561 |
-
|
| 562 |
-
|
| 563 |
-
|
| 564 |
-
|
| 565 |
-
|
| 566 |
-
|
| 567 |
-
|
| 568 |
-
|
| 569 |
-
|
| 570 |
-
|
| 571 |
-
|
|
|
|
| 572 |
camera_kind = "fisheye"
|
| 573 |
-
|
| 574 |
-
camera_kind = "perspective"
|
| 575 |
-
else:
|
| 576 |
-
camera_kind = _classify_camera(stats, args)
|
| 577 |
-
if camera_kind == "panorama":
|
| 578 |
-
out = _run_model_panorama(model, image, image_u8, distance_init_cap_m=float(args.distance_init_cap_m))
|
| 579 |
-
elif camera_kind == "fisheye":
|
| 580 |
-
render_camera_params = fit_fisheye624_params_from_rays(rays).detach().to(device=device, dtype=torch.float32)
|
| 581 |
out = _run_model_fisheye(
|
| 582 |
model,
|
| 583 |
image,
|
| 584 |
image_u8,
|
| 585 |
-
camera_params=
|
| 586 |
-
distance_init_cap_m=
|
| 587 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 588 |
else:
|
| 589 |
-
|
| 590 |
-
|
| 591 |
-
|
| 592 |
-
|
| 593 |
-
|
| 594 |
-
|
| 595 |
-
|
| 596 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 597 |
|
| 598 |
LOGGER.info(
|
| 599 |
"%s -> %s | hfov=%.1f vfov=%.1f diag=%.1f aspect=%.3f",
|
|
@@ -607,27 +779,49 @@ def _process_one(
|
|
| 607 |
|
| 608 |
src_w2c = torch.eye(4, dtype=torch.float32, device=device)
|
| 609 |
gaussians_world = transform_gaussians_to_world(out["gaussians"], src_w2c)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 610 |
forward_poses = _build_forward_poses(
|
| 611 |
-
num_views=
|
| 612 |
-
distance_m=
|
| 613 |
device=device,
|
| 614 |
)
|
| 615 |
rotate_poses = _build_rotate_poses(
|
| 616 |
-
num_views=
|
| 617 |
-
radius_m=
|
| 618 |
device=device,
|
| 619 |
)
|
| 620 |
|
| 621 |
sample_dir = out_root / _slug_from_path(image_path)
|
| 622 |
sample_dir.mkdir(parents=True, exist_ok=True)
|
| 623 |
output_crop_border_fraction = 0.0 if camera_kind == "panorama" else 0.05
|
| 624 |
-
Image.fromarray(_crop_border_u8(_to_u8_hwc(rgb_u8), output_crop_border_fraction)).save(sample_dir / "input.png")
|
| 625 |
|
| 626 |
forward_frames: list[np.ndarray] = []
|
| 627 |
rotate_frames: list[np.ndarray] = []
|
| 628 |
|
| 629 |
if camera_kind == "panorama":
|
| 630 |
-
face_w =
|
| 631 |
forward_dir = sample_dir / "forward_erp"
|
| 632 |
rotate_dir = sample_dir / "rotate_erp"
|
| 633 |
rotate_faces_dir = sample_dir / "rotate_cubemap_faces"
|
|
@@ -691,8 +885,8 @@ def _process_one(
|
|
| 691 |
forward_frames = [_crop_border_u8(frame, output_crop_border_fraction) for frame in forward_frames]
|
| 692 |
rotate_frames = [_crop_border_u8(frame, output_crop_border_fraction) for frame in rotate_frames]
|
| 693 |
|
| 694 |
-
_save_gif(forward_frames, sample_dir / "
|
| 695 |
-
_save_gif(rotate_frames, sample_dir / "
|
| 696 |
_save_ply_if_requested(gaussians_world, sample_dir / "gaussians.ply", f_px=f_px, image_h=h, image_w=w, enabled=bool(args.save_ply))
|
| 697 |
|
| 698 |
metadata = {
|
|
@@ -703,10 +897,17 @@ def _process_one(
|
|
| 703 |
"ray_stats": stats,
|
| 704 |
"camera_json": str(args.camera_json) if args.camera_json is not None else None,
|
| 705 |
"camera_json_entry": camera_json_entry,
|
|
|
|
| 706 |
"explicit_camera_intrinsics": args.camera_intrinsics,
|
| 707 |
"explicit_camera_params": args.camera_params,
|
| 708 |
-
"forward_distance_m": float(
|
| 709 |
-
"rotate_radius_m": float(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 710 |
"rotate_path": "clockwise_camera_xy_orbit_fixed_source_orientation",
|
| 711 |
"panorama_renderer": "unisharp.cli.unified_trainer.UnifiedTrainer._render_cubemap/_cube_to_erp",
|
| 712 |
"low_pass_filter_eps": float(args.low_pass_filter_eps),
|
|
@@ -727,14 +928,6 @@ def _build_argparser() -> argparse.ArgumentParser:
|
|
| 727 |
p.add_argument("--out-dir", type=Path, default=REPO_ROOT / "outputs" / "inference")
|
| 728 |
p.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
| 729 |
p.add_argument("--max-images", type=int, default=0)
|
| 730 |
-
p.add_argument("--max-long-edge", type=int, default=768)
|
| 731 |
-
p.add_argument("--forward-views", type=int, default=10)
|
| 732 |
-
p.add_argument("--forward-distance-m", type=float, default=0.2)
|
| 733 |
-
p.add_argument("--rotate-views", type=int, default=10)
|
| 734 |
-
p.add_argument("--rotate-radius-m", type=float, default=0.1)
|
| 735 |
-
p.add_argument("--gif-duration-ms", type=int, default=300)
|
| 736 |
-
p.add_argument("--face-w", type=int, default=0, help="Panorama cubemap face width. 0 uses min(H, W/4).")
|
| 737 |
-
p.add_argument("--distance-init-cap-m", type=float, default=0.0)
|
| 738 |
p.add_argument("--save-ply", action="store_true")
|
| 739 |
p.add_argument(
|
| 740 |
"--camera-json",
|
|
@@ -763,14 +956,6 @@ def _build_argparser() -> argparse.ArgumentParser:
|
|
| 763 |
choices=["auto", "perspective", "pinhole", "fisheye", "panorama", "erp"],
|
| 764 |
help="Override automatic ray-range camera classification.",
|
| 765 |
)
|
| 766 |
-
p.add_argument("--fisheye-fov-threshold-deg", type=float, default=95.0)
|
| 767 |
-
p.add_argument("--fisheye-diag-threshold-deg", type=float, default=130.0)
|
| 768 |
-
p.add_argument("--fisheye-vfov-min-deg", type=float, default=70.0)
|
| 769 |
-
p.add_argument("--fisheye-max-aspect", type=float, default=1.65)
|
| 770 |
-
p.add_argument("--panorama-hfov-threshold-deg", type=float, default=260.0)
|
| 771 |
-
p.add_argument("--panorama-vfov-threshold-deg", type=float, default=120.0)
|
| 772 |
-
p.add_argument("--panorama-aspect-min", type=float, default=1.75)
|
| 773 |
-
p.add_argument("--panorama-aspect-max", type=float, default=2.25)
|
| 774 |
p.add_argument("--low-pass-filter-eps", type=float, default=0.0)
|
| 775 |
return p
|
| 776 |
|
|
|
|
| 33 |
CameraKind = Literal["perspective", "fisheye", "panorama"]
|
| 34 |
FACE_NAMES = ["up", "back", "left", "front", "right", "down"]
|
| 35 |
|
| 36 |
+
MAX_LONG_EDGE = 0
|
| 37 |
+
PERSPECTIVE_MAX_LONG_EDGE = 768
|
| 38 |
+
PANORAMA_MAX_LONG_EDGE = 1536
|
| 39 |
+
FORWARD_VIEWS = 10
|
| 40 |
+
FORWARD_DISTANCE_M = 0.2
|
| 41 |
+
ROTATE_VIEWS = 10
|
| 42 |
+
ROTATE_RADIUS_M = 0.1
|
| 43 |
+
GIF_DURATION_MS = 300
|
| 44 |
+
VIEW_MOTION_NEAR_SCENE_DEPTH_M = 2.0
|
| 45 |
+
VIEW_MOTION_MIN_SCALE = 0.08
|
| 46 |
+
VIEW_MOTION_FORWARD_DEPTH_FRAC = 0.04
|
| 47 |
+
VIEW_MOTION_ROTATE_DEPTH_FRAC = 0.02
|
| 48 |
+
VIEW_MOTION_FAR_SCENE_MEDIAN_M = 2.5
|
| 49 |
+
VIEW_MOTION_FOREGROUND_DEPTH_QUANTILE = 0.20
|
| 50 |
+
|
| 51 |
+
FISHEYE_FOV_THRESHOLD_DEG = 120.0
|
| 52 |
+
FISHEYE_DIAG_THRESHOLD_DEG = 150.0
|
| 53 |
+
FISHEYE_VFOV_MIN_DEG = 80.0
|
| 54 |
+
FISHEYE_MAX_ASPECT = 1.65
|
| 55 |
+
PANORAMA_HFOV_THRESHOLD_DEG = 300.0
|
| 56 |
+
PANORAMA_VFOV_THRESHOLD_DEG = 120.0
|
| 57 |
+
PANORAMA_ASPECT_MIN = 1.9
|
| 58 |
+
PANORAMA_ASPECT_MAX = 2.1
|
| 59 |
+
|
| 60 |
|
| 61 |
def _configure_torchhub_cache() -> None:
|
| 62 |
torchhub_dir = REPO_ROOT / "checkpoints" / "torchhub"
|
|
|
|
| 121 |
return paths[: int(args.max_images)] if int(args.max_images) > 0 else paths
|
| 122 |
|
| 123 |
|
| 124 |
+
def _perspective_max_long_edge() -> int:
|
| 125 |
+
return int(PERSPECTIVE_MAX_LONG_EDGE)
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def _panorama_max_long_edge() -> int:
|
| 129 |
+
return int(PANORAMA_MAX_LONG_EDGE)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def _image_hw_from_path(image_path: Path) -> tuple[int, int]:
|
| 133 |
+
with Image.open(image_path) as raw:
|
| 134 |
+
image = ImageOps.exif_transpose(raw)
|
| 135 |
+
w, h = image.size
|
| 136 |
+
return int(h), int(w)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def _should_load_panorama_native(
|
| 140 |
+
*,
|
| 141 |
+
image_path: Path,
|
| 142 |
+
args: argparse.Namespace,
|
| 143 |
+
camera_json_entry: dict[str, Any] | None,
|
| 144 |
+
) -> bool:
|
| 145 |
+
forced = str(args.camera).strip().lower()
|
| 146 |
+
if forced in {"panorama", "erp"}:
|
| 147 |
+
return True
|
| 148 |
+
if forced in {"perspective", "pinhole", "fisheye"}:
|
| 149 |
+
return False
|
| 150 |
+
json_camera_name = _camera_name_from_json(camera_json_entry)
|
| 151 |
+
if json_camera_name in {"panorama", "erp", "spherical"}:
|
| 152 |
+
return True
|
| 153 |
+
if json_camera_name in {"perspective", "pinhole", "fisheye", "fisheye624", "opencv_fisheye"}:
|
| 154 |
+
return False
|
| 155 |
+
image_h, image_w = _image_hw_from_path(image_path)
|
| 156 |
+
return _camera_name_from_aspect(image_h=image_h, image_w=image_w) == "panorama"
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def _initial_max_long_edge(
|
| 160 |
+
*,
|
| 161 |
+
image_path: Path,
|
| 162 |
+
args: argparse.Namespace,
|
| 163 |
+
camera_json_entry: dict[str, Any] | None,
|
| 164 |
+
) -> int:
|
| 165 |
+
if _should_load_panorama_native(image_path=image_path, args=args, camera_json_entry=camera_json_entry):
|
| 166 |
+
return _panorama_max_long_edge()
|
| 167 |
+
return _perspective_max_long_edge()
|
| 168 |
+
|
| 169 |
+
|
| 170 |
def _load_rgb_u8(image_path: Path, max_long_edge: int) -> torch.Tensor:
|
| 171 |
with Image.open(image_path) as raw:
|
| 172 |
image = ImageOps.exif_transpose(raw).convert("RGB")
|
|
|
|
| 276 |
v_fov = float(stats["vertical_fov_deg"])
|
| 277 |
diag_fov = float(stats["diagonal_fov_deg"])
|
| 278 |
if (
|
| 279 |
+
PANORAMA_ASPECT_MIN <= aspect <= PANORAMA_ASPECT_MAX
|
| 280 |
+
and h_fov >= PANORAMA_HFOV_THRESHOLD_DEG
|
| 281 |
+
and v_fov >= PANORAMA_VFOV_THRESHOLD_DEG
|
| 282 |
):
|
| 283 |
return "panorama"
|
| 284 |
+
fishlike_aspect = aspect <= FISHEYE_MAX_ASPECT
|
| 285 |
fishlike_fov = (
|
| 286 |
+
max(h_fov, v_fov) >= FISHEYE_FOV_THRESHOLD_DEG
|
| 287 |
+
or (diag_fov >= FISHEYE_DIAG_THRESHOLD_DEG and v_fov >= FISHEYE_VFOV_MIN_DEG)
|
| 288 |
)
|
| 289 |
if fishlike_aspect and fishlike_fov:
|
| 290 |
return "fisheye"
|
|
|
|
| 391 |
return str(value).strip().lower() if value is not None and str(value).strip() else None
|
| 392 |
|
| 393 |
|
| 394 |
+
def _camera_name_from_aspect(image_h: int, image_w: int) -> str | None:
|
| 395 |
+
aspect = float(image_w) / float(max(image_h, 1))
|
| 396 |
+
if PANORAMA_ASPECT_MIN <= aspect <= PANORAMA_ASPECT_MAX:
|
| 397 |
+
return "panorama"
|
| 398 |
+
return None
|
| 399 |
+
|
| 400 |
+
|
| 401 |
@torch.no_grad()
|
| 402 |
def _predict_unik3d_rays(
|
| 403 |
model: UnisharpFeatureModel,
|
|
|
|
| 448 |
return poses
|
| 449 |
|
| 450 |
|
| 451 |
+
def _predicted_depth_samples_m(model_output: dict[str, Any]) -> torch.Tensor | None:
|
| 452 |
+
depth = model_output.get("unik3d_distance")
|
| 453 |
+
if not torch.is_tensor(depth):
|
| 454 |
+
layers = model_output.get("distance_layers")
|
| 455 |
+
if torch.is_tensor(layers) and layers.ndim >= 4 and int(layers.shape[1]) >= 1:
|
| 456 |
+
depth = layers[:, 0:1]
|
| 457 |
+
if torch.is_tensor(depth) and depth.numel() > 0:
|
| 458 |
+
values = depth.detach().reshape(-1).to(torch.float32)
|
| 459 |
+
valid = values[torch.isfinite(values) & (values > 1e-3) & (values < 1e4)]
|
| 460 |
+
if int(valid.numel()) > 0:
|
| 461 |
+
return valid
|
| 462 |
+
gaussians = model_output.get("gaussians")
|
| 463 |
+
if gaussians is not None and hasattr(gaussians, "mean_vectors"):
|
| 464 |
+
z = gaussians.mean_vectors.detach().reshape(-1, 3)[..., 2].reshape(-1).to(torch.float32)
|
| 465 |
+
valid = z[torch.isfinite(z) & (z > 1e-3) & (z < 1e4)]
|
| 466 |
+
if int(valid.numel()) > 0:
|
| 467 |
+
return valid
|
| 468 |
+
return None
|
| 469 |
+
|
| 470 |
+
|
| 471 |
+
def _scene_depth_for_motion_m(model_output: dict[str, Any]) -> tuple[float | None, float | None, float | None]:
|
| 472 |
+
"""Return (effective_depth, median_depth, foreground_depth_p25) in meters."""
|
| 473 |
+
valid = _predicted_depth_samples_m(model_output)
|
| 474 |
+
if valid is None or int(valid.numel()) == 0:
|
| 475 |
+
return None, None, None
|
| 476 |
+
median_depth_m = float(torch.median(valid).item())
|
| 477 |
+
q = float(VIEW_MOTION_FOREGROUND_DEPTH_QUANTILE)
|
| 478 |
+
foreground_depth_m = float(torch.quantile(valid, q).item())
|
| 479 |
+
if float(median_depth_m) >= float(VIEW_MOTION_FAR_SCENE_MEDIAN_M):
|
| 480 |
+
effective_depth_m = float(median_depth_m)
|
| 481 |
+
else:
|
| 482 |
+
effective_depth_m = float(min(median_depth_m, foreground_depth_m))
|
| 483 |
+
return effective_depth_m, median_depth_m, foreground_depth_m
|
| 484 |
+
|
| 485 |
+
|
| 486 |
+
def _adaptive_view_motion_distances(
|
| 487 |
+
model_output: dict[str, Any],
|
| 488 |
+
*,
|
| 489 |
+
default_forward_m: float,
|
| 490 |
+
default_radius_m: float,
|
| 491 |
+
) -> tuple[float, float, float | None, float, float | None, float | None]:
|
| 492 |
+
effective_depth_m, median_depth_m, foreground_depth_m = _scene_depth_for_motion_m(model_output)
|
| 493 |
+
near_threshold_m = float(VIEW_MOTION_NEAR_SCENE_DEPTH_M)
|
| 494 |
+
if (
|
| 495 |
+
effective_depth_m is None
|
| 496 |
+
or not math.isfinite(effective_depth_m)
|
| 497 |
+
or float(effective_depth_m) >= near_threshold_m
|
| 498 |
+
):
|
| 499 |
+
return (
|
| 500 |
+
float(default_forward_m),
|
| 501 |
+
float(default_radius_m),
|
| 502 |
+
effective_depth_m,
|
| 503 |
+
1.0,
|
| 504 |
+
median_depth_m,
|
| 505 |
+
foreground_depth_m,
|
| 506 |
+
)
|
| 507 |
+
scale = max(float(VIEW_MOTION_MIN_SCALE), float(effective_depth_m) / near_threshold_m)
|
| 508 |
+
forward_m = float(default_forward_m) * scale
|
| 509 |
+
radius_m = float(default_radius_m) * scale
|
| 510 |
+
forward_cap_m = float(effective_depth_m) * float(VIEW_MOTION_FORWARD_DEPTH_FRAC)
|
| 511 |
+
radius_cap_m = float(effective_depth_m) * float(VIEW_MOTION_ROTATE_DEPTH_FRAC)
|
| 512 |
+
forward_m = min(forward_m, forward_cap_m)
|
| 513 |
+
radius_m = min(radius_m, radius_cap_m)
|
| 514 |
+
return forward_m, radius_m, effective_depth_m, scale, median_depth_m, foreground_depth_m
|
| 515 |
+
|
| 516 |
+
|
| 517 |
def _render_pinhole_frame(
|
| 518 |
renderer: GSplatRenderer,
|
| 519 |
gaussians: Any,
|
|
|
|
| 658 |
step: int,
|
| 659 |
args: argparse.Namespace,
|
| 660 |
) -> None:
|
| 661 |
+
native_h, native_w = _image_hw_from_path(image_path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 662 |
camera_json_entry = _camera_json_for_image(getattr(args, "_camera_json_data", None), image_path)
|
| 663 |
+
load_max_long_edge = _initial_max_long_edge(
|
| 664 |
+
image_path=image_path,
|
| 665 |
+
args=args,
|
| 666 |
+
camera_json_entry=camera_json_entry,
|
| 667 |
+
)
|
| 668 |
+
|
| 669 |
+
for reload_attempt in range(2):
|
| 670 |
+
rgb_u8 = _load_rgb_u8(image_path, max_long_edge=load_max_long_edge)
|
| 671 |
+
_, h, w = rgb_u8.shape
|
| 672 |
+
if h < 4 or w < 4:
|
| 673 |
+
raise ValueError(f"Invalid image size for {image_path}: {tuple(rgb_u8.shape)}")
|
| 674 |
+
|
| 675 |
+
device = next(model.parameters()).device
|
| 676 |
+
image_u8 = rgb_u8.unsqueeze(0).to(device=device)
|
| 677 |
+
image = image_u8.to(torch.float32) / 255.0
|
| 678 |
+
json_camera_name = _camera_name_from_json(camera_json_entry)
|
| 679 |
+
aspect_camera_name = _camera_name_from_aspect(image_h=h, image_w=w)
|
| 680 |
+
forced_camera_name = str(args.camera).strip().lower()
|
| 681 |
+
forced_camera_name = None if forced_camera_name == "auto" else {"pinhole": "perspective", "erp": "panorama"}.get(forced_camera_name, forced_camera_name)
|
| 682 |
+
json_intrinsics = _values_from_camera_json(camera_json_entry, "intrinsics", "camera_intrinsics", "K")
|
| 683 |
+
json_camera_params = _values_from_camera_json(camera_json_entry, "camera_params", "fisheye624_params", "params")
|
| 684 |
+
explicit_intrinsics = _pinhole_intrinsics_from_values(json_intrinsics or args.camera_intrinsics, device=device)
|
| 685 |
+
explicit_camera_params = _fisheye624_params_from_values(json_camera_params or args.camera_params, device=device)
|
| 686 |
+
if explicit_intrinsics is not None and explicit_camera_params is not None:
|
| 687 |
+
raise ValueError("Use only one of --camera-intrinsics or --camera-params.")
|
| 688 |
+
|
| 689 |
+
rays: torch.Tensor | None
|
| 690 |
+
render_intrinsics: torch.Tensor | None = None
|
| 691 |
+
render_camera_params: torch.Tensor | None = None
|
| 692 |
+
if explicit_intrinsics is not None:
|
| 693 |
+
camera_kind: CameraKind = "panorama" if json_camera_name in {"panorama", "erp", "spherical"} else "perspective"
|
| 694 |
+
render_intrinsics = explicit_intrinsics
|
| 695 |
+
if camera_kind == "panorama":
|
| 696 |
+
out = _run_model_panorama(model, image, image_u8, distance_init_cap_m=0.0)
|
| 697 |
+
else:
|
| 698 |
+
out = _run_model_pinhole(
|
| 699 |
+
model,
|
| 700 |
+
image,
|
| 701 |
+
image_u8,
|
| 702 |
+
intrinsics=explicit_intrinsics,
|
| 703 |
+
distance_init_cap_m=0.0,
|
| 704 |
+
)
|
| 705 |
+
rays = out.get("geometry_rays", out.get("unik3d_gt_rays", out.get("unik3d_rays", None)))
|
| 706 |
+
stats = _ray_fov_stats(rays) if torch.is_tensor(rays) else _empty_ray_stats()
|
| 707 |
+
elif explicit_camera_params is not None:
|
| 708 |
camera_kind = "fisheye"
|
| 709 |
+
render_camera_params = explicit_camera_params
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 710 |
out = _run_model_fisheye(
|
| 711 |
model,
|
| 712 |
image,
|
| 713 |
image_u8,
|
| 714 |
+
camera_params=explicit_camera_params,
|
| 715 |
+
distance_init_cap_m=0.0,
|
| 716 |
)
|
| 717 |
+
rays = out.get("geometry_rays", out.get("unik3d_gt_rays", out.get("unik3d_rays", None)))
|
| 718 |
+
stats = _ray_fov_stats(rays) if torch.is_tensor(rays) else _empty_ray_stats()
|
| 719 |
+
elif forced_camera_name == "panorama" or (
|
| 720 |
+
forced_camera_name is None and (json_camera_name in {"panorama", "erp", "spherical"} or aspect_camera_name == "panorama")
|
| 721 |
+
):
|
| 722 |
+
camera_kind = "panorama"
|
| 723 |
+
out = _run_model_panorama(model, image, image_u8, distance_init_cap_m=0.0)
|
| 724 |
+
rays = out.get("geometry_rays", out.get("unik3d_gt_rays", out.get("unik3d_rays", None)))
|
| 725 |
+
stats = _ray_fov_stats(rays) if torch.is_tensor(rays) else _empty_ray_stats()
|
| 726 |
else:
|
| 727 |
+
rays = _predict_unik3d_rays(model, image_u8, image_h=h, image_w=w)
|
| 728 |
+
stats = _ray_fov_stats(rays)
|
| 729 |
+
if forced_camera_name == "fisheye":
|
| 730 |
+
camera_kind = "fisheye"
|
| 731 |
+
elif forced_camera_name == "perspective":
|
| 732 |
+
camera_kind = "perspective"
|
| 733 |
+
elif json_camera_name in {"fisheye", "fisheye624", "opencv_fisheye"}:
|
| 734 |
+
camera_kind = "fisheye"
|
| 735 |
+
elif json_camera_name in {"perspective", "pinhole"}:
|
| 736 |
+
camera_kind = "perspective"
|
| 737 |
+
else:
|
| 738 |
+
camera_kind = _classify_camera(stats, args)
|
| 739 |
+
if camera_kind == "panorama":
|
| 740 |
+
out = _run_model_panorama(model, image, image_u8, distance_init_cap_m=0.0)
|
| 741 |
+
elif camera_kind == "fisheye":
|
| 742 |
+
render_camera_params = fit_fisheye624_params_from_rays(rays).detach().to(device=device, dtype=torch.float32)
|
| 743 |
+
out = _run_model_fisheye(
|
| 744 |
+
model,
|
| 745 |
+
image,
|
| 746 |
+
image_u8,
|
| 747 |
+
camera_params=render_camera_params,
|
| 748 |
+
distance_init_cap_m=0.0,
|
| 749 |
+
)
|
| 750 |
+
else:
|
| 751 |
+
render_intrinsics = fit_pinhole_intrinsics_from_rays(rays).detach().to(device=device, dtype=torch.float32)
|
| 752 |
+
out = _run_model_pinhole(
|
| 753 |
+
model,
|
| 754 |
+
image,
|
| 755 |
+
image_u8,
|
| 756 |
+
intrinsics=render_intrinsics,
|
| 757 |
+
distance_init_cap_m=0.0,
|
| 758 |
+
)
|
| 759 |
+
|
| 760 |
+
needs_native_panorama = (
|
| 761 |
+
camera_kind == "panorama"
|
| 762 |
+
and (h < native_h or w < native_w)
|
| 763 |
+
and load_max_long_edge != _panorama_max_long_edge()
|
| 764 |
+
)
|
| 765 |
+
if needs_native_panorama and reload_attempt == 0:
|
| 766 |
+
load_max_long_edge = _panorama_max_long_edge()
|
| 767 |
+
continue
|
| 768 |
+
break
|
| 769 |
|
| 770 |
LOGGER.info(
|
| 771 |
"%s -> %s | hfov=%.1f vfov=%.1f diag=%.1f aspect=%.3f",
|
|
|
|
| 779 |
|
| 780 |
src_w2c = torch.eye(4, dtype=torch.float32, device=device)
|
| 781 |
gaussians_world = transform_gaussians_to_world(out["gaussians"], src_w2c)
|
| 782 |
+
model_output = out if isinstance(out, dict) else {"gaussians": out}
|
| 783 |
+
(
|
| 784 |
+
forward_distance_m,
|
| 785 |
+
rotate_radius_m,
|
| 786 |
+
scene_depth_m,
|
| 787 |
+
motion_scale,
|
| 788 |
+
median_depth_m,
|
| 789 |
+
foreground_depth_m,
|
| 790 |
+
) = _adaptive_view_motion_distances(
|
| 791 |
+
model_output,
|
| 792 |
+
default_forward_m=FORWARD_DISTANCE_M,
|
| 793 |
+
default_radius_m=ROTATE_RADIUS_M,
|
| 794 |
+
)
|
| 795 |
+
if float(motion_scale) < 0.999:
|
| 796 |
+
LOGGER.info(
|
| 797 |
+
"Near-scene view motion | depth_eff=%.3fm median=%.3fm p25=%.3fm scale=%.3f forward=%.3fm orbit=%.3fm",
|
| 798 |
+
float(scene_depth_m) if scene_depth_m is not None else float("nan"),
|
| 799 |
+
float(median_depth_m) if median_depth_m is not None else float("nan"),
|
| 800 |
+
float(foreground_depth_m) if foreground_depth_m is not None else float("nan"),
|
| 801 |
+
float(motion_scale),
|
| 802 |
+
float(forward_distance_m),
|
| 803 |
+
float(rotate_radius_m),
|
| 804 |
+
)
|
| 805 |
forward_poses = _build_forward_poses(
|
| 806 |
+
num_views=FORWARD_VIEWS,
|
| 807 |
+
distance_m=forward_distance_m,
|
| 808 |
device=device,
|
| 809 |
)
|
| 810 |
rotate_poses = _build_rotate_poses(
|
| 811 |
+
num_views=ROTATE_VIEWS,
|
| 812 |
+
radius_m=rotate_radius_m,
|
| 813 |
device=device,
|
| 814 |
)
|
| 815 |
|
| 816 |
sample_dir = out_root / _slug_from_path(image_path)
|
| 817 |
sample_dir.mkdir(parents=True, exist_ok=True)
|
| 818 |
output_crop_border_fraction = 0.0 if camera_kind == "panorama" else 0.05
|
|
|
|
| 819 |
|
| 820 |
forward_frames: list[np.ndarray] = []
|
| 821 |
rotate_frames: list[np.ndarray] = []
|
| 822 |
|
| 823 |
if camera_kind == "panorama":
|
| 824 |
+
face_w = max(16, int(min(h, w // 4)))
|
| 825 |
forward_dir = sample_dir / "forward_erp"
|
| 826 |
rotate_dir = sample_dir / "rotate_erp"
|
| 827 |
rotate_faces_dir = sample_dir / "rotate_cubemap_faces"
|
|
|
|
| 885 |
forward_frames = [_crop_border_u8(frame, output_crop_border_fraction) for frame in forward_frames]
|
| 886 |
rotate_frames = [_crop_border_u8(frame, output_crop_border_fraction) for frame in rotate_frames]
|
| 887 |
|
| 888 |
+
_save_gif(forward_frames, sample_dir / "forward.gif", duration_ms=GIF_DURATION_MS)
|
| 889 |
+
_save_gif(rotate_frames, sample_dir / "rotate.gif", duration_ms=GIF_DURATION_MS)
|
| 890 |
_save_ply_if_requested(gaussians_world, sample_dir / "gaussians.ply", f_px=f_px, image_h=h, image_w=w, enabled=bool(args.save_ply))
|
| 891 |
|
| 892 |
metadata = {
|
|
|
|
| 897 |
"ray_stats": stats,
|
| 898 |
"camera_json": str(args.camera_json) if args.camera_json is not None else None,
|
| 899 |
"camera_json_entry": camera_json_entry,
|
| 900 |
+
"aspect_camera_name": aspect_camera_name,
|
| 901 |
"explicit_camera_intrinsics": args.camera_intrinsics,
|
| 902 |
"explicit_camera_params": args.camera_params,
|
| 903 |
+
"forward_distance_m": float(forward_distance_m),
|
| 904 |
+
"rotate_radius_m": float(rotate_radius_m),
|
| 905 |
+
"forward_distance_m_default": float(FORWARD_DISTANCE_M),
|
| 906 |
+
"rotate_radius_m_default": float(ROTATE_RADIUS_M),
|
| 907 |
+
"scene_depth_for_motion_m": scene_depth_m,
|
| 908 |
+
"median_predicted_depth_m": median_depth_m,
|
| 909 |
+
"foreground_depth_p25_m": foreground_depth_m,
|
| 910 |
+
"view_motion_scale": float(motion_scale),
|
| 911 |
"rotate_path": "clockwise_camera_xy_orbit_fixed_source_orientation",
|
| 912 |
"panorama_renderer": "unisharp.cli.unified_trainer.UnifiedTrainer._render_cubemap/_cube_to_erp",
|
| 913 |
"low_pass_filter_eps": float(args.low_pass_filter_eps),
|
|
|
|
| 928 |
p.add_argument("--out-dir", type=Path, default=REPO_ROOT / "outputs" / "inference")
|
| 929 |
p.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
| 930 |
p.add_argument("--max-images", type=int, default=0)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 931 |
p.add_argument("--save-ply", action="store_true")
|
| 932 |
p.add_argument(
|
| 933 |
"--camera-json",
|
|
|
|
| 956 |
choices=["auto", "perspective", "pinhole", "fisheye", "panorama", "erp"],
|
| 957 |
help="Override automatic ray-range camera classification.",
|
| 958 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 959 |
p.add_argument("--low-pass-filter-eps", type=float, default=0.0)
|
| 960 |
return p
|
| 961 |
|