update raft optical flow

2026-01-31 15:22:19 +08:00
parent 2ebbe1a284
commit 58c3d97fab
7 changed files with 117 additions and 141 deletions
@@ -28,6 +28,7 @@ class Stitcher:
        self.margin=config_stitcher["margin"]
        self.bound=self.nominal_ov_pos - self.margin
        self.row_tol = config_stitcher["row_tol"]
+        self.box_size = config_stitcher["box_size"]

        if self.verbose:
            print(f"Stitcher initialized with config: {self.config_stitcher}")
@@ -486,8 +487,10 @@ class Stitcher:
        # compute common center-crop size: use minimum width/height among boxes
        widths  = [b[2] for b in boxes]
        heights = [b[3] for b in boxes]
-        target_w = max(1,int(min(widths)))
-        target_h = max(1,int(min(heights)))
+        # target_w = max(1,int(min(widths)))
+        # target_h = max(1,int(min(heights)))
+        target_w = self.box_size[0]
+        target_h = self.box_size[1]

        # center-crop each box to (target_w, target_h) inside original box bounds
        cropped_boxes = []
@@ -164,7 +164,7 @@ class Sensor:
        if not ret:
            print("Failed to get_img from video source.")
            return None
-        self.frame = frame
+        self.frame = frame.copy()
        # Convert to grayscale
        # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        gray = frame
@@ -215,8 +215,8 @@ class Sensor:
        if self.img is None:
            raise Exception("Warning: Get image first")
        
-        # flow = self.__tracker.track(self.img)
-        flow = self.__tracker.track_raft(self.img)
+        flow = self.__tracker.track(self.img)
+        # flow = self.__tracker.track_raft(self.img)
        if hasattr(flow, 'download'):
            flow = flow.download()
        self.flow = flow
@@ -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
@@ -2,7 +2,7 @@ import cv2
 import numpy as np
 from scipy.interpolate import griddata

-def draw_arrows(image, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0):
+def draw_arrows(frame, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0):
    """
    Draw colored optical flow arrows on sensor image based on flow magnitude.
    
@@ -16,7 +16,7 @@ def draw_arrows(image, optical_flow,fps=True, threshold=5, grid_spacing=40, arro
    
    """
    # Create working copies to avoid modifying originals
-    
+    image = frame.copy()
    # Convert grayscale to BGR to enable colored arrow drawing
    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)