(
base: Float[Tensor, "channel base_height base_width"],
overlay: Float[Tensor, "channel overlay_height overlay_width"],
main_axis: Axis,
main_axis_alignment: Alignment,
cross_axis_alignment: Alignment,
)
| 83 | |
| 84 | |
| 85 | def overlay( |
| 86 | base: Float[Tensor, "channel base_height base_width"], |
| 87 | overlay: Float[Tensor, "channel overlay_height overlay_width"], |
| 88 | main_axis: Axis, |
| 89 | main_axis_alignment: Alignment, |
| 90 | cross_axis_alignment: Alignment, |
| 91 | ) -> Float[Tensor, "channel base_height base_width"]: |
| 92 | # The overlay must be smaller than the base. |
| 93 | _, base_height, base_width = base.shape |
| 94 | _, overlay_height, overlay_width = overlay.shape |
| 95 | assert base_height >= overlay_height and base_width >= overlay_width |
| 96 | |
| 97 | # Compute spacing on the main dimension. |
| 98 | main_dim = _get_main_dim(main_axis) |
| 99 | main_slice = _compute_offset( |
| 100 | base.shape[main_dim], overlay.shape[main_dim], main_axis_alignment |
| 101 | ) |
| 102 | |
| 103 | # Compute spacing on the cross dimension. |
| 104 | cross_dim = _get_cross_dim(main_axis) |
| 105 | cross_slice = _compute_offset( |
| 106 | base.shape[cross_dim], overlay.shape[cross_dim], cross_axis_alignment |
| 107 | ) |
| 108 | |
| 109 | # Combine the slices and paste the overlay onto the base accordingly. |
| 110 | selector = [..., None, None] |
| 111 | selector[main_dim] = main_slice |
| 112 | selector[cross_dim] = cross_slice |
| 113 | result = base.clone() |
| 114 | result[selector] = overlay |
| 115 | return result |
| 116 | |
| 117 | |
| 118 | def cat( |
no test coverage detected