raft update

src/orisys/tracker.py

@@ -25,7 +25,7 @@ class Tracker:
         self.__frame_count = 0
         dis_preset = config_tracker["dis_preset"]
         loop_closure_interval = config_tracker["loop_closure_interval"]
-        if not cuda:
+        if not self.cuda:
             # DIS method - optimize for speed
             if dis_preset == 'ultrafast':
                 preset = cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST
@@ -53,9 +53,10 @@ class Tracker:
         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_model = self.setup_simple_raft()
         self.__raft_ref_frame = None
+        self.__raft_ref_tensor = None
 
 
     def track(self, frame):
@@ -232,61 +233,78 @@ class Tracker:
         self.__fine_flow_queue = []
         self.__prev_fine_flow = None
         self.__frame_count = 0
+        self.__raft_ref_frame = None
+        self.__raft_ref_tensor = None
 
     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')
+        # Load model and move to GPU
+        model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
         model = model.eval()
-        # Alternative: Use LiteFlowNet (lighter alternative)
-        print("RAFT loaded successfully")
+            
+        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"""
+        # Upload to device as uint8 first to save bandwidth
+        # permute returns a view, but from_numpy usually shares memory.
+        # Moving to GPU as uint8 is fast.
+        tensor = torch.from_numpy(frame).permute(2, 0, 1).to(self.__raft_device)
+        
+        # Convert to float and normalize on device
+        tensor = tensor.float() / 255.0
+        
+        # Resize if needed
+        if resize is not None:
+             tensor = T.Resize(resize)(tensor)
+
+        # Normalize to [-1, 1]
+        tensor = (tensor - 0.5) * 2.0
+        
+        # Add batch dimension
+        tensor = tensor.unsqueeze(0)
+        
+        # Keep in float32 - autocast will handle mixed precision during inference
+        return tensor
+
     def track_raft(self, frame):
+        # No need to copy if we just read
+        # frame is expected to be (H, W, 3) or (H, W)
+        
+        if frame.ndim == 2:
+            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
+        elif frame.shape[2] == 4:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2RGB)
+        elif frame.shape[2] == 3:
+             # Assume BGR input from OpenCV
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
-        gray = frame.copy()
         if self.__raft_ref_frame is None:
-            self.__raft_ref_frame = gray.copy()
-        prev_frame = self.__raft_ref_frame
+            self.__raft_ref_frame = frame  # Keep reference for potential debug/CPU use
+            # Cache the processed tensor to avoid re-processing reference frame every time
+            self.__raft_ref_tensor = self._preprocess_frame_to_tensor(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)
+        prev_tensor = self.__raft_ref_tensor
+        curr_tensor = self._preprocess_frame_to_tensor(frame)
+        
+        # Use autocast for automatic mixed precision (FP16 where supported, FP32 where needed)
         with torch.no_grad():
+            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]
 
-        flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
+        flow_np = predicted_flows[0].float().cpu().numpy() # Convert back to float32 for numpy
         
         # Transpose to (H, W, 2) format expected by draw_arrows
         flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)