app.py

from math import ceil, floor, pi
from pathlib import Path
from textwrap import dedent
from urllib.parse import quote

import gradio as gr
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image, ImageOps
from torch import Tensor

from anycalib.cameras import CameraFactory
from anycalib.model import AnyCalib
from anycalib.visualization import viz_2d

matplotlib.use("Agg")

RAD2DEG = 180 / pi
DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_DIR = Path("/models")  # mounted from https://huggingface.co/javrtg/AnyCalib

ASSETS_DIR = Path(__file__).resolve().parent / "assets"
gr.set_static_paths(paths=[ASSETS_DIR])

method_img = (
    f"/gradio_api/file={quote(str((ASSETS_DIR / 'method_dark.png').resolve()))}"
)
description = dedent(f"""
<div style="text-align: center;">
    <h1>
        AnyCalib:<br>
        On-Manifold Learning for Model-Agnostic Single-View Camera Calibration
    </h1>
    <h3>
        <a href="https://javrtg.github.io/">Javier Tirado-Garín</a>
        &emsp;&emsp;
        <a href="https://scholar.google.com/citations?user=j_sMzokAAAAJ">Javier Civera</a><br>
        I3A, University of Zaragoza
    </h3>
    <img
        src="{method_img}"
        style="display: block; width: 60%; max-width: 900px; margin: 1rem auto;"
        alt="AnyCalib method overview"
    >
    <h3>
        Camera calibration from a single perspective/edited/distorted image using a freely chosen camera model<br>
        <a href="https://github.com/javrtg/AnyCalib">Code</a>
        &emsp;
        <a href="https://arxiv.org/abs/2503.12701">Paper</a>
    </h3>
</div>
""")


def get_model(model_id: str) -> AnyCalib:
    model = AnyCalib(model_id=None)
    ckpt_path = MODEL_DIR / f"{model_id}.pt"
    state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    model.load_state_dict(state_dict, strict=True)
    model.eval()
    return model.to(DEV)


MODELS = {
    "anycalib_pinhole": get_model("anycalib_pinhole"),
    "anycalib_gen": get_model("anycalib_gen"),
    "anycalib_dist": get_model("anycalib_dist"),
    "anycalib_edit": get_model("anycalib_edit"),
}

CAM_DISTCOEFFS = {  # (min, max, default)
    "pinhole": (0, 0, 0),
    "radial": (1, 4, 1),
    "kb": (1, 4, 4),
    "ucm": (1, 1, 1),
    "eucm": (2, 2, 2),
    "division": (1, 4, 1),
}


def get_cam_model(cam_model: str, num_dist_params: int, simple: bool) -> str:
    cam_id = cam_model
    if simple:
        cam_id = "simple_" + cam_model
    if cam_model not in ("pinhole", "ucm", "eucm"):
        cam_id += f":{num_dist_params}"
    return cam_id


def log_calibration(pred: dict, cam_id: str, im_w: int, im_h: int) -> str:
    intrinsics: Tensor = pred["intrinsics"]
    cam = CameraFactory.create_from_id(cam_id)
    txt = "Intrinsics:\n"
    for p, v in zip(cam.PARAMS_IDX, intrinsics.tolist()):
        txt += f"\t{p}: {v:.6f}"
        if p in ("fx", "fy", "f", "cx", "cy"):
            txt += "\tpixels"
        txt += "\n"

    hfov, valid_h = cam.get_hfov(intrinsics, im_w)
    vfov, valid_v = cam.get_vfov(intrinsics, im_h)
    if not valid_h or not valid_v:
        raise gr.Error("Error computing FoV", duration=5)

    txt += "\nField of View:\n"
    txt += f"\thFov: {RAD2DEG * hfov:.1f}°\n"
    txt += f"\tvFov: {RAD2DEG * vfov:.1f}°\n"
    return txt


def choose_contours_every(field, base=5, target_lines=9):
    vmin = base * floor(field.amin().item() / base)
    vmax = base * ceil(field.amax().item() / base)
    ideal_spacing = (vmax - vmin) / (target_lines - 1)
    contours_every = base * ceil(ideal_spacing / base)
    return max(contours_every, base)


