awarp.cpp

// SMAGL is 0 or 2
// PixelType is uint8_t
static void warp_c(const uint8_t *srcp8, const uint8_t *edgep8, uint8_t *dstp8, int src_stride, int edge_stride, int dst_stride, int width, int height, int depth, int bits_per_sample)
{
    const uint8_t *srcp = (const uint8_t *)srcp8;
    const uint8_t *edgep = (const uint8_t *)edgep8;
    uint8_t *dstp = (uint8_t *)dstp8;

    src_stride /= sizeof(uint8_t);
    edge_stride /= sizeof(uint8_t);
    dst_stride /= sizeof(uint8_t);

    int pixel_max = (1 << bits_per_sample) - 1;

    const int x_limit_min = 0;
    const int x_limit_max = (width - 1);

    float scale = (depth << 8) / 65536.0f;

    for (int y = 0; y < height; y++) {
        float y_limit_min = -y;
        float y_limit_max = (height - y - 1) - 1e-2f;

        for (int x = 0; x < width; x++) {
            // calculate displacement

            int above, below;
            if (y == 0)
                above = edgep[x];
            else
                above = edgep[-edge_stride + x];

            if (y == height - 1)
                below = edgep[x];
            else
                below = edgep[edge_stride + x];

            int left, right;
            if (x == 0)
                left = edgep[x];
            else
                left = edgep[x - 1];

            if (x == width - 1)
                right = edgep[x];
            else
                right = edgep[x + 1];

            float _h = (left - right) * scale;
            float _v = (above - below) * scale;

            _v = std::clamp(_v, y_limit_min, y_limit_max);

            float remainder_h = fmod(_h, 1.0);
            remainder_h = remainder_h < 0.0 ? 1.0 + remainder_h : remainder_h;

            float remainder_v = fmod(_v, 1.0);
            remainder_v = remainder_v < 0.0 ? 1.0 + remainder_v : remainder_v;

            int h = floor(_h) + x, v = floor(_v);

            remainder_h = (x_limit_max > h) && !(x_limit_min > h) ? remainder_h : 0;

            h = std::min(h, x_limit_max);
            h = std::max(h, x_limit_min);

            // h and v contain the displacement now.

            int s00 = srcp[v * src_stride + h];
            int s01 = srcp[v * src_stride + h + 1];
            int s10 = srcp[(v + 1) * src_stride + h];
            int s11 = srcp[(v + 1) * src_stride + h + 1];

            float s0 = s00 * (1.0 - remainder_h) + s01 * remainder_h + 0.5;
            float s1 = s10 * (1.0 - remainder_h) + s11 * remainder_h + 0.5;
            float s  =  s0 * (1.0 - remainder_v) +  s1 * remainder_v + 0.5;

            int val = nearbyint(s); // Use floor to match original
            dstp[x] = std::min(std::max(val, 0), pixel_max);
        }

        srcp += src_stride;
        edgep += edge_stride;
        dstp += dst_stride;
    }
}

import torch

