import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 import torch import torchvision.transforms.functional as F import torchvision.transforms as T 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): """ 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(): """ 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_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): """ 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 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() # 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) 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) # 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) print(f"Flow numpy shape: {flow_np.shape}, dtype: {flow_np.dtype}") print(f"Flow range: x=[{flow_np[:,:,0].min():.2f}, {flow_np[:,:,0].max():.2f}], " f"y=[{flow_np[:,:,1].min():.2f}, {flow_np[:,:,1].max():.2f}]") arrows = orisys.util.draw_arrows( gray, flow_np, threshold=1, grid_spacing=10, 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) cv2.imshow('Optical Flow', arrows) # prev_gray = gray.copy() #步骤 3.5: 键盘控制 key = cv2.waitKey(1) & 0xFF if key == ord("q"): print("\n退出程序...") break if key == ord("r"): # 重置光流追踪器（当追踪漂移时使用） flow_track = None print("传感器已重置") if __name__ == "__main__": main_simple()