Update app.py

app.py

@@ -55,6 +55,30 @@ def write_flo(flow, filename):
     flow.tofile(f)
     f.close()
 
+def warp_flow(img, flow):
+      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 warp(frame1, frame2, flo_path, blend=0.5, weights_path=None):
+    flow21 = np.load(flo_path)
+    frame1pil = np.array(frame1.convert('RGB').resize((flow21.shape[1],flow21.shape[0])))
+    frame1_warped21 = warp_flow(frame1pil, flow21)
+    # frame2pil = frame1pil
+    frame2pil = np.array(frame2.convert('RGB').resize((flow21.shape[1],flow21.shape[0])))
+
+    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 infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
@@ -149,7 +173,9 @@ 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")
-    return "done", "predicted_flow.jpg", ["flofile.flo"]
+                         
+    warp_res = warp(frames[100], frames[101], flo_file, blend=0.5, None)
+    return "done", "predicted_flow.jpg", ["flofile.flo"], warp_res
 ####################################
 # Bonus: Creating GIFs of predicted flows
 # ---------------------------------------
@@ -177,28 +203,5 @@ def infer():
 #
 # ffmpeg -f image2 -framerate 30 -i predicted_flow_%d.jpg -loop -1 flow.gif
 
-def warp_flow(img, flow):
-      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 warp(frame1, frame2, flo_path, blend=0.5, weights_path=None):
-    flow21 = np.load(flo_path)
-    frame1pil = np.array(frame1.convert('RGB').resize((flow21.shape[1],flow21.shape[0])))
-    frame1_warped21 = warp_flow(frame1pil, flow21)
-    # frame2pil = frame1pil
-    frame2pil = np.array(frame2.convert('RGB').resize((flow21.shape[1],flow21.shape[0])))
-
-    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'))
       
-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(type="pil")]).launch()