def rgb_warp(src, mask):
    N, C, H, W = src.shape
    device = src.device

    # dst = src.clone()

    scale = (8 * 256) / 65536

    pixel_max = 1.0

    x_limit_min = 0
    x_limit_max = W - 2

    # for n in range(N):
    if True:

        # for c in range(C):
        if True:

            # for y in range(H):
            y = torch.arange(H).view(1, 1, -1, 1)
            if True:

                y_limit_min = -y
                y_limit_max = (H - 2) - y

                # for x in range(W):
                x = torch.arange(W).view(1, 1, 1, -1)
                if True:

                    # above = mask[n, c, 0, x] if y == 0 else mask[n, c, y - 1, x]
                    above = torch.cat([mask[:, :, :1, :], mask[:, :, :-1, :]], dim=2)

                    # below = mask[n, c, -1, x] if y == H - 1 else mask[n, c, y + 1, x]
                    below = torch.cat([mask[:, :, 1:, :], mask[:, :, -1:, :]], dim=2)

                    # left = mask[n, c, y, x] if x == 0 else mask[n, c, y, x - 1]
                    left = torch.cat([mask[:, :, :, :1], mask[:, :, :, :-1]], dim=3)

                    # right = mask[n, c, y, -1] if x == W - 1 else mask[n, c, y, x + 1]
                    right = torch.cat([mask[:, :, :, 1:], mask[:, :, :, -1:]], dim=3)

                    _h = (left - right) * scale
                    _v = (above - below) * scale

                    # _v = torch.clamp(_v, y_limit_min, y_limit_max)
                    _v = torch.where(_v < y_limit_min, y_limit_min, _v)
                    _v = torch.where(_v > y_limit_min, y_limit_max, _v)

                    remainder_h = _h % 1
                    remainder_v = _v % 1

                    h = (torch.floor(_h) + x)
                    v = torch.floor(_v)

                    # remainder_h = remainder_h if (x_limit_max > h) and not (x_limit_min > h) else torch.zeros(1)
                    remainder_h = torch.where((x_limit_max > h) & (x_limit_min <= h), remainder_h, torch.zeros(1))

                    # h = torch.clamp(h, x_limit_min, x_limit_max)
                    _h = torch.where(h < x_limit_min, x_limit_min, h)
                    _h = torch.where(_h > x_limit_max, x_limit_max, _h)

                    h = h.long()
                    v = v.long()

                    """
                    s00 = src[n, c, v, h]
                    s01 = src[n, c, v, h + 1]
                    s10 = src[n, c, v + 1, h]
                    s11 = src[n, c, v + 1, h + 1]
                    """
                    src = torch.flatten(src, start_dim=2) # shape: (N, C, H * W)
                    s00_indices = (v * W + h).flatten() # shape: (H * W)
                    s00 = torch.index_select(src, dim=2, index=s00_indices)
                    s01_indices = (v * W + h + 1).flatten() # shape: (H * W)
                    s01 = torch.index_select(src, dim=2, index=s01_indices)
                    s10_indices = ((v + 1) * W + h).flatten()[0, 0] # shape: (H * W)
                    s10 = torch.index_select(src, dim=2, index=s10_indices)
                    s11_indices = ((h + 1) * W + (v + 1)).flatten() # shape: (H * W)
                    s11 = torch.index_select(src, dim=2, index=s11_indices)

                    # s0 = s00 * (1 - remainder_h) + s01 * remainder_h + 0.5
                    s0 = torch.lerp(s00, s01, remainder_h) + 0.5

                    # s1 = s10 * (1 - remainder_h) + s11 * remainder_h + 0.5
                    s1 = torch.lerp(s10, s11, remainder_h) + 0.5

                    # s = s0 * (1 - remainder_v) + s1 * remainder_v + 0.5
                    s = torch.lerp(s0, s1, remainder_v) + 0.5

                    # s = torch.clamp(s, 0.0, 1.0).view(1)
                    s = torch.where(s < 0.0, 0.0, s)
                    s = torch.where(s > 1.0, 1.0, s)

                    #dst[n, c, y, x:x + 1] = s
                    dst = s.view(N, C, H, W)

    return dst

def awarp(src, mask):
    N,C,H,W = src.shape
    
    scale = (8 * 256) / 65536

    pixel_max = 1.0

    y = torch.arange(H).view(1, 1, -1, 1).float()
    x = torch.arange(W).view(1, 1, 1, -1).float()

    coord = torch.arange(H*W).view(1, 1, H, W)

    # x_limit_min = torch.zeros_like(x)
    # x_limit_max = torch.full_like(x, W) - 2

    x_limit_min = 0
    x_limit_max = W - 2

    y_limit_min = -y
    y_limit_max = (H - 2) - y

    above = torch.cat([mask[:, :,  :1, : ], mask[:, :,   :-1, :  ]], dim=2)
    below = torch.cat([mask[:, :, 1:,  : ], mask[:, :, -1:,   :  ]], dim=2)
    left  = torch.cat([mask[:, :,  :,  :1], mask[:, :,   :,   :-1]], dim=3)
    right = torch.cat([mask[:, :,  :, 1: ], mask[:, :,   :, -1:  ]], dim=3)

    _h = (left - right) * scale
    _v = (above - below) * scale

    _v = torch.where(_v < y_limit_min, y_limit_min, _v)
    _v = torch.where(_v > y_limit_min, y_limit_max, _v)

    remainder_h = _h % 1
    remainder_v = _v % 1

    h = torch.floor(_h) + x
    v = torch.floor(_v + y)

    remainder_h[(h >= x_limit_max) | (h < x_limit_min)] = 1e-4

    h = torch.clamp(h, x_limit_min, x_limit_max)

    h = h.long()
    v = v.long()

    """
    s00 = src[n, c, v, h]
    s01 = src[n, c, v, h + 1]
    s10 = src[n, c, v + 1, h]
    s11 = src[n, c, v + 1, h + 1]
    """
    src = torch.flatten(src, start_dim=2) # shape: (N, C, H * W)

    s00_indices = (v * W + h).flatten() # shape: (1, 1, H * W)
    s00 = torch.index_select(src, dim=2, index=s00_indices).reshape(-1,C,H,W)
    s01_indices = (v * W + h + 1).flatten() # shape: (1, 1, H * W)
    s01 = torch.index_select(src, dim=2, index=s01_indices).reshape(-1,C,H,W)
    s10_indices = ((v + 1) * W + h).flatten() # shape: (1, 1, H * W)
    s10 = torch.index_select(src, dim=2, index=s10_indices).reshape(-1,C,H,W)
    s11_indices = ((h + 1) * W + (v + 1)).flatten() # shape: (1, 1, H * W)
    s11 = torch.index_select(src, dim=2, index=s11_indices).reshape(-1,C,H,W)

    s0 = s00 * (1 - remainder_h) + s01 * remainder_h + 0.5
    s1 = s10 * (1 - remainder_h) + s11 * remainder_h + 0.5
    s  =  s0 * (1 - remainder_v) +  s1 * remainder_v + 0.5

    s = torch.clamp(s, 0.0, pixel_max)
    dst = s.view(N, C, H, W)
    return dst

