Update app.py

app.py

@@ -147,8 +147,8 @@ def infer():
     #img1_batch = torch.stack([frames[0]])
     #img2_batch = torch.stack([frames[1]])
 
-    img1_batch = torch.stack([input_frame_2])
-    img2_batch = torch.stack([input_diffused])
+    img1_batch = torch.stack([input_frame_1])
+    img2_batch = torch.stack([input_frame_2])
     
     print(f"FRAME AFTER stack: {img1_batch}")
     
@@ -240,7 +240,7 @@ def infer():
     # convert the tensor to PIL image using above transform
     #img = transform(frames[1])
     img = transform(input_frame_2)
-    img = img.resize((960, 520))
+    img = img.resize((520, 960))
     # display the PIL image
     #img.show()
     frame2pil = np.array(img.convert('RGB'))
@@ -248,6 +248,18 @@ def infer():
     print(f"frame1pil shape: {frame2pil.shape}")
     print(f"frame1pil dtype: {frame2pil.dtype}")
     img.save('raw_frame2.jpg')
+
+    # convert the tensor diffused to PIL image using above transform
+    #img = transform(frames[1])
+    img_diff = transform(input_diffused)
+    img_diff = img_diff.resize((520, 960))
+    # display the PIL image
+    #img.show()
+    diffpil = np.array(img_diff.convert('RGB'))
+    print(f"frame1pil: {diffpil}")
+    print(f"frame1pil shape: {diffpil.shape}")
+    print(f"frame1pil dtype: {diffpil.dtype}")
+    img_diff.save('diffused_resized.jpg')
     
 
     numpy_array_flow = predicted_flow.permute(1, 2, 0).detach().cpu().numpy()
@@ -264,7 +276,7 @@ def infer():
     numpy_array_flow*=1.
     # print('flow stats mul', flow.max(), flow.min(), flow.mean())
     # res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)
-    res = cv2.remap(frame2pil, numpy_array_flow, None, cv2.INTER_LANCZOS4)
+    res = cv2.remap(diffpil, numpy_array_flow, None, cv2.INTER_LANCZOS4)
     print(res)
     
     res = Image.fromarray(res)