@torch.inference_mode()
def run_demo(
    image: Image.Image | None,
    model_id: str,
    simple: bool,
    cam_model: str,
    num_dist_params: int,
    undist_scale: float,
    undist_proj: str,
):
    if image is None:
        raise gr.Error("Upload an image first.", duration=5)
    assert isinstance(image, Image.Image)

    im = np.array(ImageOps.exif_transpose(image).convert("RGB"))  # type: ignore
    im = torch.tensor(im, dtype=torch.float32, device=DEV).permute(2, 0, 1) / 255

    cam_id = get_cam_model(cam_model, num_dist_params, simple)
    pred = MODELS[model_id].predict(im, cam_id=cam_id)

    summary = log_calibration(pred, cam_id, *image.size)

    # fov field
    base = 5
    fig_ff, axs_ff = plt.subplots(ncols=2)
    fov_field: Tensor = pred["fov_field"].reshape((*pred["pred_size"], 2))
    im_plot = MODELS[model_id].set_im_size(im[None], pred["pred_size"])[0][0]
    for i in range(2):
        viz_2d.plot_image(axs_ff[i], im_plot.permute(1, 2, 0).cpu())
        ff = (fov_field[..., i] * RAD2DEG).cpu()
        viz_2d.plot_heatmap(
            axs_ff[i],
            ff,
            contours_every=choose_contours_every(ff, base=base, target_lines=15),
            cmap="Spectral_r",
            vmin=base * floor(ff.amin().item() / base),
            vmax=base * ceil(ff.amax().item() / base),
            label_decimals=0,
        )
        viz_2d.plot_text(
            axs_ff[i],
            r"$\theta_x$" if i == 0 else r"$\theta_y$",
            pos=(0.01, 0.99),
            ha="left",
            va="top",
            color="k",
            lw=0,
            bbox=dict(
                facecolor="white",
                alpha=0.5,
                ec="none",
                boxstyle="round,pad=0.2",
            ),
            zorder=10,
        )
        axs_ff[i].patch.set_alpha(0)
    fig_ff.patch.set_alpha(0)
    fig_ff.tight_layout(pad=0.5)

    # polar angles
    base = 5
    fig_pa, ax_pa = plt.subplots()
    angles = torch.linalg.norm(fov_field, dim=-1).cpu() * RAD2DEG
    viz_2d.plot_image(ax_pa, im_plot.permute(1, 2, 0).cpu())
    viz_2d.plot_heatmap(
        ax_pa,
        angles,
        contours_every=choose_contours_every(angles, base=base, target_lines=15),
        cmap="Spectral_r",
        vmin=0,
        vmax=base * ceil(angles.amax().item() / base),
        label_decimals=0,
    )
    ax_pa.patch.set_alpha(0)
    fig_pa.patch.set_alpha(0)
    fig_pa.tight_layout(pad=0.5)

    # undistortion figure
    cam = CameraFactory.create_from_id(cam_id)
    undist_im = cam.undistort_image(
        im,
        pred["intrinsics"],
        scale=undist_scale,
        target_proj=(undist_proj if undist_proj != "pinhole" else "perspective"),
        add_alpha=True,
    )
    undist_im = undist_im.cpu().permute(1, 2, 0).numpy()

    plt.close(fig_ff)
    plt.close(fig_pa)
    return (fig_ff, fig_pa, undist_im, summary)