Deprecated

https://gitlab.com/sr11/gist/blob/master/awarp/awarp.py

def awarp(src, mask):
    N,C,H,W = src.shape
    
    scale = (8 * 256) / 65536

    pixel_max = 1.0

    y = torch.arange(H).view(1, 1, -1, 1).float()
    x = torch.arange(W).view(1, 1, 1, -1).float()

    # coord = torch.arange(H*W).view(1, 1, H, W)

    # x_limit_min = torch.zeros_like(x)
    # x_limit_max = torch.full_like(x, W) - 2

    x_limit_min = 0
    x_limit_max = W - 2

    y_limit_min = -y
    y_limit_max = (H - 2) - y

    above = torch.cat([mask[:, :,  :1, : ], mask[:, :,   :-1, :  ]], dim=2)
    below = torch.cat([mask[:, :, 1:,  : ], mask[:, :, -1:,   :  ]], dim=2)
    left  = torch.cat([mask[:, :,  :,  :1], mask[:, :,   :,   :-1]], dim=3)
    right = torch.cat([mask[:, :,  :, 1: ], mask[:, :,   :, -1:  ]], dim=3)

    _h = (left - right) * scale
    _v = (above - below) * scale

    _v = torch.where(_v < y_limit_min, y_limit_min, _v)
    _v = torch.where(_v > y_limit_min, y_limit_max, _v)

    remainder_h = _h % 1
    remainder_v = _v % 1

    h = torch.floor(_h) + x
    v = torch.floor(_v + y)

    remainder_h[(h >= x_limit_max) | (h < x_limit_min)] = 1e-4

    h = torch.clamp(h, x_limit_min, x_limit_max)

    h = h.long()
    v = v.long()

    """
    s00 = src[n, c, v, h]
    s01 = src[n, c, v, h + 1]
    s10 = src[n, c, v + 1, h]
    s11 = src[n, c, v + 1, h + 1]
    """
    src = torch.flatten(src, start_dim=2) # shape: (N, C, H * W)

    # please also try replacing ".expand()" with ".repeat()"
    s00_indices = (v * W + h).flatten(start_dim=2).expand_as(src) # shape: (N, C, H * W)
    s00 = torch.gather(src, dim=2, index=s00_indices).reshape(-1,C,H,W)
    s01_indices = (v * W + h + 1).flatten(start_dim=2).expand_as(src) # shape: (N, C, H * W)
    s01 = torch.gather(src, dim=2, index=s01_indices).reshape(-1,C,H,W)
    s10_indices = ((v + 1) * W + h).flatten(start_dim=2).expand_as(src) # shape: (N, C, H * W)
    s10 = torch.gather(src, dim=2, index=s10_indices).reshape(-1,C,H,W)
    s11_indices = ((h + 1) * W + (v + 1)).flatten(start_dim=2).expand(-1, C, -1) # shape: (N, C, H * W)
    s11 = torch.gather(src, dim=2, index=s11_indices).expand_as(src)

    s0 = torch.lerp(s00, s01, remainder_h) # + 0.5
    s1 = torch.lerp(s10, s11, remainder_h) # + 0.5
    s = torch.lerp(s0, s1, remainder_v) # + 0.5

    s = torch.clamp(s, 0.0, pixel_max)
    dst = s.view(N, C, H, W)
    return dst