update raft optical flow

src/orisys/tracker.py

@@ -1,5 +1,8 @@
 import cv2
 import numpy as np
+import torch
+import torchvision.transforms as T
+from torchvision.models.optical_flow import raft_large
 
 class Tracker:
     def __init__(self,
@@ -12,9 +15,9 @@ class Tracker:
                  dis_preset='medium',
                  verbose=False):
         # Check if OpenCV is built with CUDA support
-        if cuda and not cv2.cuda.getCudaEnabledDeviceCount():
-            # print("Warning: CUDA requested but OpenCV not built with CUDA support or no CUDA device found. Falling back to CPU.")
-            cuda = False
+        # if cuda and not cv2.cuda.getCudaEnabledDeviceCount():
+        #     # print("Warning: CUDA requested but OpenCV not built with CUDA support or no CUDA device found. Falling back to CPU.")
+        #     cuda = False
         
         self.cuda = cuda
         self.verbose = verbose
@@ -36,15 +39,12 @@ class Tracker:
             
             # Aggressive optimization for CPU
             self.__tracking_inst.setFinestScale(config_tracker["of_mode"])      # Skip finest pyramid level (2x faster)
-            # self.__tracking_inst.setPatchSize(8)        # Smaller patches (default 8)
-            # self.__tracking_inst.setPatchStride(4)      # Larger stride = fewer patches
-            # self.__tracking_inst.setGradientDescentIterations(12)  # Reduce iterations (default 25)
-            # self.__tracking_inst.setVariationalRefinementIterations(3)  # Reduce refinement (default 5)
-        else:
-            # CUDA flow method with stream support
-            self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create()
-            # Create CUDA stream for async operations
-            self.__cuda_stream = cv2.cuda.Stream()
+
+        # else:
+        #     # CUDA flow method with stream support
+        #     self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create()
+        #     # Create CUDA stream for async operations
+        #     self.__cuda_stream = cv2.cuda.Stream()
             
         self.__fine_base_frame_queue = []
         self.__fine_flow_queue = []
@@ -52,12 +52,17 @@ class Tracker:
         self.__adaptive = adaptive
         self.__lower_threshold = config_tracker["lower_threshold"]  # error lower than threshold will close the loop
         self.__upper_threshold = config_tracker["upper_threshold"]  # error greater than threshold will add new reference frame
+        ### RAFT
+        self.__raft_model = self.setup_simple_raft()
+        self.__raft_device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.__raft_ref_frame = None
 
 
     def track(self, frame):
         """Main tracking function that dispatches to CUDA or CPU implementation"""
         if self.cuda:
-            return self.track_cuda(frame)
+            # return self.track_cuda(frame)
+            return self.track_raft(frame)
         else:
             if frame.ndim == 3:
                 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
@@ -230,4 +235,59 @@ class Tracker:
 
     def getFrameQueueLength(self):
         return len(self.__fine_base_frame_queue)
-    
+    
+### RAFT ###
+
+    def preprocess(self,batch, resize=None):
+        """
+        Preprocess image batch for RAFT model.
+        
+        Args:
+            batch: Tensor of shape (N, C, H, W)
+            resize: Optional tuple (height, width) to resize. If None, keeps original size.
+        
+        Returns:
+            Preprocessed tensor
+        """
+        transform_list = [
+            T.ConvertImageDtype(torch.float32),
+            T.Normalize(mean=0.5, std=0.5),  # map [0, 1] into [-1, 1]
+        ]
+        
+        if resize is not None:
+            transform_list.append(T.Resize(size=resize))
+        
+        transforms = T.Compose(transform_list)
+        batch = transforms(batch)
+        return batch
+    
+    def setup_simple_raft(self):
+        """
+        Setup a simpler version using available models
+        """
+        model = raft_large(pretrained=True, progress=False).to('cuda')
+        model = model.eval()
+        # Alternative: Use LiteFlowNet (lighter alternative)
+        print("RAFT loaded successfully")
+        return model
+
+    def track_raft(self, frame):
+
+        gray = frame.copy()
+        if self.__raft_ref_frame is None:
+            self.__raft_ref_frame = gray.copy()
+        prev_frame = self.__raft_ref_frame
+
+        prev_tensor = torch.from_numpy(prev_frame.transpose(2, 0, 1).astype(np.float32) / 255.0)
+        curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
+        prev_tensor = self.preprocess(prev_tensor.unsqueeze(0), resize=None).to(self.__raft_device)
+        curr_tensor = self.preprocess(curr_tensor.unsqueeze(0), resize=None).to(self.__raft_device)
+        with torch.no_grad():
+            flow = self.__raft_model(prev_tensor, curr_tensor)
+            predicted_flows = flow[-1]
+
+        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)
+        return flow_np