From 58c3d97fabf006f6ac4c8fe8fe075645b21ab025 Mon Sep 17 00:00:00 2001 From: bluepulsar Date: Sat, 31 Jan 2026 15:22:19 +0800 Subject: [PATCH] update raft optical flow --- RAFT_test.py | 146 ++++++++----------------------------- config/prototype_of3.json | 3 +- example1_basic_pipeline.py | 4 +- src/orisys/Stitcher.py | 7 +- src/orisys/sensor.py | 6 +- src/orisys/tracker.py | 88 ++++++++++++++++++---- src/orisys/util.py | 4 +- 7 files changed, 117 insertions(+), 141 deletions(-) diff --git a/RAFT_test.py b/RAFT_test.py index 020ae58..6896f4a 100644 --- a/RAFT_test.py +++ b/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): - """ - 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 - ) +def track_raft(prev_frame, frame, model, device): + + gray = frame.copy() + start_time = time.time() + 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] + + flow_np = predicted_flows[0].cpu().numpy() # or .squeeze(0) - # Convert flow to visualization - 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 + # 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 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,65 +93,22 @@ def main_simple(): prev_gray = frame.copy() gray = frame # Compute optical flow - if prev_gray is not None: - start_time = time.time() - - # 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: - fps_buffer.pop(0) + 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: + fps_buffer.pop(0) avg_fps = np.mean(fps_buffer) print(f"avg_fps: {avg_fps}") - - - # grid = [[img1, flow_img] for (img1, flow_img) in zip(img1_batch, flow_imgs)] - # plot(grid) - # 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) + # # 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 + # # Transpose to (H, W, 2) format expected by draw_arrows + # flow_np = flow_np.transpose(1, 2, 0) # (2, H, W) -> (H, W, 2) # 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__": diff --git a/config/prototype_of3.json b/config/prototype_of3.json index ac2758c..c8a0345 100644 --- a/config/prototype_of3.json +++ b/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, diff --git a/example1_basic_pipeline.py b/example1_basic_pipeline.py index a345a82..23de99d 100644 --- a/example1_basic_pipeline.py +++ b/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) diff --git a/src/orisys/Stitcher.py b/src/orisys/Stitcher.py index 259067d..6498204 100644 --- a/src/orisys/Stitcher.py +++ b/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_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 = [] diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 8e70cf7..86b3412 100644 --- a/src/orisys/sensor.py +++ b/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_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 diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index c30c28a..a799161 100644 --- a/src/orisys/tracker.py +++ b/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) - \ No newline at end of file + +### 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 \ No newline at end of file diff --git a/src/orisys/util.py b/src/orisys/util.py index cbeb8cd..ef1a1a1 100644 --- a/src/orisys/util.py +++ b/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)