Update app.py

app.py

@@ -24,8 +24,7 @@ import torchvision.transforms.functional as F
 from torchvision.io import read_video
 from torchvision.models.optical_flow import Raft_Large_Weights
 from torchvision.models.optical_flow import raft_large
-
-
+from torchvision.io import write_jpeg
 
 import tempfile
 from pathlib import Path
@@ -115,12 +114,16 @@ def infer():
 
     from torchvision.utils import flow_to_image
 
-    flow_imgs = flow_to_image(predicted_flows)
+    #flow_imgs = flow_to_image(predicted_flows)
 
-    print(flow_imgs)
+    #print(flow_imgs)
 
-    return "done"
+    predicted_flow = list_of_flows[-1][0]
+    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_{i}.jpg")
 
+    return "done", "predicted_flow_{i}.jpg"
 ####################################
 # Bonus: Creating GIFs of predicted flows
 # ---------------------------------------
@@ -166,4 +169,4 @@ def write_flo(filename, flow):
     flow.tofile(f)
     f.close() 
 
-gr.Interface(fn=infer, inputs=[], outputs=gr.Textbox()).launch()
+gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image()]).launch()