initial commit

This commit is contained in:
unknown
2025-12-21 17:15:00 +08:00
commit f38b2138cb
86 changed files with 3748 additions and 0 deletions
+486
View File
@@ -0,0 +1,486 @@
"""
Stch Calibration OK
"""
import os
import numpy as np
from scipy import ndimage
import cv2
from .Stitcher import Stitcher
from .tracker import Tracker
from .nHHD import *
import time
from scipy.interpolate import griddata
from enum import Enum, auto
from typing import Optional, Tuple, Union, Any
class Sensor:
"""
Sensor helper that can decompose optical flow via Fourier methods.
Stores the decomposed components and global norms as attributes after compute().
"""
def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True):
"""
Initialize Sensor with video source and configuration.
Args:
vid_src: Video source (int for camera ID, str for file path, None for default camera)
isstitch: Whether to use stitching mode for multi-tile sensor (default: False)
verbose: Whether to print initialization logs (default: True)
cuda: Whether to use CUDA acceleration for optical flow & decomposition (default: True)
"""
try:
import cupy as cp
import cupyx.scipy.signal as cp_signal
self.CUPY_AVAILABLE = True
except ImportError:
self.CUPY_AVAILABLE = False
# Initialize settings
if vid_src is not None:
vid_src = vid_src
else:
print("User must specify a video source.")
self.vid_src, self.isstitch, self.cuda, self.verbose = vid_src, isstitch, cuda, verbose
self.index = np.array([1,2,3,6,7,8,11,12,13])
# Initialize components
self.__stitcher = Stitcher(isstitch=isstitch, index=self.index)
self.tracker = Tracker(adaptive=True, cuda=cuda)
self.hhd = None
if self.verbose:
print("Components initialized")
# Setup video capture
self.cap = cv2.VideoCapture(vid_src)
self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G'))
self.cap.set(cv2.CAP_PROP_FPS, 30)
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
if not self.cap.isOpened():
raise RuntimeError("Could not open video source")
# Wait for first valid frame
if self.verbose:
print("Sensor connecting to video source...")
frame = None
while frame is None:
ret, frame = self.cap.read()
if not ret:
time.sleep(1)
if self.verbose:
print("Video source connected.")
# Initialize and calibrate
self.__find_stitch_boxes(frame, show=False)
self.__get_tiles(frame)
self.stitch_cal()
if self.verbose:
print("Sensor initialized and calibrated.")
# Initialize state variables
self.stitch_valid = False
self.n_boxes = None
self.stitch_calibration = []
self.img = self.flow = self.normal = self.shear = self.force = None
self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0
self.imgupdate = self.vupdate = False
self.pos_default = [(260, 56), (577, 62), (897, 69), (1214, 75), (1530, 81), (256, 391), (570, 397), (889, 403), (1208, 409), (1523, 416), (248, 724), (563, 730), (882, 738), (1201, 743), (1516, 751)]
# Logging state
self.last_log_time = time.time()
self.fps = 0
if self.verbose:
print("Settings:")
print(f" Video source: {vid_src}")
print(f" Stitch mode: {isstitch}")
print(f" CUDA acceleration: {cuda}")
print(f" Cupy available: {self.CUPY_AVAILABLE}")
print(f" OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")
class info(Enum):
VRAW = auto()
VNORMAL = auto()
VSHEAR = auto()
FPS = auto()
TIMESTAMP = auto()
FNORMAL = auto()
FSHEARX = auto()
FSHEARY = auto()
def get_img(self):
"""
Capture frame from video source, convert to grayscale, and perform stitching.
Returns:
np.ndarray: Stitched grayscale image (contiguous array) or None if capture fails
"""
ret, frame = self.cap.read()
if not ret:
print("Failed to get_img from video source.")
self.cap.release()
exit(1)
return None
gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
self.__stitcher.get_stitch_unit(gray_frame)
self.__get_tiles(gray_frame)
img = np.ascontiguousarray(self.__stitcher.stitch_rt_optimized())
self.img = cv2.resize(img, (400, 400), interpolation=cv2.INTER_CUBIC)
self.imgupdate = True
return self.img
def compute_deformation(self):
"""
Track optical flow and decompose into force field.
Returns:
tuple: (flow, force) where:
- flow: Optical flow field (2D vector field)
- force: Decomposed force components from nHHD, ([normal, shear_x, shear_y, torque]). Note: torque is experimental.
"""
flow = self.tracker.track(self.img)
if hasattr(flow, 'download'):
flow = flow.download()
if self.hhd is None:
self.hhd = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18))
self.hhd.decompose(flow)
self.imgupdate = False
self.vupdate = True
self.flow = flow
self.normal = self.hhd.get_normal()
self.shear = self.hhd.get_shear()
flag, [self.fnormal,self.fshearx,self.fsheary,self.ftor] = self.hhd.get_force(flow)
self.fps = 1/(time.time() - self.last_log_time)
self.last_log_time = time.time()
return self.flow
def read_info(self, *args: info) -> Tuple:
"""
Get sensor data according to specified info types.
Args:
*args: Variable number of info enum values specifying desired data
Returns:
Tuple: Data in the order specified by args. Each element corresponds to:
- VRAW: Raw optical flow field (np.ndarray)
- VNORMAL: Normal component of flow (np.ndarray)
- VSHEAR: Shear component of flow (np.ndarray)
- FNORMAL: Normal force magnitude (float)
- FSHEAR: Shear force magnitude (float)
- FPS: Current frames per second (float)
- TIMESTAMP: Current timestamp string (str)
Example:
normal, shear, fps = sensor.read_info(sensor.info.VNORMAL, sensor.info.VSHEAR, sensor.info.FPS)
Note:
Must call compute_deformation() before reading flow/force data
"""
if not self.vupdate:
print("Warning: Update field first via compute_deformation() before reading flow/force info.")
results = []
for arg in args:
if not isinstance(arg, self.info):
print(f"Warning: Invalid argument type {type(arg)}, skipping")
continue
if arg == self.info.VRAW:
results.append(self.flow if self.flow is not None else np.array([]))
elif arg == self.info.VNORMAL:
results.append(self.normal if self.normal is not None else np.array([]))
elif arg == self.info.VSHEAR:
results.append(self.shear if self.shear is not None else np.array([]))
elif arg == self.info.FNORMAL:
results.append(self.fnormal)
elif arg == self.info.FSHEARX:
results.append(self.fshearx)
elif arg == self.info.FSHEARY:
results.append(self.fsheary)
elif arg == self.info.FPS:
results.append(self.fps)
elif arg == self.info.TIMESTAMP:
results.append(time.strftime('%H:%M:%S'))
return tuple(results) if results else ()
def disconnect(self):
"""
Release video capture resources.
"""
if self.cap.isOpened():
self.cap.release()
if self.verbose:
print("Sensor disconnected.")
cv2.destroyAllWindows()
def __get_tiles(self, img_src):
"""
Extract individual tiles from source image for stitching.
Args:
img_src: Source image (grayscale) containing multiple vision tiles
Returns:
bool: True (always succeeds)
"""
self.__stitcher.get_stitch_unit(img_src)
self.__stitcher.extract_tiles()
return True
def __find_stitch_boxes(self, img_src, show=False):
"""
Detect vision tile boxes in the frame, waiting for camera stabilization.
Args:
img_src: Initial source image (not actively used, boxes detected from live capture)
show: Whether to display detected boxes visually (default: False)
Returns:
int: Number of boxes detected (should be 15 for valid stitching)
"""
n_boxes, v_mean0, timeout = 0, 0, 10
start_time = cv2.getTickCount()
# print("Waiting for camera stabilization...")
while n_boxes < 15:
ret, frame = self.cap.read()
v_mean = np.sum(cv2.mean(frame))
if abs(v_mean - v_mean0) < 5:
n_boxes = self.__stitcher.find_stitch_boxes(frame, show=show)
if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout:
print("Timeout for finding valid vision units.")
self.isstitch = False
self.__stitcher.isstitch = False
break
v_mean0 = v_mean
# print("Camera stabilized.")
# print("Checking validabity for stitching...")
self.n_boxes = n_boxes
self.stitch_valid = self.isstitch and n_boxes == 15
if self.isstitch:
print("Sensor is valid for stithcing." if self.stitch_valid else "Sensor not valid to stitching, switching to non-stitch mode")
else:
print("Using non-stitch mode.")
return n_boxes
def stitch_cal(self):
"""
Load or create stitch calibration for aligning multiple tiles.
Returns:
np.ndarray: Stitched image after calibration
Note:
- Loads calibration from ./config/cal.npy if exists
- Creates new calibration and saves to ./config/cal.npy if not found
- Calibration contains transformation parameters for tile alignment
"""
cal_path = "./config/cal.npy"
if os.path.exists(cal_path):
print("Loading existing calibration from", cal_path)
self.stitch_calibration = np.load(cal_path, allow_pickle=True)
self.__stitcher.stitch_calibration = self.stitch_calibration
img_stitched_height = self.__stitcher.stitch_rt_optimized()
else:
print("Calibration file not found, creating new calibration...")
img_stitched_height, self.stitch_calibration = self.__stitcher.stitch_cal()
os.makedirs("./config", exist_ok=True)
np.save(cal_path, self.stitch_calibration)
print("Calibration saved to", cal_path)
# print("Image stitched with shape:", img_stitched_height.shape)
return img_stitched_height
def draw_arrows(self, info_type: info,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0) -> np.ndarray:
"""
Draw colored optical flow arrows on sensor image based on flow magnitude.
Args:
threshold: Minimum flow magnitude to display arrow (default: 10.0)
grid_spacing: Spacing between arrows in pixels (default: 20)
arrow_scale: Scale factor for arrow length (default: 3.0)
Returns:
np.ndarray: 400x400 BGR image with colored arrows (JET colormap)
"""
# Create working copies to avoid modifying originals
image = self.img.copy()
if info_type is None:
print("Flow field type not specified, defaulting to show raw vector field.")
optical_flow = self.flow.copy()
elif info_type == self.info.VRAW:
optical_flow = self.flow.copy()
elif info_type == self.info.VNORMAL:
optical_flow = self.normal.copy()
elif info_type == self.info.VSHEAR:
optical_flow = self.shear.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)
# Calculate scaling factors between image and flow dimensions
img_height, img_width = image.shape[:2]
flow_height, flow_width = optical_flow.shape[:2]
scale_x, scale_y = img_width / flow_width, img_height / flow_height
# Create sampling grid for flow field (used for interpolation source points)
sample_step = 19 # Spacing between sample points
flow_start = np.stack(np.meshgrid(range(0, flow_width, sample_step), range(0, flow_height, sample_step)), 2)
flow_end = (optical_flow[flow_start[:, :, 1], flow_start[:, :, 0], :] * 2.0 + flow_start).astype(np.int32)
# Create regular grid for arrow placement in image coordinates
grid_x, grid_y = np.mgrid[grid_spacing//2:flow_height:grid_spacing, grid_spacing//2:flow_width:grid_spacing]
x_axis = (grid_y.reshape(-1) * scale_x).astype(np.int64) # Scale to image width
y_axis = (grid_x.reshape(-1) * scale_y).astype(np.int64) # Scale to image height
# Interpolate flow vectors from sample grid to regular arrow grid using cubic interpolation
grid_vector = griddata(
flow_start.reshape(-1, 2),
(flow_end - flow_start).reshape(-1, 2), # Flow from start to end (original to new position)
(grid_y, grid_x), method='cubic', fill_value=0
)
grid_vector[np.isnan(grid_vector)] = 0 # Replace NaN with zero for stability
grid_flat_scaled = grid_vector.reshape(-1, 2) * np.array([scale_x, scale_y])
# Compute arrow endpoints and ensure they stay within image bounds
end_x = (grid_flat_scaled[:, 0] / arrow_scale + x_axis).astype(np.int64)
end_y = (grid_flat_scaled[:, 1] / arrow_scale + y_axis).astype(np.int64)
valid_mask = (y_axis >= 0) & (y_axis < img_height) & (x_axis >= 0) & (x_axis < img_width)
valid_mask &= (end_x >= 0) & (end_x < img_width) & (end_y >= 0) & (end_y < img_height)
# Filter arrows by magnitude threshold and compute colors using JET colormap
magnitudes = np.linalg.norm(grid_flat_scaled, axis=1)
valid_mask &= (magnitudes >= 0) # Only show arrows above threshold
valid_indices = np.where(valid_mask)[0]
mag_range = magnitudes.max() - magnitudes.min()
# Normalize magnitudes to 0-255 range for colormap application, but set below-threshold values to 0 (blue)
norm_magnitudes = ((magnitudes - magnitudes.min()) / mag_range * 255.0).astype(np.uint8) if mag_range > 0 else np.zeros_like(magnitudes, dtype=np.uint8)
norm_magnitudes[magnitudes < threshold] = 0 # Keep lowest color (blue) for below-threshold magnitudes
colors_array = cv2.applyColorMap(norm_magnitudes.reshape(-1, 1), cv2.COLORMAP_JET) # JET: blue (low) -> red (high)
# Draw arrows at each valid grid point with magnitude-based coloring
for i in valid_indices:
color = tuple(map(int, colors_array[i, 0]))
# Scale arrow length and thickness based on threshold
if magnitudes[i] < threshold:
# Below threshold: draw as small dot (zero-length arrow with larger tip)
cv2.circle(image, (x_axis[i], y_axis[i]), radius=1, color=color, thickness=-1)
else:
# Above threshold: draw full arrow with thickness proportional to magnitude
# Arrow rooted at original position (x_axis, y_axis), head points to new position (end_x, end_y)
thickness = max(1, int(2 * magnitudes[i] / magnitudes.max())) if magnitudes.max() > 0 else 2
cv2.arrowedLine(image, (x_axis[i], y_axis[i]), (end_x[i], end_y[i]), color, thickness=thickness, tipLength=0.3)
out_img = cv2.resize(image, (400, 400), interpolation=cv2.INTER_CUBIC)
# if fps:
# out_img = self.draw_fps(image=out_img)
return out_img
def draw_magnitude_map(self, info_type: info = None, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.ndarray:
"""
Draw color map of pointwise magnitude of velocity field.
Args:
info_type: Type of flow field to visualize (VRAW, VNORMAL, VSHEAR, or None for total)
threshold: Minimum magnitude to display (below threshold shows as zero/blue)
colormap: OpenCV colormap to use (default: COLORMAP_JET)
Returns:
np.ndarray: 400x400 BGR image with colored magnitude map
"""
# Get the appropriate flow field
if info_type is None or info_type == self.info.VRAW:
flow_field = self.flow.copy() if self.flow is not None else None
elif info_type == self.info.VNORMAL:
flow_field = self.normal.copy() if self.normal is not None else None
elif info_type == self.info.VSHEAR:
flow_field = self.shear.copy() if self.shear is not None else None
else:
# For combined normal + shear magnitude
if self.normal is not None and self.shear is not None:
flow_field = self.normal + self.shear
else:
flow_field = None
if flow_field is None:
# Return blank image if no flow data
return np.zeros((400, 400, 3), dtype=np.uint8)
# Compute magnitude at each point
magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
# Smooth by downsampling and interpolating back
downsample_factor = 10
h, w = magnitude.shape
magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)
# Apply threshold - set below-threshold values to zero
magnitude_thresholded = magnitude_smoothed.copy()
magnitude_thresholded[magnitude_smoothed < threshold] = 0
# Normalize to 0-255 range
mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
if mag_max > mag_min:
magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
else:
magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)
# Apply colormap
magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
# Resize to output size
out_img = cv2.resize(magnitude_colored, (400, 400), interpolation=cv2.INTER_CUBIC)
return out_img
def draw_fps(self, image=None):
"""
Draw FPS counter on the top left corner of the image.
"""
font_scale=0.5
thickness=2
color=(0, 255, 0)
image = image.copy()
fps = self.fps
# Draw FPS text
fps_text = f"FPS: {fps:.1f}"
position = (10, 30) # Top left corner with margin
font = cv2.FONT_HERSHEY_SIMPLEX
# Draw FPS text
cv2.putText(image, fps_text, position, font, font_scale, color, thickness, cv2.LINE_AA)
return image