satflow/models/utils.py
def reverse_space_to_depth(
    frames: np.ndarray, temporal_block_size: int = 1, spatial_block_size: int = 1
) -> np.ndarray:
    """Reverse space to depth transform."""
    if len(frames.shape) == 4:
        return einops.rearrange(
            frames,
            "b h w (dh dw c) -> b (h dh) (w dw) c",
            dh=spatial_block_size,
            dw=spatial_block_size,
        )
    if len(frames.shape) == 5:
        return einops.rearrange(
            frames,
            "b t h w (dt dh dw c) -> b (t dt) (h dh) (w dw) c",
            dt=temporal_block_size,
            dh=spatial_block_size,
            dw=spatial_block_size,
        )
    raise ValueError(
        "Frames should be of rank 4 (batch, height, width, channels)"
        " or rank 5 (batch, time, height, width, channels)"
    )
Reverse space to depth transform.
Reverse space to depth transform.
Reverse space to depth transform .
def reverse_space_to_depth(
    frames: np.ndarray, temporal_block_size: int = 1, spatial_block_size: int = 1
) -> np.ndarray:
    if len(frames.shape) == 4:
        return einops.rearrange(
            frames,
            "b h w (dh dw c) -> b (h dh) (w dw) c",
            dh=spatial_block_size,
            dw=spatial_block_size,
        )
    if len(frames.shape) == 5:
        return einops.rearrange(
            frames,
            "b t h w (dt dh dw c) -> b (t dt) (h dh) (w dw) c",
            dt=temporal_block_size,
            dh=spatial_block_size,
            dw=spatial_block_size,
        )
    raise ValueError(
        "Frames should be of rank 4 (batch, height, width, channels)"
        " or rank 5 (batch, time, height, width, channels)"
    )
def reverse_space_to_depth ( frames : np . ndarray , temporal_block_size : int = 1 , spatial_block_size : int = 1 ) -> np . ndarray : if len ( frames . shape ) == 4 : return einops . rearrange ( frames , "b h w (dh dw c) -> b (h dh) (w dw) c" , dh = spatial_block_size , dw = spatial_block_size , ) if len ( frames . shape ) == 5 : return einops . rearrange ( frames , "b t h w (dt dh dw c) -> b (t dt) (h dh) (w dw) c" , dt = temporal_block_size , dh = spatial_block_size , dw = spatial_block_size , ) raise ValueError ( "Frames should be of rank 4 (batch, height, width, channels)" " or rank 5 (batch, time, height, width, channels)" )
Reverse space to depth transform.
Reverse space to depth transform .
"""Reverse space to depth transform."""
frames
np.ndarray
temporal_block_size
int
spatial_block_size
int