Update app.py

app.py

@@ -30,6 +30,25 @@ import tempfile
 from pathlib import Path
 from urllib.request import urlretrieve
 
+def write_flo(filename, flow):
+    """
+    write optical flow in Middlebury .flo format
+    :param flow: optical flow map
+    :param filename: optical flow file path to be saved
+    :return: None
+    """
+    f = open(filename, 'wb')
+    magic = np.array([202021.25], dtype=np.float32)
+    (height, width) = flow.shape[0:2]
+    w = np.array([width], dtype=np.int32)
+    h = np.array([height], dtype=np.int32)
+    magic.tofile(f)
+    w.tofile(f)
+    h.tofile(f)
+    flow.tofile(f)
+    f.close() 
+
+
 def infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
     video_path = Path(tempfile.mkdtemp()) / "basketball.mp4"
@@ -122,8 +141,8 @@ def infer():
     flow_img = flow_to_image(predicted_flow).to("cpu")
 #     output_folder = "/tmp/"  # Update this to the folder of your choice
     write_jpeg(flow_img, f"predicted_flow.jpg")
-
-    return "done", "predicted_flow.jpg"
+    flo_file = write_flo("flofile.flo", predicted_flow)
+    return "done", "predicted_flow.jpg", flo_file
 ####################################
 # Bonus: Creating GIFs of predicted flows
 # ---------------------------------------
@@ -151,22 +170,5 @@ def infer():
 #
 # ffmpeg -f image2 -framerate 30 -i predicted_flow_%d.jpg -loop -1 flow.gif
 
-def write_flo(filename, flow):
-    """
-    write optical flow in Middlebury .flo format
-    :param flow: optical flow map
-    :param filename: optical flow file path to be saved
-    :return: None
-    """
-    f = open(filename, 'wb')
-    magic = np.array([202021.25], dtype=np.float32)
-    (height, width) = flow.shape[0:2]
-    w = np.array([width], dtype=np.int32)
-    h = np.array([height], dtype=np.int32)
-    magic.tofile(f)
-    w.tofile(f)
-    h.tofile(f)
-    flow.tofile(f)
-    f.close() 
 
-gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image()]).launch()
+gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(), gr.File()]).launch()