<gh_stars>1-10
import torch
import torch.nn.functional as F


def normalize_coords(grid):
    """Normalize coordinates of image scale to [-1, 1]
    Args:
        grid: [B, 2, H, W]
    """
    assert grid.size(1) == 2
    h, w = grid.size()[2:]
    grid[:, 0, :, :] = 2 * (grid[:, 0, :, :].clone() / (w - 1)) - 1  # x: [-1, 1]
    grid[:, 1, :, :] = 2 * (grid[:, 1, :, :].clone() / (h - 1)) - 1  # y: [-1, 1]
    grid = grid.permute((0, 2, 3, 1))  # [B, H, W, 2]
    return grid


def meshgrid(img, homogeneous=False):
    """Generate meshgrid in image scale
    Args:
        img: [B, _, H, W]
        homogeneous: whether to return homogeneous coordinates
    Return:
        grid: [B, 2, H, W]
    """
    b, _, h, w = img.size()

    x_range = torch.arange(0, w).view(1, 1, w).expand(1, h, w).type_as(img)  # [1, H, W]
    y_range = torch.arange(0, h).view(1, h, 1).expand(1, h, w).type_as(img)

    grid = torch.cat((x_range, y_range), dim=0)  # [2, H, W], grid[:, i, j] = [j, i]
    grid = grid.unsqueeze(0).expand(b, 2, h, w)  # [B, 2, H, W]

    if homogeneous:
        ones = torch.ones_like(x_range).unsqueeze(0).expand(b, 1, h, w)  # [B, 1, H, W]
        grid = torch.cat((grid, ones), dim=1)  # [B, 3, H, W]
        assert grid.size(1) == 3
    return grid


def disp_warp(img, disp, padding_mode='border'):
    """Warping by disparity
    Args:
        img: [B, 3, H, W]
        disp: [B, 1, H, W], positive
        padding_mode: 'zeros' or 'border'
    Returns:
        warped_img: [B, 3, H, W]
        valid_mask: [B, 3, H, W]
    """
    # assert disp.min() >= 0

    grid = meshgrid(img)  # [B, 2, H, W] in image scale
    # Note that -disp here
    offset = torch.cat((-disp, torch.zeros_like(disp)), dim=1)  # [B, 2, H, W]
    sample_grid = grid + offset
    sample_grid = normalize_coords(sample_grid)  # [B, H, W, 2] in [-1, 1]
    warped_img = F.grid_sample(img, sample_grid, mode='bilinear', padding_mode=padding_mode)

    mask = torch.ones_like(img)
    valid_mask = F.grid_sample(mask, sample_grid, mode='bilinear', padding_mode='zeros')
    valid_mask[valid_mask < 0.9999] = 0
    valid_mask[valid_mask > 0] = 1
    return warped_img, valid_mask


def disp_warp_ms(x, disp):
    size = disp.size()
    if x.size()[-2:] != disp.size()[-2:]:   #resize disp to size of x
        # sf = disp.size()[-2:]/x.size()[-2:]
        # print("Warp Scale factor: ", sf)
        disp = disp.unsqueeze(1)
        disp = F.interpolate(disp, size=x.size()[-2:], mode='bilinear', align_corners=False)

        disp = disp.squeeze(1)

    bs, ch, h, w = x.size()
    bg, hg, wg = torch.meshgrid(torch.arange(0, bs), torch.arange(0, h), torch.arange(0, w))

    grid_b, grid_h, grid_w = bg.cuda(), hg.cuda(), wg.cuda()

    warped_gw = torch.sub(grid_w.float(), disp)
    grid = torch.stack([warped_gw, grid_h.float()], dim=-1)
    grid_normalized = ((grid * 2) / torch.Tensor([w, h]).cuda()) - 1
    output = F.grid_sample(x, grid_normalized, mode='bilinear', padding_mode='zeros')

    if x.size()[-2:] != disp.size()[-2:]:  # resize output back to original disp size
        output = F.interpolate(output, size=size[-2:], mode='bilinear', align_corners=False)

    return output


def pyramid_warp(img, disp):
    """Warping image features toowards disp disparity
    Args:
        img: feature pyramid
        disp: disparity pyramid
    Returns:
        warped_img: [B, 3, H, W]
        valid_mask: [B, 3, H, W]
    """
    warped_pyramid = []
    if len(disp) == 1:
        for im in img:
            warped_pyramid.append(disp_warp_ms(im, disp[0]))
    else:
        for p, im in zip(reversed(disp), img):
            warped_pyramid.append(disp_warp_ms(im, p))

    return warped_pyramid

import torch
import torch.nn.functional as F


def normalize_coords(grid):
    """Normalize coordinates of image scale to [-1, 1]
    Args:
        grid: [B, 2, H, W]
    """
    assert grid.size(1) == 2
    h, w = grid.size()[2:]
    grid[:, 0, :, :] = 2 * (grid[:, 0, :, :].clone() / (w - 1)) - 1  # x: [-1, 1]
    grid[:, 1, :, :] = 2 * (grid[:, 1, :, :].clone() / (h - 1)) - 1  # y: [-1, 1]
    grid = grid.permute((0, 2, 3, 1))  # [B, H, W, 2]
    return grid


