Update app.py

app.py

@@ -28,6 +28,7 @@ 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 torchvision.transforms as T
 
 import tempfile
 from pathlib import Path
@@ -193,11 +194,21 @@ 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")
-    write_jpeg(frames[100], f"input_image.jpg")
+    #write_jpeg(frames[100], f"input_image.jpg")
     #res = warp_image(img1_batch, predicted_flow)
-    res = get_warp_res("input_image.jpg", "flofile.flo", fname_output='warped.png')
-    print(res)
-    return "done", "predicted_flow.jpg", ["flofile.flo"]
+    
+    # define a transform to convert a tensor to PIL image
+    transform = T.ToPILImage()
+    
+    # convert the tensor to PIL image using above transform
+    img = transform(frames[100])
+    
+    # display the PIL image
+    #img.show()
+    img.save('frame_input.jpg')
+    #res = get_warp_res("input_image.jpg", "flofile.flo", fname_output='warped.png')
+    #print(res)
+    return "done", "predicted_flow.jpg", ["flofile.flo"], 'frame_input.jpg'
 ####################################
 # Bonus: Creating GIFs of predicted flows
 # ---------------------------------------
@@ -226,4 +237,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()]).launch()
+gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(), gr.Files(), gr.Image()]).launch()