raft testing

.gitignore

@@ -10,4 +10,5 @@ orisys.egg-info/
 /old_orisys/
 /dist_nuitka/
 /build_nuitka
-/data/
+/data/
+/models/

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)
     

config/prototype_of3.json

@@ -2,11 +2,11 @@
   "name": "prototype",
   "version": "1.0",
   "camera": {
-    "resolution_raw": [1920, 1080],
+    "resolution_raw": [960, 540],
     "fps": 30,
     "units": 15,
     "units_used": 9, 
-    "resolution_stitch": [400, 400]
+    "resolution_stitch": [200, 200]
   },
   "stitcher": {
     "min_area_frac": 0.01,

src/orisys/tracker.py

@@ -2,7 +2,8 @@ import cv2
 import numpy as np
 import torch
 import torchvision.transforms as T
-from torchvision.models.optical_flow import raft_large
+from torchvision.models.optical_flow import raft_large, raft_small
+import time
 
 class Tracker:
     def __init__(self,
@@ -54,6 +55,7 @@ class Tracker:
         self.__upper_threshold = config_tracker["upper_threshold"]  # error greater than threshold will add new reference frame
         ### RAFT
         self.__raft_device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"RAFT device: {self.__raft_device}")
         self.__raft_model = self.setup_simple_raft()
         self.__raft_ref_frame = None
         self.__raft_ref_tensor = None
@@ -250,7 +252,7 @@ class Tracker:
         model = model.eval()
         
         # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs)
-        torch.set_float32_matmul_precision('medium')
+        # torch.set_float32_matmul_precision('medium')
         
         # Enable CUDA-specific optimizations
         if self.__raft_device == 'cuda':
@@ -263,7 +265,7 @@ class Tracker:
             
         print(f"RAFT model loaded on {self.__raft_device}")
         return model
-
+    
     def _preprocess_frame_to_tensor(self, frame, resize=None):
         """Efficiently convert numpy frame to RAFT-ready tensor"""
         # Ensure contiguous memory for faster tensor conversion
@@ -296,7 +298,6 @@ class Tracker:
         # Add batch dimension
         tensor = tensor.unsqueeze(0)
         
-        # Keep in float32 - autocast will handle mixed precision during inference
         return tensor
     
     def track_raft(self, frame):
@@ -319,16 +320,22 @@ class Tracker:
         prev_tensor = self.__raft_ref_tensor
         curr_tensor = self._preprocess_frame_to_tensor(frame)
         
-        # OPTIMIZATION: Remove CUDA stream since immediate sync eliminates async benefit
-        # The bottleneck is inference, not data transfer (already optimized with pin_memory)
-        with torch.no_grad():
+        # Synchronize before timing to ensure preprocessing is complete
+        if self.__raft_device == 'cuda':
+            torch.cuda.synchronize()
+        
+        with torch.inference_mode():
             if self.__raft_device == 'cuda':
                 with torch.amp.autocast('cuda'):
                     flow = self.__raft_model(prev_tensor, curr_tensor)
             else:
                 flow = self.__raft_model(prev_tensor, curr_tensor)
             predicted_flows = flow[-1]
-
+        
+        # Synchronize after inference to get accurate timing
+        if self.__raft_device == 'cuda':
+            torch.cuda.synchronize()
+        
         # Optimize tensor->numpy conversion: contiguous() ensures optimal memory layout
         flow_np = predicted_flows[0].contiguous().cpu().numpy()