def build_demo() -> gr.Blocks:
    def update_distortion_range(camera_model):
        min_dist, max_dist, default_dist = CAM_DISTCOEFFS[camera_model]
        num_dist_params = gr.Number(
            value=default_dist,
            minimum=min_dist,
            maximum=max_dist,
            precision=0,
            label="# distortion coefficients",
            interactive=min_dist != max_dist,
        )
        return num_dist_params

    with gr.Blocks(
        title="AnyCalib demo",
        fill_height=True,
        delete_cache=(3600, 3600),
    ) as demo:
        gr.HTML(description)
        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(
                    label="Image",
                    type="pil",
                    image_mode="RGB",
                    height=420,
                )
                with gr.Row():
                    with gr.Group():
                        gr.Markdown("#### AnyCalib model")
                        model_id = gr.Dropdown(
                            choices=list(MODELS),
                            value="anycalib_gen",
                            show_label=False,
                        )

                with gr.Row():
                    with gr.Group():
                        gr.Markdown("#### Camera model configuration")
                        with gr.Column():
                            simple = gr.Checkbox(
                                value=False,
                                label="one focal length",
                                scale=1,
                            )

                            cam_model = gr.Dropdown(
                                choices=list(CAM_DISTCOEFFS),
                                value="pinhole",
                                label="model",
                                scale=1,
                            )

                            num_dist_params = gr.Number(
                                value=0,
                                minimum=0,
                                maximum=0,
                                precision=0,
                                label="# distortion coefficients",
                                scale=1,
                            )

                            cam_model.change(
                                fn=update_distortion_range,
                                inputs=cam_model,
                                outputs=num_dist_params,
                            )

                    with gr.Group():
                        gr.Markdown("#### Undistortion configuration")
                        with gr.Column():
                            undist_scale = gr.Number(
                                value=1,
                                minimum=0.2,
                                maximum=5,
                                step=0.1,
                                precision=1,
                                label="focal length factor",
                                scale=1,
                            )

                            undist_proj = gr.Dropdown(
                                choices=(
                                    "pinhole",
                                    "equisolid",
                                    "equidistant",
                                    "stereographic",
                                    "orthographic",
                                ),
                                value="pinhole",
                                label="target projection",
                                scale=1,
                            )

                run = gr.Button("Calibrate", variant="primary")
                # fmt:off
                gr.Examples(
                    examples=[
                        ["assets/seo.jpg", "anycalib_pinhole", False, "pinhole", 0, 1, "pinhole"],
                        ["assets/cat.jpg", "anycalib_pinhole", False, "radial", 1, 1, "pinhole"],
                        ["assets/park.jpg", "anycalib_gen", False, "eucm", 2, 0.5, "pinhole"],
                        ["assets/road.jpg", "anycalib_gen", False, "kb", 4, 0.5, "pinhole"],
                    ],
                    inputs=[image, model_id, simple, cam_model, num_dist_params, undist_scale, undist_proj],
                    label="Load example images and configurations",
                )
                # fmt:on
                gr.Markdown(
                    """
                    **Image credits**, in order:
                    1. [Sergei Gussev](https://www.flickr.com/photos/sergeigussev/49521010167/), [CC BY 2.0](https://creativecommons.org/licenses/by/2.0/deed.en)
                    2. [Michele_Sacchet](https://explorecams.com/photos/pDWj91ExFb?lens=iphone-xs-back-dual-camera-6mm-f-2-4)
                    3. [crystal~♥ Yang](https://explorecams.com/photos/epidmHVSwD?lens=samsung-nx-10mm-f3-5-fisheye)
                    4. [Imre Farago](https://explorecams.com/photos/dYXHUevDXw?lens=samsung-nx-10mm-f3-5-fisheye)
                    """
                )

            with gr.Column(scale=1):
                with gr.Tabs():
                    with gr.Tab("FoV Field"):
                        fov_field_plot = gr.Plot(show_label=False, container=False)
                    with gr.Tab("Polar Angles"):
                        polar_angles_plot = gr.Plot(show_label=False, container=False)
                    with gr.Tab("Undistortion"):
                        undist_im = gr.Image(
                            type="numpy",
                            height=420,
                            show_label=False,
                            container=False,
                        )

                summary = gr.Textbox(
                    label="Prediction log",
                    lines=10,
                    buttons=["copy"],
                    interactive=False,
                )

        run.click(
            fn=run_demo,
            inputs=[
                image,
                model_id,
                simple,
                cam_model,
                num_dist_params,
                undist_scale,
                undist_proj,
            ],
            outputs=[
                fov_field_plot,
                polar_angles_plot,
                undist_im,
                summary,
            ],
            concurrency_limit=1,
        )

    return demo


demo = build_demo()


if __name__ == "__main__":
    demo.queue(max_size=8, default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        theme=gr.themes.Soft(),  # type: ignore
    )