model¶

def predict_pose(
    model: PoseRegressor,
    config: dict,
    img: Float[torch.Tensor, "1 1 H W"],
    sdd: float,
    delx: float,
    dely: float,
    x0: float,
    y0: float,
) -> RigidTransform:
    """Regress the camera pose of an X-ray with a trained pose regressor.

    The X-ray is resampled to the intrinsics the model was trained on, center
    cropped, and normalized before being passed through the network. This is how
    a checkpoint is used as an initializer for iterative registration.

    Args:
        model: Trained `PoseRegressor`.
        config: The checkpoint's saved training config, supplying the model's
            assumed `sdd`, `height`, and `delx`.
        img: X-ray image tensor of shape ``(1, 1, H, W)``.
        sdd: Source-to-detector distance of the X-ray, in millimeters.
        delx: Pixel spacing of the X-ray along the x-axis, in millimeters.
        dely: Pixel spacing of the X-ray along the y-axis, in millimeters.
            Must equal `delx`; non-square pixels are not supported.
        x0: Detector origin offset along the x-axis, in millimeters.
        y0: Detector origin offset along the y-axis, in millimeters.

    Returns:
        The predicted camera pose.
    """
    # Resample the X-ray image to match the model's assumed intrinsics
    img, height, width = _resample_xray(img, sdd, delx, dely, x0, y0, config)
    height = min(height, width)
    img = center_crop(img, (height, height))

    # Resize the image and normalize pixel intensities
    transforms = XrayTransforms(config["height"])
    img = transforms(img).cuda()

    # Predict pose
    with torch.no_grad():
        init_pose = model(img)

    return init_pose

def load_model(ckptpath, meta=False):
    """Load a pretrained pose regression model"""
    ckpt = torch.load(ckptpath, weights_only=False)
    config = ckpt["config"]

    model_state_dict = ckpt["model_state_dict"]
    model = PoseRegressor(
        model_name=config["model_name"],
        parameterization=config["parameterization"],
        convention=config["convention"],
        norm_layer=config["norm_layer"],
        height=config["height"],
        unit_conversion_factor=config.get("unit_conversion_factor", 1.0),
    ).cuda()
    model.load_state_dict(model_state_dict)
    model.eval()

    if meta:
        return model, config, ckpt["date"]
    else:
        return model, config

def get_random_pose(
    alphamin,
    alphamax,
    betamin,
    betamax,
    gammamin,
    gammamax,
    txmin,
    txmax,
    tymin,
    tymax,
    tzmin,
    tzmax,
    batch_size,
):
    """Generate a batch of random poses in SE(3) using specified ranges."""
    alpha = uniform(alphamin, alphamax, batch_size, circle_shift=True)
    beta = uniform(betamin, betamax, batch_size, circle_shift=True)
    gamma = uniform(gammamin, gammamax, batch_size, circle_shift=True)
    tx = uniform(txmin, txmax, batch_size)
    ty = uniform(tymin, tymax, batch_size)
    tz = uniform(tzmin, tzmax, batch_size)
    rot = torch.concat([alpha, beta, gamma], dim=1)
    xyz = torch.concat([tx, ty, tz], dim=1)
    return convert(
        rot, xyz, parameterization="euler_angles", convention="ZXY", degrees=True
    )

def initialize_subjects(
    volpath: str,  # A single CT or a directory with multiple volumes
    maskpath: str | None,  # Optional labelmaps corresponding to the CTs
    orientation: str | None,  # "AP", "PA", or None
    patch_size: tuple | None,  # Tuple for random crop sizes (h, w, d)
    num_samples: int,  # Total number of training iterations
    num_workers: int,  # Number of workers for the dataloader
    pin_memory: bool,  # Pin memory for the dataloader
    weights: tuple[float, ...] | None = None,  # Sampling probability for each volume
    replacement: bool = True,  # Sample with replacement
) -> tuple[Subject | SubjectsLoader, bool]:
    """Load the CT (or CTs) to render training DRRs from.

    A single volume is read eagerly. A directory is loaded lazily into a
    weighted-sampling dataloader, optionally serving random patches rather than
    whole volumes. Volumes and masks are matched by filename across the two
    directories.

    Returns:
        subjects: A single `Subject` if `volpath` is one file, else a loader
            yielding subjects or patches.
        single_subject: True if `volpath` was a single volume, which tells the
            trainer whether the DRR module can hold the volume itself.
    """
    # If only a single subject is passed, load it and return
    single_subject = False
    if Path(volpath).is_file():
        single_subject = True
        subject = read(volpath, maskpath, orientation=orientation)
        return subject, single_subject

    # Else, construct a list of all volumes and masks
    # We assume volumes and masks have the same name, but are in different folders
    volumes = sorted(Path(volpath).glob("[!.]*.nii.gz"))
    masks = sorted(Path(maskpath).glob("[!.]*.nii.gz")) if maskpath is not None else []
    itr = zip_longest(volumes, masks)
    pbar = tqdm(itr, desc="Lazily loading CTs...", total=len(volumes), ncols=200)

    # Lazily load a list of all subjects in the dataset
    subjects = []
    for volpath, maskpath in pbar:
        volume = ScalarImage(volpath)
        mask = LabelMap(maskpath) if maskpath is not None else None
        subject = Subject(volume=volume, mask=mask)
        subjects.append(subject)

    # Construct a dataloader with efficient IO
    subjects = SubjectsDataset(subjects)
    if weights is None:
        weights = [1 / len(volumes) for _ in range(len(volumes))]
    subject_sampler = WeightedRandomSampler(
        weights, replacement=replacement, num_samples=num_samples
    )

    # Return entire volumes
    if patch_size is None:
        subjects = SubjectsLoader(
            subjects,
            sampler=subject_sampler,
            num_workers=num_workers,
            pin_memory=pin_memory,
        )
        return subjects, single_subject

    # Return random crops
    patch_sampler = UniformSampler(patch_size)
    patches_queue = Queue(
        subjects,
        max_length=64,
        samples_per_volume=4,
        sampler=patch_sampler,
        subject_sampler=subject_sampler,
        shuffle_subjects=False,
        num_workers=num_workers,
    )

    subjects = SubjectsLoader(
        patches_queue,
        batch_size=1,
        num_workers=0,
        pin_memory=pin_memory,
    )

    return subjects, single_subject

