update raft optical flow

RAFT_test.py

@@ -8,25 +8,6 @@ from torchvision.models.optical_flow import raft_large, raft_small
 import numpy as np
 import time
 import orisys
-# from torchvision.utils import flow_to_image
-# import matplotlib.pyplot as plt
-
-# def plot(imgs, **imshow_kwargs):
-#     if not isinstance(imgs[0], list):
-#         # Make a 2d grid even if there's just 1 row
-#         imgs = [imgs]
-
-#     num_rows = len(imgs)
-#     num_cols = len(imgs[0])
-#     _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
-#     for row_idx, row in enumerate(imgs):
-#         for col_idx, img in enumerate(row):
-#             ax = axs[row_idx, col_idx]
-#             img = F.to_pil_image(img.to("cpu"))
-#             ax.imshow(np.asarray(img), **imshow_kwargs)
-#             ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
-
-#     plt.tight_layout()
 
 def preprocess(batch, resize=None):
     """
@@ -55,44 +36,29 @@ def setup_simple_raft():
     """
     Setup a simpler version using available models
     """
-    # Check for torchvision RAFT (available in torchvision 0.13+)
-    # try:
-    #     import torchvision
-    #     if hasattr(torchvision.models, 'optical_flow'):
-    #         raft = torchvision.models.optical_flow.raft_small(pretrained=True)
-    #         raft = raft.to('cuda')
-    #         raft.eval()
-    #         print("Loaded RAFT small from torchvision")
-    #         return raft
-    # except:
-    #     pass
     model = raft_large(pretrained=True, progress=False).to('cuda')
     model = model.eval()
     # Alternative: Use LiteFlowNet (lighter alternative)
     print("RAFT loaded successfully")
     return model
 
-def compute_dense_optical_flow_simple(prev_gray, curr_gray):
+def track_raft(prev_frame, frame, model, device):
-    """
-    Fallback to OpenCV optical flow if RAFT is not available
-    """
-    # Use OpenCV's Farneback optical flow
-    flow = cv2.calcOpticalFlowFarneback(
-        prev_gray, curr_gray, None,
-        pyr_scale=0.5, levels=3, winsize=15,
-        iterations=3, poly_n=5, poly_sigma=1.2, flags=0
-    )
 
-    # Convert flow to visualization
+    gray = frame.copy()
-    hsv = np.zeros((curr_gray.shape[0], curr_gray.shape[1], 3), dtype=np.uint8)
+    start_time = time.time()
-    hsv[..., 1] = 255
+    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 = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)
+    curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)
+    with torch.no_grad():
+        flow = model(prev_tensor, curr_tensor)
+        predicted_flows = flow[-1]
 
-    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
+    flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
-    hsv[..., 0] = ang * 180 / np.pi / 2
-    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
     
-    flow_rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+    # Transpose to (H, W, 2) format expected by draw_arrows
-    return flow_rgb
+    flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)
+    return flow_np
 
 def main_simple():
     # Initialize camera
@@ -106,15 +72,10 @@ def main_simple():
     
     prev_gray = None
     fps_buffer = []
-    flow_track = None
+
     
     while True:
-        # ret, frame = cap.read()
-        # if not ret:
-        #     break
 
-        # # Convert to grayscale
-        # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
         ret, frame = cap.read()
         if not ret:
             break
@@ -132,34 +93,8 @@ def main_simple():
             prev_gray = frame.copy()
         gray = frame
         # Compute optical flow
-        if prev_gray is not None:
         start_time = time.time()
-            
+        flow_np = track_raft(prev_gray, gray, model, device)
-            # Use simple optical flow
-            # flow_viz = compute_dense_optical_flow_simple(prev_gray, gray)
-            
-            # Convert numpy arrays from (H, W, C) to (C, H, W) format
-            # Normalize to [0, 1] range
-            prev_tensor = torch.from_numpy(prev_gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
-            curr_tensor = torch.from_numpy(gray.transpose(2, 0, 1).astype(np.float32) / 255.0)
-            
-            # Preprocess: add batch dimension, preprocess (keeps batch dimension)
-            # preprocess expects (N, C, H, W), returns (N, C, H, W)
-            # Pass resize=None to keep original size (400x400)
-            prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)  # (1, C, H, W)
-            curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)  # (1, C, H, W)
-            print(f"prev_tensor shape: {prev_tensor.shape}")
-            
-            # RAFT expects tensors with batch dimension: (batch_size, C, H, W)
-            with torch.no_grad():
-                flow = model(prev_tensor, curr_tensor)
-                predicted_flows = flow[-1]
-                print(f"Flow shape: {predicted_flows.shape}")
-            print(f"dtype = {predicted_flows.dtype}")
-            print(f"shape = {predicted_flows.shape} = (N, 2, H, W)")
-            print(f"min = {predicted_flows.min()}, max = {predicted_flows.max()}")
-                        
-            # Calculate FPS
         fps = 1.0 / (time.time() - start_time)
         fps_buffer.append(fps)
         if len(fps_buffer) > 30:
@@ -169,28 +104,11 @@ def main_simple():
             print(f"avg_fps: {avg_fps}")
             
             
+            # # Remove batch dimension: (1, 2, H, W) -> (2, H, W)
+            # flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
             
-            # grid = [[img1, flow_img] for (img1, flow_img) in zip(img1_batch, flow_imgs)]
+            # # Transpose to (H, W, 2) format expected by draw_arrows
-            # plot(grid)
+            # flow_np = flow_np.transpose(1, 2, 0)  # (2, H, W) -> (H, W, 2)
-            # plt.show()
-            # Convert RAFT flow tensor to numpy array for draw_arrows
-            # RAFT flow shape: (1, 2, H, W) where:
-            #   - Channel 0: x-component (horizontal flow)
-            #   - Channel 1: y-component (vertical flow)
-            # Flow values are already in pixel coordinates (not normalized)
-            
-            # Remove batch dimension: (1, 2, H, W) -> (2, H, W)
-            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)
-            # if flow_track is None:
-            #     flow_track = np.zeros_like(flow_np)
-            # else:
-            #     if abs(np.mean(flow_np)) >0.1:
-            #         flow_track = flow_track + flow_np
-
-            # flow_np = flow_track
 
             # Flow is already in pixel coordinates, no denormalization needed
             # Keep as float32 for draw_arrows (it can handle float values)
@@ -206,12 +124,6 @@ def main_simple():
                 arrow_scale=1.0
             )
             
-            # Display results
-            # cv2.putText(frame, f"FPS: {avg_fps:.1f}", (10, 30),
-            #            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
-            
-            # cv2.imshow('Original', frame)
-            # cv2.imshow('Grayscale', gray)
             arrows = cv2.resize(arrows, (1200, 400))
             cv2.imshow('Optical Flow', arrows)
         
@@ -224,7 +136,7 @@ def main_simple():
             break
         if key == ord("r"):
             # 重置光流追踪器（当追踪漂移时使用）
-            flow_track = None
+            prev_gray = None
             print("传感器已重置")
     
 if __name__ == "__main__":

config/prototype_of3.json

@@ -16,7 +16,8 @@
     "margin": 10,
     "nominal_ov_pos": 100,
     "n": [3,3],
-    "row_tol": 0.6
+    "row_tol": 0.6,
+    "box_size": [250, 250]
   },
   "tracker": {
     "adaptive": true,

example1_basic_pipeline.py

@@ -21,7 +21,7 @@ def main():
     # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
 
     # 方式1: 从视频文件读取
-    sensor = orisys.Sensor("s1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=False, verbose=True)
+    sensor = orisys.Sensor("s1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=True, verbose=True)
 
     # 方式2: 从摄像头读取（取消注释使用） 
     # sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
@@ -58,7 +58,7 @@ def main():
             img,
             flow,
             threshold=4,
-            grid_spacing=10,
+            grid_spacing=40,
             arrow_scale=1.0
         )
         cv2.imshow("Deformation", arrows)

src/orisys/Stitcher.py

@@ -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_w = max(1,int(min(widths)))
-        target_h = max(1,int(min(heights)))
+        # 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 = []

src/orisys/sensor.py

@@ -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(self.img)
-        flow = self.__tracker.track_raft(self.img)
+        # flow = self.__tracker.track_raft(self.img)
         if hasattr(flow, 'download'):
             flow = flow.download()
         self.flow = 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():
+        # 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.")
+        #     # print("Warning: CUDA requested but OpenCV not built with CUDA support or no CUDA device found. Falling back to CPU.")
-            cuda = False
+        #     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
+        # else:
-            # self.__tracking_inst.setGradientDescentIterations(12)  # Reduce iterations (default 25)
+        #     # CUDA flow method with stream support
-            # self.__tracking_inst.setVariationalRefinementIterations(3)  # Reduce refinement (default 5)
+        #     self.__tracking_inst = cv2.cuda.FarnebackOpticalFlow_create()
-        else:
+        #     # Create CUDA stream for async operations
-            # CUDA flow method with stream support
+        #     self.__cuda_stream = cv2.cuda.Stream()
-            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)
@@ -231,3 +236,58 @@ 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

src/orisys/util.py

@@ -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)