raft testing

RAFT_test.py

@@ -50,14 +50,27 @@ def track_raft(prev_frame, frame, model, device):
     curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
     prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)
     curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)
+    t0 = time.perf_counter()
     with torch.no_grad():
         flow = model(prev_tensor, curr_tensor)
         predicted_flows = flow[-1]
-
+    t1 = time.perf_counter()
+    print(f"RAFT time: {t1 - t0}")
     flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
     
     # Transpose to (H, W, 2) format expected by draw_arrows
     flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)
+
+
+    # torch.onnx.export(
+    #     model,                  # model to export
+    #     (prev_tensor, curr_tensor),        # inputs of the model,
+    #     "my_model.onnx",        # filename of the ONNX model
+    #     input_names=["prev_tensor", "curr_tensor"],  # Rename inputs for the ONNX model
+    #     dynamo=True             # True or False to select the exporter to use
+    # )
+    # print("ONNX model exported successfully")
+
     return flow_np
 
 def main_simple():
@@ -66,7 +79,7 @@ def main_simple():
     model = setup_simple_raft()
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Using device: {device}")
-    cap = cv2.VideoCapture(0)
+    cap = cv2.VideoCapture("s1.mp4")
     cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)