def initialize_modules(
    model_name,
    pretrained,
    parameterization,
    convention,
    norm_layer,
    unit_conversion_factor,
    sdd,
    height,
    delx,
    orientation,
    reverse_x_axis,
    lr,
    n_total_itrs,
    n_warmup_itrs,
    n_grad_accum_itrs,
    subject,
    disable_scheduler,
    ckptpath,
    reuse_optimizer,
) -> tuple[
    PoseRegressor,
    DRR,
    Compose,
    Optimizer,
    IdentitySchedule | WarmupCosineSchedule,
    int,
    int,
]:
    """Construct the model, renderer, and optimizer for a training run.

    Restores model weights from `ckptpath` when given, and the optimizer and
    scheduler state as well when `reuse_optimizer` is set. With more than one
    subject the DRR module is built against a dummy CT, since each step swaps in
    its own volume.

    Returns:
        model: The pose regressor.
        drr: The differentiable renderer.
        transforms: X-ray preprocessing shared with inference.
        optimizer: The optimizer.
        scheduler: The learning rate scheduler.
        start_itr: Iteration to resume from, 0 unless reusing an optimizer.
        model_number: Checkpoint number to resume from.
    """
    # Initialize the pose regression model
    model = PoseRegressor(
        model_name=model_name,
        pretrained=pretrained,
        parameterization=parameterization,
        convention=convention,
        norm_layer=norm_layer,
        height=height,
        unit_conversion_factor=unit_conversion_factor,
    ).cuda()

    # If a checkpoint is passed, reload the model state
    ckpt, start_itr, model_number = _load_checkpoint(ckptpath, reuse_optimizer)
    if ckpt is not None:
        print("Loading previous model weights...")
        model.load_state_dict(ckpt["model_state_dict"])

    # Initialize the optimizer and learning rate scheduler
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    if disable_scheduler:
        scheduler = IdentitySchedule(optimizer)
    else:
        warmup_itrs = n_warmup_itrs / n_grad_accum_itrs
        total_itrs = n_total_itrs / n_grad_accum_itrs
        scheduler = WarmupCosineSchedule(optimizer, warmup_itrs, total_itrs)

    # Optionally, reload the optimizer and scheduler
    if ckpt is not None and reuse_optimizer:
        print("Reinitializing optimizer...")
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    model.train()

    # If more than one subject is provided, initialize the DRR module with a dummy CT
    drr = DRR(
        subject if subject is not None else load_example_ct(orientation=orientation),
        sdd=sdd,
        height=height,
        delx=delx,
        reverse_x_axis=reverse_x_axis,
        renderer="trilinear",
    )
    drr.density = None  # Unload the precomputed density map to free up memory
    if not hasattr(drr, "mask"):
        drr.mask = None
    transforms = XrayTransforms(height)

    # If a single subject was provided, expose their volume and isocenter in the DRR module
    if subject is not None:
        drr.register_buffer("volume", subject.volume.data.squeeze())
        drr.register_buffer("center", torch.tensor(subject.volume.get_center())[None])
    drr = drr.cuda().to(torch.float32)

    return model, drr, transforms, optimizer, scheduler, start_itr, model_number

def initialize_coordinate_frame(
    warp: str | Path | None,
    img: str | Path,
    invert: bool,
) -> RigidTransform | None:
    """Read the rigid transform reframing a CT into a checkpoint's frame.

    Args:
        warp: Path to a SimpleITK transform, or None for no reframing.
        img: Path to the CT the transform is defined against.
        invert: If True, invert the transform.

    Returns:
        The transform, or None if `warp` is None.
    """
    if warp is None:
        return None
    return read_rigid_transform(warp, img, invert).cuda()