python將光流apply/warp（映射）到圖像上

實驗目的

在圖像上選一個點，對這一點進行膨脹，使原圖空出一片區域，然後填充其他東西（這裏爲了演示，填充了一個黑圓，半徑大小爲100個pixel），要求是保證填充後圖像儘量自然，且保證原有的邊界連續

實驗結果（左邊爲原圖，右邊爲填充後的結果）：

代碼

warp（基於TensorFlow）

image_warp.py

import tensorflow as tf


def image_warp(im, flow):
    """Performs a backward warp of an image using the predicted flow.

    Args:
        im: Batch of images. [num_batch, height, width, channels]
        flow: Batch of flow vectors. [num_batch, height, width, 2]
    Returns:
        warped: transformed image of the same shape as the input image.
    """
    with tf.variable_scope('image_warp'):

        num_batch, height, width, channels = im.shape.as_list()
        max_x = tf.cast(width - 1, 'int32')

        max_y = tf.cast(height - 1, 'int32')
        zero = tf.zeros([], dtype='int32')

        # We have to flatten our tensors to vectorize the interpolation
        im_flat = tf.reshape(im, [-1, channels])
        flow_flat = tf.reshape(flow, [-1, 2])

        # Floor the flow, as the final indices are integers
        # The fractional part is used to control the bilinear interpolation.
        flow_floor = tf.to_int32(tf.floor(flow_flat))
        bilinear_weights = flow_flat - tf.floor(flow_flat)

        # Construct base indices which are displaced with the flow
        pos_x = tf.tile(tf.range(width), [height * num_batch])
        grid_y = tf.tile(tf.expand_dims(tf.range(height), 1), [1, width])
        pos_y = tf.tile(tf.reshape(grid_y, [-1]), [num_batch])

        x = flow_floor[:, 0]
        y = flow_floor[:, 1]
        xw = bilinear_weights[:, 0]
        yw = bilinear_weights[:, 1]

        # Compute interpolation weights for 4 adjacent pixels
        # expand to num_batch * height * width x 1 for broadcasting in add_n below
        wa = tf.expand_dims((1 - xw) * (1 - yw), 1) # top left pixel
        wb = tf.expand_dims((1 - xw) * yw, 1) # bottom left pixel
        wc = tf.expand_dims(xw * (1 - yw), 1) # top right pixel
        wd = tf.expand_dims(xw * yw, 1) # bottom right pixel

        x0 = pos_x + x
        x1 = x0 + 1
        y0 = pos_y + y
        y1 = y0 + 1

        x0 = tf.clip_by_value(x0, zero, max_x)
        x1 = tf.clip_by_value(x1, zero, max_x)
        y0 = tf.clip_by_value(y0, zero, max_y)
        y1 = tf.clip_by_value(y1, zero, max_y)

        dim1 = width * height
        batch_offsets = tf.range(num_batch) * dim1
        base_grid = tf.tile(tf.expand_dims(batch_offsets, 1), [1, dim1])
        base = tf.reshape(base_grid, [-1])

        base_y0 = base + y0 * width
        base_y1 = base + y1 * width
        idx_a = base_y0 + x0
        idx_b = base_y1 + x0
        idx_c = base_y0 + x1
        idx_d = base_y1 + x1

        Ia = tf.gather(im_flat, idx_a)
        Ib = tf.gather(im_flat, idx_b)
        Ic = tf.gather(im_flat, idx_c)
        Id = tf.gather(im_flat, idx_d)

        warped_flat = tf.add_n([wa * Ia, wb * Ib, wc * Ic, wd * Id])
        warped = tf.reshape(warped_flat, [num_batch, height, width, channels])

        return warped

主代碼

import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"
import tensorflow as tf 
import numpy as np 
import cv2
from PIL import Image
import scipy.sparse as ss
from image_warp import image_warp

def create_position_index(height, width):
    position_indexes = np.rollaxis(np.indices(dimensions=(width, height)), 0, 3).transpose((1,0,2))
    position_indexes = np.reshape(position_indexes, (-1, 2))
    return position_indexes

def distance(pos, center):
    center = np.asarray(center)
    dif = pos - center
    dis = np.sqrt(np.square(dif[:,0]) + np.square(dif[:,1]))
    return dis, dif

def gen_angle(size, center):
    h, w, _ = size
    pos = create_position_index(h, w)
    dis, dif = distance(pos, center)
    sin = dif[:,1] / dis
    cos = dif[:,0] / dis
    x_list = pos[:,0]
    y_list = pos[:,1]
    sin_matrix = ss.coo_matrix((sin,(x_list,y_list)),shape=(h, w))
    sin_matrix = sin_matrix.toarray()
    cos_matrix = ss.coo_matrix((cos,(x_list,y_list)),shape=(h, w))
    cos_matrix = cos_matrix.toarray()
    sin_matrix = sin_matrix[:,:,np.newaxis]
    cos_matrix = cos_matrix[:,:,np.newaxis]
    tan = np.concatenate([sin_matrix, cos_matrix], axis=2)
    return tan

def gen_flow(size, center, radius):
    flow = np.ones((size[0], size[1], 2), dtype=np.float32)
    angle = gen_angle(flow.shape, center)
    flow = -flow * angle * radius
    return flow

def deform(img, flow):
    img_dims = (1, img.shape[0], img.shape[1], 1)
    flow_dims = (1, img.shape[0], img.shape[1], 2)
    image_moving = tf.placeholder(tf.float32, shape=img_dims)
    flow_moving = tf.placeholder(tf.float32, shape=flow_dims)
    image_warped = image_warp(image_moving, flow_moving)
    img = img[np.newaxis, :, :, np.newaxis]
    flow = flow[np.newaxis, :]
    sess_config = tf.ConfigProto(allow_soft_placement=True)
    with tf.Session(config=sess_config) as sess:
        warped_label = sess.run(image_warped, 
                        feed_dict={image_moving: img,
                                    flow_moving: flow})
    warped_label = np.squeeze(warped_label).astype(np.uint8)
    return warped_label

def draw_mask(deformed, center, radius):
    h, w = deformed.shape
    pos = create_position_index(h, w)
    dis, _ = distance(pos, center)
    dis[dis <= radius] = 0
    dis[dis > radius] = 1
    x_list = pos[:,0]
    y_list = pos[:,1]
    mask_matrix = ss.coo_matrix((dis,(x_list,y_list)),shape=(h, w))
    mask_matrix = mask_matrix.toarray()
    deformed = np.multiply(deformed, mask_matrix)
    deformed = deformed.astype(np.uint8)
    return deformed


img = np.asarray(Image.open('./0000.png'))
# img = img[:512, :512]

center_point = [600, 600]
radius = 100
flow = gen_flow(img.shape, center_point, radius)

deformed = deform(img.copy(), flow)

deformed = draw_mask(deformed, center_point, radius)

img_cat = np.concatenate([img, deformed], axis=1)
Image.fromarray(img_cat).save('./deformed.png')