Update app.py

app.py

@@ -58,20 +58,95 @@ def write_flo(flow, filename):
     flow.tofile(f)
     f.close()
 
-def get_flow(filename):
-    with open(filename, 'rb') as f:
-        magic = np.fromfile(f, np.float32, count=1)
-        if 202021.25 != magic:
-            print('Magic number incorrect. Invalid .flo file')
-        else:
-            w = np.fromfile(f, np.int32, count=1)
-            h = np.fromfile(f, np.int32, count=1)
-            print('Reading %d x %d flo file' % (w, h))
-            data = np.fromfile(f, np.float32, count=2*w*h)
-            # Reshape data into 3D array (columns, rows, bands)
-            data2D = np.resize(data, (1, h[0], w[0],2))
-            data2D = np.transpose(data2D,[0, 3,1,2])
-            return data2D
+def get_pixel_value(img, x, y):
+    """
+    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)
+    """
+    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
 
 def infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
@@ -166,9 +241,19 @@ def infer():
 #     output_folder = "/tmp/"  # Update this to the folder of your choice
     write_jpeg(flow_img, f"predicted_flow.jpg")
     flo_file = write_flo(predicted_flow, "flofile.flo")
-                        
-    test = get_flow("flofile.flo")
-    print(test)
+    
+    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')    
     return "done", "predicted_flow.jpg", ["flofile.flo"]
 ####################################
 # Bonus: Creating GIFs of predicted flows