| 1 | #!/usr/bin/env python3 |
| 2 | """Estimate a conditioning image's horizontal FOV with MoGe-2 (prints JSON). |
| 3 | |
| 4 | The action model consumes camera intrinsics (``K_pix``) alongside the pose |
| 5 | trajectory; ``action_string_camera.default_k_pix`` otherwise guesses a fixed 70° |
| 6 | horizontal FOV. Feeding the *image's own* FOV makes the UCPE Plücker embedding |
| 7 | consistent with the first frame, so a given camera translation produces the pixel |
| 8 | flow the model expects instead of a systematically wide/narrow one. |
| 9 | |
| 10 | Run as a subprocess so the ~1.4 GB ViT-L never shares |
| 11 | the generator's VRAM: |
| 12 | |
| 13 | python echo_wm/helpers/moge_fov.py --image frame.png --target-width 1280 --target-height 704 |
| 14 | |
| 15 | Output (stdout, one JSON object): |
| 16 | {"fov_x_deg": 63.2, "fov_x_raw_deg": 71.5, "crop_factor": 0.83, ...} |
| 17 | |
| 18 | ``fov_x_raw_deg`` is MoGe's estimate for the image as given. ``fov_x_deg`` is the |
| 19 | value to actually use: the trainer's ``_encode_conditioning_image`` fits the image |
| 20 | to the target resolution by **resize-to-cover + center-crop**, so when the input |
| 21 | aspect is wider than the target, the sides are cropped away and the effective |
| 22 | horizontal FOV shrinks by ``a_target / a_input``. Inputs taller than the target |
| 23 | keep their full horizontal FOV (only the top/bottom are cropped). |
| 24 | """ |
| 25 | from __future__ import annotations |
| 26 | |
| 27 | import argparse |
| 28 | import json |
| 29 | import math |
| 30 | import sys |
| 31 | from pathlib import Path |
| 32 | |
| 33 | import numpy as np |
| 34 | import torch |
| 35 | from PIL import Image |
| 36 | |
| 37 | DEFAULT_MODEL = "Ruicheng/moge-2-vitl-normal" |
| 38 | |
| 39 | |
| 40 | def effective_fov_x(fov_x_raw_deg: float, in_w: int, in_h: int, target_w: int, target_h: int) -> tuple[float, float]: |
| 41 | """Horizontal FOV after resize-to-cover + center-crop to the target size. |
| 42 | |
| 43 | Returns ``(fov_x_eff_deg, crop_factor)`` where ``crop_factor`` is the retained |
| 44 | fraction of the original horizontal extent (1.0 = nothing cropped away). |
| 45 | """ |
| 46 | a_in = in_w / in_h |
| 47 | a_target = target_w / target_h |
| 48 | crop_factor = min(1.0, a_target / a_in) # >1 would mean *adding* FOV, impossible |
| 49 | fov_eff = 2.0 * math.degrees(math.atan(crop_factor * math.tan(math.radians(fov_x_raw_deg) / 2.0))) |
| 50 | return fov_eff, crop_factor |
| 51 | |
| 52 | |
| 53 | def estimate(image_path: Path, model_name: str, device: str, target_w: int, target_h: int, |
| 54 | resolution_level: int = 9) -> dict: |
| 55 | from moge.model.v2 import MoGeModel |
| 56 | |
| 57 | image = Image.open(image_path).convert("RGB") |
| 58 | in_w, in_h = image.size |
| 59 | dev = torch.device(device) |
| 60 | model = MoGeModel.from_pretrained(model_name).to(dev).eval() |
| 61 | |
| 62 | tensor = torch.tensor(np.asarray(image) / 255.0, dtype=torch.float32, device=dev).permute(2, 0, 1) |
| 63 | with torch.inference_mode(): |
| 64 | out = model.infer(tensor, resolution_level=resolution_level) |
| 65 | |
| 66 | # MoGe returns NORMALIZED intrinsics (fx = focal / width, principal point 0.5). |
| 67 | k = out["intrinsics"].float().cpu().numpy() |
| 68 | fx_n, fy_n = float(k[0, 0]), float(k[1, 1]) |
| 69 | fov_x_raw = 2.0 * math.degrees(math.atan(0.5 / fx_n)) |
| 70 | fov_y_raw = 2.0 * math.degrees(math.atan(0.5 / fy_n)) |
| 71 | fov_x_eff, crop_factor = effective_fov_x(fov_x_raw, in_w, in_h, target_w, target_h) |
| 72 | |
| 73 | return { |
| 74 | "fov_x_deg": round(fov_x_eff, 3), |
| 75 | "fov_x_raw_deg": round(fov_x_raw, 3), |
| 76 | "fov_y_raw_deg": round(fov_y_raw, 3), |
| 77 | "crop_factor": round(crop_factor, 4), |
| 78 | "fx_normalized": round(fx_n, 6), |
| 79 | "input_width": in_w, |
| 80 | "input_height": in_h, |
| 81 | "target_width": target_w, |
| 82 | "target_height": target_h, |
| 83 | # Focal the trainer will derive from fov_x_deg at the target resolution. |
| 84 | "fx_pixels_at_target": round((target_w / 2.0) / math.tan(math.radians(fov_x_eff) / 2.0), 2), |
| 85 | "model": model_name, |
| 86 | "device": str(dev), |
| 87 | } |
| 88 | |
| 89 | |
| 90 | def main() -> None: |
| 91 | p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| 92 | p.add_argument("--image", type=Path, required=True) |
| 93 | p.add_argument("--model", default=DEFAULT_MODEL, help="HF repo id or local model.pt") |
| 94 | p.add_argument("--device", default="cuda", choices=["cuda", "cpu"]) |
| 95 | p.add_argument("--target-width", type=int, default=1280) |
| 96 | p.add_argument("--target-height", type=int, default=704) |
| 97 | p.add_argument("--resolution-level", type=int, default=9, |
| 98 | help="MoGe inference resolution level (higher = finer, slower).") |
| 99 | args = p.parse_args() |
| 100 | |
| 101 | if not args.image.exists(): |
| 102 | raise SystemExit(f"image not found: {args.image}") |
| 103 | result = estimate(args.image, args.model, args.device, args.target_width, args.target_height, |
| 104 | args.resolution_level) |
| 105 | # JSON on stdout only; progress/warnings from torch go to stderr. |
| 106 | print(json.dumps(result), flush=True) |
| 107 | |
| 108 | |
| 109 | if __name__ == "__main__": |
| 110 | try: |
| 111 | main() |
| 112 | except KeyboardInterrupt: |
| 113 | sys.exit(130) |
| 114 |