prepare for calibration data collection
This commit is contained in:
unknown
+218
@@ -0,0 +1,218 @@
|
||||
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="prototype", 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)
|
||||
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()
|
||||
Reference in New Issue
Block a user