Update app.py

app.py

@@ -58,95 +58,43 @@ def write_flo(flow, filename):
     flow.tofile(f)
     f.close()
 
-def get_pixel_value(img, x, y):
+#warp using scipy
+def warp_image(im, flow):
     """
-    Utility function to get pixel value for coordinate
-    vectors x and y from a  4D tensor image.
-    Input
-    -----
-    - img: tensor of shape (B, H, W, C)
-    - x: flattened tensor of shape (B*H*W, )
-    - y: flattened tensor of shape (B*H*W, )
-    Returns
-    -------
-    - output: tensor of shape (B, H, W, C)
+    Use optical flow to warp image to the next
+    :param im: image to warp
+    :param flow: optical flow
+    :return: warped image
     """
-    shape = tf.shape(x)
-    batch_size = shape[0]
-    height = shape[1]
-    width = shape[2]
-
-    batch_idx = tf.range(0, batch_size)
-    batch_idx = tf.reshape(batch_idx, (batch_size, 1, 1))
-    b = tf.tile(batch_idx, (1, height, width))
-
-    indices = tf.stack([b, y, x], 3)
-
-    return tf.gather_nd(img, indices)
-
-def tf_warp(img, flow, H, W):
-#    H = 256
-#    W = 256
-    x,y = tf.meshgrid(tf.range(W), tf.range(H))
-    x = tf.expand_dims(x,0)
-    x = tf.expand_dims(x,0)
-
-    y  =tf.expand_dims(y,0)
-    y = tf.expand_dims(y,0)
-
-    x = tf.cast(x, tf.float32)
-    y = tf.cast(y, tf.float32)
-    grid  = tf.concat([x,y],axis = 1)
-#    print grid.shape
-    flows = grid+flow
-    print(flows.shape)
-    max_y = tf.cast(H - 1, tf.int32)
-    max_x = tf.cast(W - 1, tf.int32)
-    zero = tf.zeros([], dtype=tf.int32)
-
-    x = flows[:,0,:,:]
-    y = flows[:,1,:,:]
-    x0 = x
-    y0 = y
-    x0 = tf.cast(x0, tf.int32)
-    x1 = x0 + 1
-    y0 = tf.cast(y0,  tf.int32)
-    y1 = y0 + 1
-
-    # clip to range [0, H/W] to not violate img boundaries
-    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)
-
-    # get pixel value at corner coords
-    Ia = get_pixel_value(img, x0, y0)
-    Ib = get_pixel_value(img, x0, y1)
-    Ic = get_pixel_value(img, x1, y0)
-    Id = get_pixel_value(img, x1, y1)
-
-    # recast as float for delta calculation
-    x0 = tf.cast(x0, tf.float32)
-    x1 = tf.cast(x1, tf.float32)
-    y0 = tf.cast(y0, tf.float32)
-    y1 = tf.cast(y1, tf.float32)
-
-
-    # calculate deltas
-    wa = (x1-x) * (y1-y)
-    wb = (x1-x) * (y-y0)
-    wc = (x-x0) * (y1-y)
-    wd = (x-x0) * (y-y0)
-
-    # add dimension for addition
-    wa = tf.expand_dims(wa, axis=3)
-    wb = tf.expand_dims(wb, axis=3)
-    wc = tf.expand_dims(wc, axis=3)
-    wd = tf.expand_dims(wd, axis=3)
-
-    # compute output
-    out = tf.add_n([wa*Ia, wb*Ib, wc*Ic, wd*Id])
-    return out
+    from scipy import interpolate
+    image_height = im.shape[0]
+    image_width = im.shape[1]
+    flow_height = flow.shape[0]
+    flow_width = flow.shape[1]
+    n = image_height * image_width
+    (iy, ix) = np.mgrid[0:image_height, 0:image_width]
+    (fy, fx) = np.mgrid[0:flow_height, 0:flow_width]
+    fx = fx.astype(np.float64)
+    fy = fy.astype(np.float64)
+    fx += flow[:,:,0]
+    fy += flow[:,:,1]
+    mask = np.logical_or(fx <0 , fx > flow_width)
+    mask = np.logical_or(mask, fy < 0)
+    mask = np.logical_or(mask, fy > flow_height)
+    fx = np.minimum(np.maximum(fx, 0), flow_width)
+    fy = np.minimum(np.maximum(fy, 0), flow_height)
+    points = np.concatenate((ix.reshape(n,1), iy.reshape(n,1)), axis=1)
+    xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n,1)), axis=1)
+    warp = np.zeros((image_height, image_width, im.shape[2]))
+    for i in range(im.shape[2]):
+        channel = im[:, :, i]
+        values = channel.reshape(n, 1)
+        new_channel = interpolate.griddata(points, values, xi, method='cubic')
+        new_channel = np.reshape(new_channel, [flow_height, flow_width])
+        new_channel[mask] = 1
+        warp[:, :, i] = new_channel.astype(np.uint8)
+
+    return warp.astype(np.uint8)
 
 def infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
@@ -242,18 +190,9 @@ def infer():
     write_jpeg(flow_img, f"predicted_flow.jpg")
     flo_file = write_flo(predicted_flow, "flofile.flo")
     
-    with tf.Session() as sess:
-        a = tf.placeholder(tf.float32, shape = [None,None,None,3])
-        flow_vec = tf.placeholder(tf.float32, shape = [None, 2, None, None])
-        init = tf.global_variables_initializer()
-        sess.run(init)
-
-        output = tf_warp(a, predicted_flow, 520, 960)
-        out = sess.run(output, feed_dict = {a:img, flow_vec:flow})
-        out = np.clip(out,0,255).astype('uint8')
-#        print out.shape
-        im = Image.fromarray(out[0].astype('uint8'))
-        im.save('output.jpg')    
+    res = warp_image(img1_batch, predicted_flow)
+    im = Image.fromarray(res)
+    im.save('output.jpg')    
     return "done", "predicted_flow.jpg", ["flofile.flo"]
 ####################################
 # Bonus: Creating GIFs of predicted flows