Update app.py

app.py

@@ -19,6 +19,8 @@ predicted flows to RGB images for visualization.
 
 import cv2
 import numpy as np
+import os
+import sys
 import torch
 import matplotlib.pyplot as plt
 import torchvision.transforms.functional as F
@@ -56,52 +58,20 @@ def write_flo(flow, filename):
     flow.tofile(f)
     f.close()
 
-def warp_flow(img, flow):
-      #flow = flow.cpu().data.numpy()
-      h, w = flow.shape[:2]
-      flow = flow.copy()
-      flow[:, :, 0] += np.arange(w)
-      flow[:, :, 1] += np.arange(h)[:, np.newaxis]
-      res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)
-      return res
-
-def read_flo_file(filename):
-    """
-    Read from Middlebury .flo file
-    :param flow_file: name of the flow file
-    :return: optical flow data in matrix
-    """
-    f = open(filename, 'rb')
-    magic = np.fromfile(f, np.float32, count=1)
-    data2d = None
-
-    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 flow file in .flo format" % (h, w))
-        data2d = np.fromfile(f, np.float32, count=2 * w * h)
-        # reshape data into 3D array (columns, rows, channels)
-        data2d = np.resize(data2d, (h[0], w[0], 2))
-    f.close()
-    return data2d
-
-def warp(frame1, frame2, flo_path, blend=0.5, weights_path=None):
-    flo_path = flo_path.cpu().data.numpy()
-    flow21 = np.array(flo_path)
-    frame1pil = np.array(frame1)
-    frame1_warped21 = warp_flow(frame1pil, flow21)
-    # frame2pil = frame1pil
-    frame2pil = np.array(frame2)
-
-    if weights_path: 
-      forward_weights = load_cc(weights_path)
-      blended_w = frame2pil*(1-blend) + blend*(frame1_warped21*forward_weights+frame2pil*(1-forward_weights))
-    else: blended_w = frame2pil*(1-blend) + frame1_warped21*(blend)
-
-    return  PIL.Image.fromarray(blended_w.astype('uint8'))
-
+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 infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
@@ -197,8 +167,9 @@ def infer():
     write_jpeg(flow_img, f"predicted_flow.jpg")
     flo_file = write_flo(predicted_flow, "flofile.flo")
                         
-    warp_res = warp(img1_batch, img2_batch, predicted_flow, blend=0.5)
-    return "done", "predicted_flow.jpg", ["flofile.flo"], warp_res
+    test = get_flow("flofile.flo")
+    print(test)
+    return "done", "predicted_flow.jpg", ["flofile.flo"]
 ####################################
 # Bonus: Creating GIFs of predicted flows
 # ---------------------------------------
@@ -227,4 +198,4 @@ def infer():
 # ffmpeg -f image2 -framerate 30 -i predicted_flow_%d.jpg -loop -1 flow.gif
 
       
-gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(), gr.Files(), gr.Image(type="pil")]).launch()
+gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(), gr.Files()]).launch()