Update app.py

app.py

@@ -1,3 +1,5 @@
+import gradio as gr
+
 """
 =====================================================
 Optical Flow: Predicting movement with the RAFT model
@@ -74,10 +76,8 @@ _ = urlretrieve(video_url, video_path)
 from torchvision.io import read_video
 frames, _, _ = read_video(str(video_path), output_format="TCHW")
 
-img1_batch = torch.stack([frames[100], frames[150]])
-img2_batch = torch.stack([frames[101], frames[151]])
-
-plot(img1_batch)
+img1= [frames[100]
+img2 = [frames[101]
 
 #########################
 # The RAFT model accepts RGB images. We first get the frames from
@@ -92,15 +92,15 @@ weights = Raft_Large_Weights.DEFAULT
 transforms = weights.transforms()
 
 
-def preprocess(img1_batch, img2_batch):
-    img1_batch = F.resize(img1_batch, size=[520, 960])
-    img2_batch = F.resize(img2_batch, size=[520, 960])
-    return transforms(img1_batch, img2_batch)
+def preprocess(img, img2):
+    img1 = F.resize(img1, size=[520, 960])
+    img2 = F.resize(img2, size=[520, 960])
+    return transforms(img1, img2)
 
 
-img1_batch, img2_batch = preprocess(img1_batch, img2_batch)
+img1, img2 = preprocess(img1, img2)
 
-print(f"shape = {img1_batch.shape}, dtype = {img1_batch.dtype}")
+print(f"shape = {img1.shape}, dtype = {img1.dtype}")
 
 
 ####################################
@@ -120,7 +120,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).to(device)
 model = model.eval()
 
-list_of_flows = model(img1_batch.to(device), img2_batch.to(device))
+list_of_flows = model(img1.to(device), img2.to(device))
 print(f"type = {type(list_of_flows)}")
 print(f"length = {len(list_of_flows)} = number of iterations of the model")
 
@@ -160,9 +160,9 @@ from torchvision.utils import flow_to_image
 flow_imgs = flow_to_image(predicted_flows)
 
 # The images have been mapped into [-1, 1] but for plotting we want them in [0, 1]
-img1_batch = [(img1 + 1) / 2 for img1 in img1_batch]
+img1 = [(img1 + 1) / 2 for img1 in img1]
 
-grid = [[img1, flow_img] for (img1, flow_img) in zip(img1_batch, flow_imgs)]
+grid = [[img1, flow_img] for (img1, flow_img) in zip(img1, flow_imgs)]
 plot(grid)
 
 ####################################