gistfile1.txt

import cv2
import torch
import torch.nn.functional as F
import scipy
import numpy as np


def theta_for_patch_center(img_shape, window_size, patch_center):
    theta = torch.FloatTensor([[
        [window_size[0] / img_shape[0], 0, 2 * patch_center[0] / (img_shape[0] - 1) - 1],
        [0, window_size[1] / img_shape[1], 2 * patch_center[1] / (img_shape[1] - 1) - 1],
    ]])
    return theta


def main():
    canvas1 = np.zeros((5, 5), dtype=np.uint8)
    canvas1[0, 0] = 128
    canvas1[2, 2] = 255
    canvas1[2, 3] = 255
    canvas1[3, 1] = 255    

    window_size = (3, 5)

    theta = theta_for_patch_center(canvas1.shape, window_size, (2, 1))

    canvas1_torch = torch.FloatTensor(canvas1.astype(np.float32))
    canvas1_torch = canvas1_torch.unsqueeze(0).unsqueeze(0)
    grid = F.affine_grid(theta, [1, 1, window_size[0], window_size[1]])
    sampled = F.grid_sample(canvas1_torch, grid, mode="nearest")
    
    print(canvas1_torch)
    print(grid)
    print(sampled)


if __name__ == '__main__':
    main()