def meshgrid(img, homogeneous=False):
    """Generate meshgrid in image scale
    Args:
        img: [B, _, H, W]
        homogeneous: whether to return homogeneous coordinates
    Return:
        grid: [B, 2, H, W]
    """
    b, _, h, w = img.size()

    x_range = torch.arange(0, w).view(1, 1, w).expand(1, h, w).type_as(img)  # [1, H, W]
    y_range = torch.arange(0, h).view(1, h, 1).expand(1, h, w).type_as(img)

    grid = torch.cat((x_range, y_range), dim=0)  # [2, H, W], grid[:, i, j] = [j, i]
    grid = grid.unsqueeze(0).expand(b, 2, h, w)  # [B, 2, H, W]

    if homogeneous:
        ones = torch.ones_like(x_range).unsqueeze(0).expand(b, 1, h, w)  # [B, 1, H, W]
        grid = torch.cat((grid, ones), dim=1)  # [B, 3, H, W]
        assert grid.size(1) == 3
    return grid


def disp_warp(img, disp, padding_mode='border'):
    """Warping by disparity
    Args:
        img: [B, 3, H, W]
        disp: [B, 1, H, W], positive
        padding_mode: 'zeros' or 'border'
    Returns:
        warped_img: [B, 3, H, W]
        valid_mask: [B, 3, H, W]
    """
    # assert disp.min() >= 0

    grid = meshgrid(img)  # [B, 2, H, W] in image scale
    # Note that -disp here
    offset = torch.cat((-disp, torch.zeros_like(disp)), dim=1)  # [B, 2, H, W]
    sample_grid = grid + offset
    sample_grid = normalize_coords(sample_grid)  # [B, H, W, 2] in [-1, 1]
    warped_img = F.grid_sample(img, sample_grid, mode='bilinear', padding_mode=padding_mode)

    mask = torch.ones_like(img)
    valid_mask = F.grid_sample(mask, sample_grid, mode='bilinear', padding_mode='zeros')
    valid_mask[valid_mask < 0.9999] = 0
    valid_mask[valid_mask > 0] = 1
    return warped_img, valid_mask


def disp_warp_ms(x, disp):
    size = disp.size()
    if x.size()[-2:] != disp.size()[-2:]:   #resize disp to size of x
        # sf = disp.size()[-2:]/x.size()[-2:]
        # print("Warp Scale factor: ", sf)
        disp = disp.unsqueeze(1)
        disp = F.interpolate(disp, size=x.size()[-2:], mode='bilinear', align_corners=False)

        disp = disp.squeeze(1)

    bs, ch, h, w = x.size()
    bg, hg, wg = torch.meshgrid(torch.arange(0, bs), torch.arange(0, h), torch.arange(0, w))

    grid_b, grid_h, grid_w = bg.cuda(), hg.cuda(), wg.cuda()

    warped_gw = torch.sub(grid_w.float(), disp)
    grid = torch.stack([warped_gw, grid_h.float()], dim=-1)
    grid_normalized = ((grid * 2) / torch.Tensor([w, h]).cuda()) - 1
    output = F.grid_sample(x, grid_normalized, mode='bilinear', padding_mode='zeros')

    if x.size()[-2:] != disp.size()[-2:]:  # resize output back to original disp size
        output = F.interpolate(output, size=size[-2:], mode='bilinear', align_corners=False)

    return output


def pyramid_warp(img, disp):
    """Warping image features toowards disp disparity
    Args:
        img: feature pyramid
        disp: disparity pyramid
    Returns:
        warped_img: [B, 3, H, W]
        valid_mask: [B, 3, H, W]
    """
    warped_pyramid = []
    if len(disp) == 1:
        for im in img:
            warped_pyramid.append(disp_warp_ms(im, disp[0]))
    else:
        for p, im in zip(reversed(disp), img):
            warped_pyramid.append(disp_warp_ms(im, p))

    return warped_pyramid
Normalize coordinates of image scale to [-1, 1]     Args:         grid: [B, 2, H, W] # x: [-1, 1] # y: [-1, 1] # [B, H, W, 2] Generate meshgrid in image scale     Args:         img: [B, _, H, W]         homogeneous: whether to return homogeneous coordinates     Return:         grid: [B, 2, H, W] # [1, H, W] # [2, H, W], grid[:, i, j] = [j, i] # [B, 2, H, W] # [B, 1, H, W] # [B, 3, H, W] Warping by disparity     Args:         img: [B, 3, H, W]         disp: [B, 1, H, W], positive         padding_mode: 'zeros' or 'border'     Returns:         warped_img: [B, 3, H, W]         valid_mask: [B, 3, H, W] # assert disp.min() >= 0 # [B, 2, H, W] in image scale # Note that -disp here # [B, 2, H, W] # [B, H, W, 2] in [-1, 1] #resize disp to size of x # sf = disp.size()[-2:]/x.size()[-2:] # print("Warp Scale factor: ", sf) # resize output back to original disp size Warping image features toowards disp disparity     Args:         img: feature pyramid         disp: disparity pyramid     Returns:         warped_img: [B, 3, H, W]         valid_mask: [B, 3, H, W]