CropOrPad

class CropOrPad(SpatialTransform):
    r"""Crop and/or pad to a target spatial shape.

    If the current spatial size along an axis is larger than the target, that
    axis is cropped symmetrically from both sides. If it is smaller, it is
    padded symmetrically. The affine matrix is updated so that physical
    positions of the voxels are maintained.

    The target shape can be specified in voxels (the default), millimetres,
    or centimetres. When physical units are used, the target is converted to
    voxels at transform time using the image spacing.

    When the input is a `Subject` or `Image`, the transform operates
    lazily (data is not loaded from disk until it is actually accessed).

    Args:
        target_shape: Desired spatial shape. A single `int` broadcasts
            to all three axes. When `units` is `"mm"` or `"cm"`,
            values may be floats representing the physical extent along
            each axis. Use `None` for an axis to leave it unchanged,
            e.g., `(256, 256, None)`.
        units: Coordinate system for `target_shape`. One of
            `"voxels"` (default), `"mm"`, or `"cm"`.
        padding_mode: One of `'constant'`, `'reflect'`,
            `'replicate'`, or `'circular'`. See
            [`torch.nn.functional.pad`](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html).
        fill: Fill value when `padding_mode='constant'`.
        only_crop: If `True`, padding is never applied. Mutually
            exclusive with `only_pad`.
        only_pad: If `True`, cropping is never applied. Mutually
            exclusive with `only_crop`.
        location: Where to place the crop window when the image is
            larger than the target. `"center"` (default) centres the
            window; `"random"` picks a uniformly random position.
            Padding is always centred regardless of this parameter.
        **kwargs: See [`Transform`][torchio.Transform] for additional
            keyword arguments.

    Examples:
        >>> import torchio as tio
        >>> transform = tio.CropOrPad(target_shape=(120, 80, 180))
        >>> transform = tio.CropOrPad(target_shape=256)
        >>> transform = tio.CropOrPad(target_shape=(150.0, 200.0, 180.0), units='mm')
        >>> transform = tio.CropOrPad(target_shape=(15.0, 20.0, 18.0), units='cm')
        >>> transform = tio.CropOrPad(target_shape=256, only_pad=True)
        >>> transform = tio.CropOrPad(target_shape=(256, 256, None))  # keep depth
        >>> transform = tio.CropOrPad(target_shape=96, location='random')
    """

    def __init__(
        self,
        target_shape: TargetShapeParam,
        *,
        units: Units = "voxels",
        padding_mode: str = "constant",
        fill: float = 0,
        only_crop: bool = False,
        only_pad: bool = False,
        location: Location = "center",
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if only_crop and only_pad:
            msg = "only_crop and only_pad cannot both be True"
            raise ValueError(msg)
        if units not in ("voxels", "mm", "cm"):
            msg = f"units must be 'voxels', 'mm', or 'cm', got {units!r}"
            raise ValueError(msg)
        if location not in ("center", "random"):
            msg = f"location must be 'center' or 'random', got {location!r}"
            raise ValueError(msg)
        self.target_shape = _parse_target_shape(target_shape)
        self.units: Units = units
        self.padding_mode = padding_mode
        self.fill = fill
        self.only_crop = only_crop
        self.only_pad = only_pad
        self.location: Location = location

    def forward(self, data):
        """Apply the transform.

        For `Subject` and `Image` inputs, operates lazily per-image
        without loading data from disk. For batched inputs, falls back
        to the standard `SubjectsBatch` path.
        """
        if isinstance(data, (Subject, Image)):
            return self._forward_lazy(data)
        return super().forward(data)

    def _forward_lazy(self, data: Subject | Image) -> Subject | Image:
        is_image = isinstance(data, Image)
        if is_image:
            subject = Subject(tio_default_image=data)
        else:
            assert isinstance(data, Subject)
            subject = data

        if self.copy:
            subject = _copy.deepcopy(subject)

        if torch.rand(1).item() > self.p:
            return subject.tio_default_image if is_image else subject

        first_image = next(iter(subject.images.values()))
        current_shape: TypeThreeInts = first_image.spatial_shape
        target_voxels = _to_voxels(
            self.target_shape,
            self.units,
            first_image.affine.spacing,
            current_shape,
        )

        padding, cropping = _compute_crop_and_pad(
            current_shape,
            target_voxels,
            only_crop=self.only_crop,
            only_pad=self.only_pad,
            location=self.location,
        )

        self._apply_lazy_ops(subject, padding, cropping)

        return subject.tio_default_image if is_image else subject

    def _apply_lazy_ops(
        self,
        subject: Subject,
        padding: TypeSixInts | None,
        cropping: TypeSixInts | None,
    ) -> None:
        """Apply lazy pad/crop and record history."""
        images = _get_images(subject, self.include, self.exclude)

        if padding is not None:
            for name, image in images.items():
                subject._images[name] = _pad_image_lazy(
                    image,
                    padding,
                    self.padding_mode,
                    self.fill,
                )
            subject.applied_transforms.append(
                AppliedTransform(
                    name="Pad",
                    params={
                        "padding": padding,
                        "padding_mode": self.padding_mode,
                        "fill": self.fill,
                    },
                ),
            )

        if cropping is not None:
            images = _get_images(subject, self.include, self.exclude)
            for name, image in images.items():
                subject._images[name] = _crop_image_lazy(image, cropping)
            subject.applied_transforms.append(
                AppliedTransform(name="Crop", params={"cropping": cropping}),
            )

        subject.applied_transforms.append(
            AppliedTransform(
                name="CropOrPad",
                params={"padding": padding, "cropping": cropping},
            ),
        )

    # --- Standard batch path (for SubjectsBatch, Tensor, etc.) ---

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        first_images = next(iter(batch.images.values()))
        spacing = first_images.affines[0].spacing

        data_tensor = first_images.data
        current_shape: TypeThreeInts = (
            data_tensor.shape[-3],
            data_tensor.shape[-2],
            data_tensor.shape[-1],
        )

        target_voxels = _to_voxels(
            self.target_shape,
            self.units,
            spacing,
            current_shape,
        )

        if self.units != "voxels":
            logger.debug(
                "CropOrPad target {} {} → {} voxels (spacing {} mm)",
                self.target_shape,
                self.units,
                target_voxels,
                spacing,
            )

        padding, cropping = _compute_crop_and_pad(
            current_shape,
            target_voxels,
            only_crop=self.only_crop,
            only_pad=self.only_pad,
            location=self.location,
        )

        return {"padding": padding, "cropping": cropping}

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        padding: TypeSixInts | None = params["padding"]
        cropping: TypeSixInts | None = params["cropping"]

        transforms: list[SpatialTransform] = []
        if padding is not None:
            transforms.append(
                Pad(
                    padding=padding,
                    padding_mode=self.padding_mode,
                    fill=self.fill,
                    include=self.include,
                    exclude=self.exclude,
                )
            )
        if cropping is not None:
            transforms.append(
                Crop(
                    cropping=cropping,
                    include=self.include,
                    exclude=self.exclude,
                )
            )

        if transforms:
            pipeline = Compose(transforms, copy=False)
            batch = pipeline(batch)

        return batch