bluepulsar
524 lines
22 KiB
Python
524 lines
22 KiB
Python
"""
|
||
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
|
||
from .getDensity import getFilteredDensity
|
||
from . import configs
|
||
|
||
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, config_name="prototype"):
|
||
"""
|
||
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.config_name = config_name
|
||
|
||
# Initialize Sensor Configuration
|
||
self.get_config()
|
||
|
||
# Initialize components
|
||
self.__stitcher = Stitcher(isstitch=isstitch, index=self.config_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, self.config_fps)
|
||
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.config_resolution_raw[0])
|
||
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.config_resolution_raw[1])
|
||
|
||
# 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 = self.depth = None
|
||
self.img_filtered = None # For temporal noise reduction
|
||
self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0
|
||
self.imgupdate = self.vupdate = False
|
||
|
||
# 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()
|
||
DEPTH = auto()
|
||
|
||
def get_config(self):
|
||
"""
|
||
Load and return sensor configuration from JSON file.
|
||
|
||
Returns:
|
||
dict: Dictionary containing sensor configuration
|
||
"""
|
||
self.config = configs.load_config(self.config_name)
|
||
self.config_resolution_raw = self.config["camera"]["resolution_raw"]
|
||
self.config_fps = self.config["camera"]["fps"]
|
||
self.config_units = self.config["camera"]["units"]
|
||
self.config_resolution_stitch = self.config["camera"]["resolution_stitch"]
|
||
self.config_index = self.config["camera"]["index"]
|
||
self.config_pos_default = self.config["camera"]["pos_default"]
|
||
|
||
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())
|
||
|
||
img_resized = cv2.resize(img, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
|
||
|
||
# Apply temporal noise reduction using exponential moving average
|
||
# Blend current frame with previous filtered frame to reduce temporal noise
|
||
alpha = 0.7 # Weight for current frame (0.7 = 70% current, 30% previous)
|
||
if self.img_filtered is None:
|
||
# First frame: no previous frame to blend with
|
||
self.img_filtered = img_resized.astype(np.float32)
|
||
else:
|
||
# Temporal filtering: blend current frame with previous filtered frame
|
||
self.img_filtered = alpha * img_resized.astype(np.float32) + (1 - alpha) * self.img_filtered
|
||
|
||
# Convert back to uint8
|
||
self.img = self.img_filtered.astype(np.uint8)
|
||
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.depth = getFilteredDensity(self.flow, use_cuda=self.cuda)
|
||
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)
|
||
- DEPTH: Depth image (np.ndarray)
|
||
|
||
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'))
|
||
elif arg == self.info.DEPTH:
|
||
results.append(self.depth if self.depth is not None else np.array([]))
|
||
return tuple(results) if results else ()
|
||
|
||
def reset(self):
|
||
self.tracker.reset()
|
||
|
||
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 < self.config_units:
|
||
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 == self.config_units
|
||
|
||
|
||
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, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), 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((self.config_resolution_stitch[0], self.config_resolution_stitch[1], 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, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), 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
|
||
|
||
|
||
|
||
|