Skip to content

rslearn.train.transforms.pad

pad

Pad transform.

Pad

Bases: Transform

Pad (or crop) inputs to a fixed size.

Source code in rslearn/train/transforms/pad.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
class Pad(Transform):
    """Pad (or crop) inputs to a fixed size."""

    def __init__(
        self,
        size: int | tuple[int, int],
        mode: str = "topleft",
        image_selectors: list[str] = ["image"],
        box_selectors: list[str] = [],
        skip_missing: bool = False,
    ):
        """Initialize a new Crop.

        Result will be (input - mean) / std.

        Args:
            size: the size to pad to, or a min/max range of pad sizes. If the image is
                larger than this size, then it is cropped instead.
            mode: "topleft" (default) to only apply padding on the bottom and right
                sides, or "center" to apply padding equally on all sides.
            image_selectors: image items to transform.
            box_selectors: boxes items to transform.
            skip_missing: if True, skip selectors that don't exist in the input/target
                dicts. Useful when working with optional inputs.
        """
        super().__init__(skip_missing=skip_missing)
        if isinstance(size, int):
            self.size = (size, size + 1)
        else:
            self.size = size

        self.mode = mode
        self.image_selectors = image_selectors
        self.box_selectors = box_selectors

    def sample_state(self) -> dict[str, Any]:
        """Randomly decide how to transform the input.

        Returns:
            dict of sampled choices
        """
        return {"size": torch.randint(low=self.size[0], high=self.size[1], size=())}

    def apply_image(self, image: RasterImage, state: dict[str, bool]) -> RasterImage:
        """Apply the sampled state on the specified image.

        Args:
            image: the image to transform.
            state: the sampled state.
        """
        size = state["size"]
        horizontal_extra = size - image.shape[-1]
        vertical_extra = size - image.shape[-2]

        def apply_padding(
            im: torch.Tensor, horizontal: bool, before: int, after: int
        ) -> torch.Tensor:
            # Before/after must either be both non-negative or both negative.
            # >=0 indicates padding while <0 indicates cropping.
            assert (before < 0 and after <= 0) or (before >= 0 and after >= 0)
            if before > 0:
                # Padding.
                if horizontal:
                    padding_tuple: tuple = (before, after)
                else:
                    padding_tuple = (before, after, 0, 0)
                return torch.nn.functional.pad(im, padding_tuple)
            else:
                # Cropping.
                if horizontal:
                    return torchvision.transforms.functional.crop(
                        im,
                        top=0,
                        left=-before,
                        height=im.shape[-2],
                        width=im.shape[-1] + before + after,
                    )
                else:
                    return torchvision.transforms.functional.crop(
                        im,
                        top=-before,
                        left=0,
                        height=im.shape[-2] + before + after,
                        width=im.shape[-1],
                    )

        if self.mode == "topleft":
            horizontal_pad = (0, horizontal_extra)
            vertical_pad = (0, vertical_extra)

        elif self.mode == "center":
            horizontal_half = horizontal_extra // 2
            vertical_half = vertical_extra // 2
            horizontal_pad = (horizontal_half, horizontal_extra - horizontal_half)
            vertical_pad = (vertical_half, vertical_extra - vertical_half)

        image.image = apply_padding(
            image.image, True, horizontal_pad[0], horizontal_pad[1]
        )
        image.image = apply_padding(
            image.image, False, vertical_pad[0], vertical_pad[1]
        )
        return image

    def apply_boxes(self, boxes: Any, state: dict[str, bool]) -> torch.Tensor:
        """Apply the sampled state on the specified image.

        Args:
            boxes: the boxes to transform.
            state: the sampled state.
        """
        raise NotImplementedError

    def forward(
        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Apply transform over the inputs and targets.

        Args:
            input_dict: the input
            target_dict: the target

        Returns:
            transformed (input_dicts, target_dicts) tuple
        """
        state = self.sample_state()
        self.apply_fn(
            self.apply_image, input_dict, target_dict, self.image_selectors, state=state
        )
        self.apply_fn(
            self.apply_boxes, input_dict, target_dict, self.box_selectors, state=state
        )
        return input_dict, target_dict

sample_state 

sample_state() -> dict[str, Any]

Randomly decide how to transform the input.

Returns:

Type Description
dict[str, Any]

dict of sampled choices
Source code in rslearn/train/transforms/pad.py 
48
49
50
51
52
53
54
def sample_state(self) -> dict[str, Any]:
    """Randomly decide how to transform the input.

    Returns:
        dict of sampled choices
    """
    return {"size": torch.randint(low=self.size[0], high=self.size[1], size=())}

apply_image 

apply_image(image: RasterImage, state: dict[str, bool]) -> RasterImage

Apply the sampled state on the specified image.

Parameters:

Name Type Description Default
image RasterImage

the image to transform.

 required
state dict[str, bool]

the sampled state.

 required
Source code in rslearn/train/transforms/pad.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def apply_image(self, image: RasterImage, state: dict[str, bool]) -> RasterImage:
    """Apply the sampled state on the specified image.

    Args:
        image: the image to transform.
        state: the sampled state.
    """
    size = state["size"]
    horizontal_extra = size - image.shape[-1]
    vertical_extra = size - image.shape[-2]

    def apply_padding(
        im: torch.Tensor, horizontal: bool, before: int, after: int
    ) -> torch.Tensor:
        # Before/after must either be both non-negative or both negative.
        # >=0 indicates padding while <0 indicates cropping.
        assert (before < 0 and after <= 0) or (before >= 0 and after >= 0)
        if before > 0:
            # Padding.
            if horizontal:
                padding_tuple: tuple = (before, after)
            else:
                padding_tuple = (before, after, 0, 0)
            return torch.nn.functional.pad(im, padding_tuple)
        else:
            # Cropping.
            if horizontal:
                return torchvision.transforms.functional.crop(
                    im,
                    top=0,
                    left=-before,
                    height=im.shape[-2],
                    width=im.shape[-1] + before + after,
                )
            else:
                return torchvision.transforms.functional.crop(
                    im,
                    top=-before,
                    left=0,
                    height=im.shape[-2] + before + after,
                    width=im.shape[-1],
                )

    if self.mode == "topleft":
        horizontal_pad = (0, horizontal_extra)
        vertical_pad = (0, vertical_extra)

    elif self.mode == "center":
        horizontal_half = horizontal_extra // 2
        vertical_half = vertical_extra // 2
        horizontal_pad = (horizontal_half, horizontal_extra - horizontal_half)
        vertical_pad = (vertical_half, vertical_extra - vertical_half)

    image.image = apply_padding(
        image.image, True, horizontal_pad[0], horizontal_pad[1]
    )
    image.image = apply_padding(
        image.image, False, vertical_pad[0], vertical_pad[1]
    )
    return image

apply_boxes 

apply_boxes(boxes: Any, state: dict[str, bool]) -> Tensor

Apply the sampled state on the specified image.

Parameters:

Name Type Description Default
boxes Any

the boxes to transform.

 required
state dict[str, bool]

the sampled state.

 required
Source code in rslearn/train/transforms/pad.py 
117
118
119
120
121
122
123
124
def apply_boxes(self, boxes: Any, state: dict[str, bool]) -> torch.Tensor:
    """Apply the sampled state on the specified image.

    Args:
        boxes: the boxes to transform.
        state: the sampled state.
    """
    raise NotImplementedError

forward 

forward(input_dict: dict[str, Any], target_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]

Apply transform over the inputs and targets.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input

 required
target_dict dict[str, Any]

the target

 required

Returns:

Type Description
tuple[dict[str, Any], dict[str, Any]]

transformed (input_dicts, target_dicts) tuple
Source code in rslearn/train/transforms/pad.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def forward(
    self, input_dict: dict[str, Any], target_dict: dict[str, Any]
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Apply transform over the inputs and targets.

    Args:
        input_dict: the input
        target_dict: the target

    Returns:
        transformed (input_dicts, target_dicts) tuple
    """
    state = self.sample_state()
    self.apply_fn(
        self.apply_image, input_dict, target_dict, self.image_selectors, state=state
    )
    self.apply_fn(
        self.apply_boxes, input_dict, target_dict, self.box_selectors, state=state
    )
    return input_dict, target_dict