Update app.py

app.py

@@ -80,7 +80,7 @@ def warpImage(im, vx, vy, cast_uint8=True):
     #XX = XX + vx
     XX = np.concatenate([XX, vx], axis = 1)
     #YY = YY + vy
-    YY = np.concatenate([YY, vy], axis = 0)
+    YY = np.concatenate([YY, vy], axis = 1)
     mask = (XX < 1) | (XX > width2) | (YY < 1) | (YY > height2)
     XX = np.clip(XX, 1, width2)
     YY = np.clip(XX, 1, height2)
@@ -195,21 +195,23 @@ 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")
+    input_image = flow_to_image(img1_batch).to("cpu")
+    write_jpeg(input_image, f"frame_input.jpg")
     flo_file = write_flo(predicted_flow, "flofile.flo")
     #write_jpeg(frames[100], f"input_image.jpg")
     #res = warp_image(img1_batch, predicted_flow)
     
     # define a transform to convert a tensor to PIL image
-    transform = T.ToPILImage()
+    #transform = T.ToPILImage()
     
     # convert the tensor to PIL image using above transform
-    img = transform(frames[100])
+    #img = transform(frames[100])
     
     # display the PIL image
     #img.show()
-    img.save('frame_input.jpg')
-    res = get_warp_res("frame_input.jpg", predicted_flow, fname_output='warped.png')
-    print(res)
+    #img.save('frame_input.jpg')
+    #res = get_warp_res("frame_input.jpg", predicted_flow, fname_output='warped.png')
+    #print(res)
     return "done", "predicted_flow.jpg", ["flofile.flo"], 'frame_input.jpg'
 ####################################
 # Bonus: Creating GIFs of predicted flows