register¶

def load_loss_function(loss: str | torch.nn.Module, **kwargs: Any) -> torch.nn.Module:
    """Initialize a loss function for 2D/3D registration.

    Args:
        loss: Either the name of a loss function in `METRICS` or a callable
            Module whose forward method implements the loss.
        **kwargs: Optional arguments for the constructor of the loss function.
    """
    if isinstance(loss, torch.nn.Module):
        return loss
    try:
        return METRICS[loss](**kwargs)
    except KeyError:
        raise ValueError(f"Unknown metric '{loss}'. Available: {list(METRICS.keys())}")

def gif(
    logger: OptimizationLogger,
    width: int = 600,
    duration: int = 500,
) -> RegistrationOutput:
    """Animate between GT and predictions as an infinitely looping GIF.

    Args:
        logger: OptimizationLogger with .gt, .init_pose, .final_pose, and .drr.
        width: Display width in pixels when rendered in Jupyter.
        duration: Time per frame in milliseconds.

    Returns:
        RegistrationOutput that renders inline in Jupyter and can be saved to disk.
    """
    with torch.inference_mode():
        img_init = logger.drr(logger.init_pose)
        img_final = logger.drr(logger.final_pose)

    gt_pil = _to_pil(logger.gt)
    init_pil = _to_pil(img_init)
    fin_pil = _to_pil(img_final)

    frame1 = Image.fromarray(np.hstack([np.array(gt_pil), np.array(gt_pil)]))
    frame2 = Image.fromarray(np.hstack([np.array(init_pil), np.array(fin_pil)]))

    buf = BytesIO()
    frame1.save(buf, format="GIF", save_all=True, append_images=[frame2], loop=0, duration=duration)
    buf.seek(0)
    return RegistrationOutput(buf.read(), fmt="gif", width=width)

def plot(
    logger: OptimizationLogger,
    method: VizMethod = "edges",
    histeq_strength: float = 0.5,
    width: int = 900,
    **kwargs: Any,
) -> RegistrationOutput:
    """Visualise registration results (GT | GT vs initial pred | GT vs final pred).

    Args:
        logger: OptimizationLogger with .gt, .init_pose, .final_pose, and .drr.
        method: Comparison method, one of "overlay", "checkerboard", "edges".
        histeq_strength: Histogram equalisation strength (0.0 = off, 1.0 = full).
        width: Display width in pixels when rendered in Jupyter.
        **kwargs: Forwarded to the chosen method function. "checkerboard" accepts
            n_patches (int, default 8); "edges" accepts sigma (float, default 2.0).

    Returns:
        RegistrationOutput that renders inline in Jupyter and can be saved to disk.
    """
    if method not in {"overlay", "checkerboard", "edges"}:
        raise ValueError(f"method must be 'overlay', 'checkerboard', or 'edges'; got {method!r}")

    with torch.inference_mode():
        img_init = logger.drr(logger.init_pose)
        img_final = logger.drr(logger.final_pose)

    gt_pil, init_pil, fin_pil = _histeq_joint(
        [_to_pil(t) for t in [logger.gt, img_init, img_final]],
        strength=histeq_strength,
    )

    viz_fn = {"overlay": _overlay, "checkerboard": _checkerboard, "edges": _edges}[method]

    _dpi = 100  # internal only — cancels out in figsize, controls rendering sharpness
    h, w = logger.gt.squeeze().shape[-2:]
    fig, axs = plt.subplots(ncols=3, figsize=(3 * w / _dpi, h / _dpi), dpi=_dpi)

    plot_drr(logger.gt, axs=axs[0], ticks=False)
    axs[0].set_aspect("auto")

    for ax, pred_pil in zip(axs[1:], [init_pil, fin_pil]):
        ax.imshow(viz_fn(pred_pil, gt_pil, **kwargs), aspect="auto")
        ax.set_xticks([])
        ax.set_yticks([])

    fig.subplots_adjust(wspace=0, hspace=0, top=1, bottom=0, left=0, right=1)
    buf = BytesIO()
    fig.savefig(buf, format="png", dpi=_dpi)
    plt.close(fig)
    buf.seek(0)
    return RegistrationOutput(buf.read(), fmt="png", width=width)

def parse_scales(scales: list[float], crop: int, height: int) -> list[float]:
    """Convert absolute downscale factors to sequential rescale ratios for an image pyramid.

    Args:
        scales: Absolute downscale factors relative to the original image (e.g. [8, 4, 2, 1]).
            A scale of 1.0 snaps back to full cropped resolution.
        crop: Total pixels cropped from the original image (crop/2 from each side).
        height: Height of the cropped image in pixels.

    Returns:
        Sequential rescale ratios between consecutive pyramid levels.
    """
    pyramid = [1.0] + [1.0 if x == 1.0 else x * (height / (height + crop)) for x in scales]
    return [pyramid[idx] / pyramid[idx + 1] for idx in range(len(pyramid) - 1)]