Merge branch 'remote' into dev

.gitignore

@@ -11,3 +11,4 @@ orisys.egg-info/
 /dist_nuitka/
 /build_nuitka
 /data/
+/models/

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,101 +36,78 @@ def setup_simple_raft():
     """
     Setup a simpler version using available models
     """
-    # Check for torchvision RAFT (available in torchvision 0.13+)
+    model = raft_large(pretrained=True, progress=False).to('cuda')
-    # 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_small(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)
-    hsv[..., 1] = 255
-    
-    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
-    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)
-    return flow_rgb
-
-def main_simple():
-    # Initialize camera
-    sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
-    model = setup_simple_raft()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    # cap = cv2.VideoCapture(0)
-    # cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
-    # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
-    
-    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)
-        sensor.get_img() 
-        frame = sensor.img.copy()
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        if prev_gray is None:
-            prev_gray = frame.copy
-        gray = frame
-        # Compute optical flow
-        if prev_gray is not None:
     start_time = time.time()
-            
+    prev_tensor = torch.from_numpy(prev_frame.transpose(2, 0, 1).astype(np.float32) / 255.0)
-            # 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)
-            
+    prev_tensor = preprocess(prev_tensor.unsqueeze(0), resize=None).to(device)
-            # Preprocess: add batch dimension, preprocess (keeps batch dimension)
+    curr_tensor = preprocess(curr_tensor.unsqueeze(0), resize=None).to(device)
-            # preprocess expects (N, C, H, W), returns (N, C, H, W)
+    t0 = time.perf_counter()
-            # 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}")
+    t1 = time.perf_counter()
-            print(f"dtype = {predicted_flows.dtype}")
+    print(f"RAFT time: {t1 - t0}")
-            print(f"shape = {predicted_flows.shape} = (N, 2, H, W)")
+    flow_np = predicted_flows[0].cpu().numpy()  # or .squeeze(0)
-            print(f"min = {predicted_flows.min()}, max = {predicted_flows.max()}")
     
-            # Calculate FPS
+    # 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():
+    # Initialize camera
+    # sensor = orisys.Sensor('s1.mp4', config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
+    model = setup_simple_raft()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    cap = cv2.VideoCapture("s1.mp4")
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
+    
+    prev_gray = None
+    fps_buffer = []
+
+    
+    while True:
+
+        ret, frame = cap.read()
+        if not ret:
+            break
+        dx = 300
+        dy = 0
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame1 = frame[420:660, 840:1080, :]
+        frame2 = frame[420+dy:660+dy, 840+dx:1080+dx, :]
+        frame3 = frame[420-dy:660-dy, 840-dx:1080-dx, :]
+
+        frame = np.concatenate([frame3, frame1, frame2], axis=1)
+        # cv2.imshow("frame", frame3)
+        # cv2.waitKey(0)
+        if prev_gray is None:
+            prev_gray = frame.copy()
+        gray = frame
+        # Compute optical flow
+        start_time = time.time()
+        flow_np = track_raft(prev_gray, gray, model, device)
         fps = 1.0 / (time.time() - start_time)
         fps_buffer.append(fps)
         if len(fps_buffer) > 30:
@@ -159,28 +117,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)
@@ -191,17 +132,12 @@ def main_simple():
             arrows = orisys.util.draw_arrows(
                 gray,
                 flow_np,
-                threshold=1,
+                threshold=2,
-                grid_spacing=10,
+                grid_spacing=40,
                 arrow_scale=1.0
             )
             
-            # Display results
+            arrows = cv2.resize(arrows, (1200, 400))
-            # 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)
             cv2.imshow('Optical Flow', arrows)
         
         # prev_gray = gray.copy()
@@ -213,7 +149,7 @@ def main_simple():
             break
         if key == ord("r"):
             # 重置光流追踪器（当追踪漂移时使用）
-            flow_track = None
+            prev_gray = None
             print("传感器已重置")
     
 if __name__ == "__main__":

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,
@@ -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}")
@@ -170,9 +171,10 @@ class Stitcher:
     
     def extract_tiles(self):
         if not self.isstitch:
-            tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1])
+            # tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1])
-            if tile_idx is None:
             tile_idx = 7
+            # if tile_idx is None:
+            #     tile_idx = 7
             # Ensure index is always iterable (list or array)
             self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx
             self.ncol = 1
@@ -485,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,10 +164,10 @@ 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 = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-
+        gray = frame
         # Extract and stitch tiles
         self.stitcher.get_stitch_unit(gray)
         self._extract_tiles(gray)
@@ -216,6 +216,7 @@ class Sensor:
             raise Exception("Warning: Get image first")
         
         flow = self.__tracker.track(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,9 @@
 import cv2
 import numpy as np
+import torch
+import torchvision.transforms as T
+from torchvision.models.optical_flow import raft_large, raft_small
+import time
 
 class Tracker:
     def __init__(self,
@@ -12,9 +16,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
@@ -22,7 +26,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
@@ -36,15 +40,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,13 +53,22 @@ 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_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
 
 
     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)
             return self.track_cpu(frame)
     
     def track_cuda(self, frame):
@@ -225,7 +235,110 @@ 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 setup_simple_raft(self):
+        """
+        Setup a simpler version using available models with performance optimizations
+        """
+        # Load model and move to GPU
+        model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
+        model = model.eval()
+        
+        # Optimize matmul precision for faster inference (TF32 on Ampere+ GPUs)
+        # torch.set_float32_matmul_precision('medium')
+        
+        # Enable CUDA-specific optimizations
+        if self.__raft_device == 'cuda':
+            # cuDNN optimizations
+            torch.backends.cudnn.benchmark = True  # Auto-select fastest algorithms
+            torch.backends.cudnn.deterministic = False  # Allow non-deterministic for speed
+            # TensorFloat-32 for Ampere+ GPUs (faster with minimal accuracy loss)
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            
+        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
+        if not frame.flags['C_CONTIGUOUS']:
+            frame = np.ascontiguousarray(frame)
+        
+        # Upload to device as uint8 first to save bandwidth (4x less data than float32)
+        tensor = torch.from_numpy(frame).permute(2, 0, 1)
+        
+        # Move to device with non_blocking for async transfer (if CUDA)
+        # Pin memory for faster CPU-GPU transfer (similar to DataLoader pin_memory=True)
+        if self.__raft_device == 'cuda':
+            # Pin memory before transferring to GPU (must be on CPU first)
+            if not tensor.is_pinned():
+                tensor = tensor.pin_memory()
+            tensor = tensor.to(self.__raft_device, non_blocking=True)
+        else:
+            tensor = tensor.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)
+        
+        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)
+
+        if self.__raft_ref_frame is None:
+            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 = self.__raft_ref_tensor
+        curr_tensor = self._preprocess_frame_to_tensor(frame)
+        
+        # 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()
+        
+        # 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)