orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update raft optical flow

Browse Source
This commit is contained in:
bluepulsar
2026-01-31 15:22:19 +08:00
parent 2ebbe1a284
commit 58c3d97fab
7 changed files with 117 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
RAFT_test.py
+29 -117
View File
@@ -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__":
config/prototype_of3.json
+2 -1
View File
@@ -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
+2 -2
View File
@@ -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
+5 -2
View File
@@ -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 = []
src/orisys/sensor.py
+3 -3
View File
@@ -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
src/orisys/tracker.py
+74 -14
View File
@@ -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
src/orisys/util.py
+2 -2
View File
@@ -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)