Skip to content

rslearn.train.transforms.flip

flip

Flip transform.

Flip

Bases: Transform

Flip inputs horizontally and/or vertically.

Source code in rslearn/train/transforms/flip.py 
 12
 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
class Flip(Transform):
    """Flip inputs horizontally and/or vertically."""

    def __init__(
        self,
        horizontal: bool = True,
        vertical: bool = True,
        image_selectors: list[str] = ["image"],
        box_selectors: list[str] = [],
        skip_missing: bool = False,
    ):
        """Initialize a new Flip.

        Args:
            horizontal: whether to randomly flip horizontally
            vertical: whether to randomly flip vertically
            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)
        self.horizontal = horizontal
        self.vertical = vertical
        self.image_selectors = image_selectors
        self.box_selectors = box_selectors

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

        Returns:
            dict of sampled choices
        """
        horizontal = False
        if self.horizontal:
            horizontal = torch.randint(low=0, high=2, size=()) == 0
        vertical = False
        if self.vertical:
            vertical = torch.randint(low=0, high=2, size=()) == 0
        return {
            "horizontal": horizontal,
            "vertical": vertical,
        }

    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.
        """
        if state["horizontal"]:
            image.image = torch.flip(image.image, dims=[-1])
        if state["vertical"]:
            image.image = torch.flip(image.image, dims=[-2])
        return image

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

        Args:
            boxes: the boxes to transform.
            state: the sampled state.
        """
        if state["horizontal"]:
            boxes["boxes"] = torch.stack(
                [
                    boxes["width"] - boxes["boxes"][:, 2],
                    boxes["boxes"][:, 1],
                    boxes["width"] - boxes["boxes"][:, 0],
                    boxes["boxes"][:, 3],
                ],
                dim=1,
            )
        if state["vertical"]:
            boxes["boxes"] = torch.stack(
                [
                    boxes["boxes"][:, 0],
                    boxes["height"] - boxes["boxes"][:, 3],
                    boxes["boxes"][:, 2],
                    boxes["height"] - boxes["boxes"][:, 1],
                ],
                dim=1,
            )
        return boxes

    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, bool]

Randomly decide how to transform the input.

Returns:

Type Description
dict[str, bool]

dict of sampled choices
Source code in rslearn/train/transforms/flip.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def sample_state(self) -> dict[str, bool]:
    """Randomly decide how to transform the input.

    Returns:
        dict of sampled choices
    """
    horizontal = False
    if self.horizontal:
        horizontal = torch.randint(low=0, high=2, size=()) == 0
    vertical = False
    if self.vertical:
        vertical = torch.randint(low=0, high=2, size=()) == 0
    return {
        "horizontal": horizontal,
        "vertical": vertical,
    }

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/flip.py 
56
57
58
59
60
61
62
63
64
65
66
67
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.
    """
    if state["horizontal"]:
        image.image = torch.flip(image.image, dims=[-1])
    if state["vertical"]:
        image.image = torch.flip(image.image, dims=[-2])
    return image

apply_boxes 

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

Apply the sampled state on the specified image.

Parameters:

Name Type Description Default
boxes dict[str, Tensor]

the boxes to transform.

 required
state dict[str, bool]

the sampled state.

 required
Source code in rslearn/train/transforms/flip.py 
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
def apply_boxes(
    self, boxes: dict[str, torch.Tensor], state: dict[str, bool]
) -> dict[str, torch.Tensor]:
    """Apply the sampled state on the specified image.

    Args:
        boxes: the boxes to transform.
        state: the sampled state.
    """
    if state["horizontal"]:
        boxes["boxes"] = torch.stack(
            [
                boxes["width"] - boxes["boxes"][:, 2],
                boxes["boxes"][:, 1],
                boxes["width"] - boxes["boxes"][:, 0],
                boxes["boxes"][:, 3],
            ],
            dim=1,
        )
    if state["vertical"]:
        boxes["boxes"] = torch.stack(
            [
                boxes["boxes"][:, 0],
                boxes["height"] - boxes["boxes"][:, 3],
                boxes["boxes"][:, 2],
                boxes["height"] - boxes["boxes"][:, 1],
            ],
            dim=1,
        )
    return boxes

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/flip.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
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