io¶

def read_xray(
    filename: Path,
    crop: int = 0,
    subtract_background: bool = False,
    linearize: bool = True,
    reducefn: str | int | Callable = "max",
    radius: float | None = None,
) -> tuple[Float[torch.Tensor, "1 1 H W"], Intrinsics, bool]:
    """Read and preprocess an X-ray image from a DICOM file.

    Loads the pixel array, parses imaging system intrinsics, optionally reorients
    the image, reduces any temporal dimension, and applies configurable preprocessing.

    Args:
        filename: Path to the DICOM file.
        crop: Total number of pixels to remove from each spatial dimension.
            For example, ``crop=10`` removes 5 pixels from each edge.
        subtract_background: If True, subtract the mode image intensity and
            clamp the result to [0, 1].
        linearize: If True, convert the image from exponential to linear form
            by applying a log transform.
        reducefn: Controls how a multiframe (5D) DICOM is collapsed to a single
            2D image. Accepted values:

            - ``"max"``: max intensity projection over the temporal dimension.
            - ``"sum"``: sum projection over the temporal dimension.
            - ``int``: index of the frame to extract.
            - ``Callable``: arbitrary function applied to the 5D tensor.
            - ``None``: no reduction; tensor remains 5D.
        radius: If given, treat the detector as circular with this radius (in
            pixels), centered on the image. Pixels outside the circle are set
            to the minimum intensity of the processed image.

    Returns:
        img: Preprocessed image tensor of shape ``(1, 1, H, W)``.
        intrinsics: Imaging system intrinsic parameters.
        pf_to_af: True if the image was reoriented from posterior-foot
            to anterior-foot orientation.
    """
    # Get the image
    ds = dcmread(filename)
    img = torch.from_numpy(ds.pixel_array.astype(np.int32)).to(torch.float32)[None, None]

    # Get the C-arm intrinsics and flip the columns if image is PF
    intrinsics, pf_to_af = _parse_dicom_intrinsics(ds)
    if pf_to_af:
        img = img.flip(-1)

    # Preprocess the X-ray image
    img = _preprocess_xray(img, crop, subtract_background, linearize)

    # Reduce a temporal dimension
    if img.ndim == 5:
        img = _reduce_frames(img, reducefn)

    # Isolate the circular detector area
    if radius is not None:
        img = _mask_outside_circle(img, radius)

    return img, intrinsics, pf_to_af

def parse_dicom_pose(
    filename: str | Path,
    orientation: str | None = "AP",
    device: str = "cuda",
):
    """Convert DICOM pose params to a C-arm SE(3) pose."""
    multiplier = -1 if orientation == "PA" else 1

    ds = dcmread(filename)
    alpha = float(ds.PositionerPrimaryAngle)
    beta = float(ds.PositionerSecondaryAngle)
    sid = multiplier * float(ds.DistanceSourceToPatient)

    return convert(
        torch.tensor([[alpha, beta, 0.0]], device=device),
        torch.tensor([[0.0, sid, 0.0]], device=device),
        parameterization="euler_angles",
        convention="ZXY",
        degrees=True,
    )