Update app.py

app.py

@@ -33,6 +33,10 @@ import tempfile
 from pathlib import Path
 from urllib.request import urlretrieve
 import tensorflow as tf
+from scipy.interpolate import interp2d
+from imageio import imread, imwrite
+from flowio import readFlowFile
+
 def write_flo(flow, filename):
     """
     Write optical flow in Middlebury .flo format
@@ -58,43 +62,57 @@ def write_flo(flow, filename):
     flow.tofile(f)
     f.close()
 
-#warp using scipy
-def warp_image(im, flow):
-    im = [(img1 + 1) / 2 for img1 in im]
-    
-    print(f"IMAGE: {im}")
-    #im = im[0].cpu().data.numpy()
-    #im = im.astype(np.float32)
-    im = torch.FloatTensor(im).permute(0, 3, 1, 2).contiguous()
-    
+def warpImage(im, vx, vy, cast_uint8=True):
+    '''
+    function to warp images with different dimensions
+    '''
+
+    height2, width2, nChannels = im.shape
+    height1, width1 = vx.shape
+
+    x = np.linspace(1, width2, width2)
+    y = np.linspace(1, height2, height2)
+    X = np.linspace(1, width1, width1)
+    Y = np.linspace(1, height1, height1)
+    xx, yy = np.meshgrid(x, y)
+    XX, YY = np.meshgrid(X, Y)
+    XX = XX + vx
+    YY = YY + vy
+    mask = (XX < 1) | (XX > width2) | (YY < 1) | (YY > height2)
+    XX = np.clip(XX, 1, width2)
+    YY = np.clip(XX, 1, height2)
+
+    warpI2 = np.zeros((height1, width1, nChannels))
+    for i in range(nChannels):
+        f = interp2d(x, y, im[:, :, i], 'cubic')
+        foo = f(X, Y)
+        foo[mask] = 0.6
+        warpI2[:, :, i] = foo
+
+    mask = 1 - mask
+
+    if cast_uint8:
+        warpI2 = warpI2.astype(np.uint8)
+
+    return warpI2, mask
+
+
+def get_warp_res(fname_image, fname_flow, fname_output='warped.png'):
+    im2 = imread(fname_image)
+    flow = readFlowFile(fname_flow)
+    im_warped, _ = warpImage(im2, flow[:, :, 0], flow[:, :, 1])
+    imwrite(fname_output, im_warped)
     
-    print(f"IMAGE AP: {im}")
-    print(f"FLOW AV: {flow}")
-    #flow = flow.cpu().data.numpy()
-    #flow = flow.astype(np.float32)
-    flow = torch.FloatTensor(flow)
-    
-    print(f"FLOW AP: {flow}")
-    h, w = flow.shape[:2]
-    flow[:,:,0] += np.arange(w)
-    flow[:,:,1] += np.arange(h)[:,np.newaxis]
-    prevImg = cv2.remap(im, flow, None, cv2.INTER_LINEAR)
-
 def infer():
     video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4"
     video_path = Path(tempfile.mkdtemp()) / "basketball.mp4"
     _ = urlretrieve(video_url, video_path)
 
-
-
-    
     frames, _, _ = read_video(str(video_path), output_format="TCHW")
 
     img1_batch = torch.stack([frames[100]])
     img2_batch = torch.stack([frames[101]])
 
-    
-
     weights = Raft_Large_Weights.DEFAULT
     transforms = weights.transforms()
 
@@ -173,10 +191,10 @@ 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")
-    
-    res = warp_image(img1_batch, predicted_flow)
-    im = Image.fromarray(res)
-    im.save('output.jpg')    
+    write_jpeg(img1_batch, 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"]
 ####################################
 # Bonus: Creating GIFs of predicted flows