Update app.py

app.py

@@ -41,8 +41,8 @@ def infer():
     
     frames, _, _ = read_video(str(video_path), output_format="TCHW")
 
-    img1= frames[100]
-    img2 = frames[101]
+    img1_batch = torch.stack([frames[100]])
+    img2_batch = torch.stack([frames[101]])
 
     
 
@@ -50,15 +50,15 @@ def infer():
     transforms = weights.transforms()
 
 
-    def preprocess(img1, img2):
-        img1 = F.resize(img1, size=[520, 960])
-        img2 = F.resize(img2, size=[520, 960])
-        return transforms(img1, img2)
+    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)
 
 
-    img1, img2 = preprocess(img1, img2)
+    img1_batch, img2_batch = preprocess(img1_batch, img2_batch)
 
-    print(f"shape = {img1.shape}, dtype = {img1.dtype}")
+    print(f"shape = {img1_batch.shape}, dtype = {img1_batch.dtype}")
 
 
 ####################################
@@ -78,7 +78,7 @@ def infer():
     model = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).to(device)
     model = model.eval()
 
-    list_of_flows = model(img1.to(device), img2.to(device))
+    list_of_flows = model(img1_batch.to(device), img2_batch.to(device))
     print(f"type = {type(list_of_flows)}")
     print(f"length = {len(list_of_flows)} = number of iterations of the model")