function rearrange for maintainability
This commit is contained in:
bluepulsar
+402
-440
@@ -1,87 +1,41 @@
|
||||
"""
|
||||
Sensor - Core sensor class for optical tactile sensing.
|
||||
|
||||
This module provides the Sensor class for capturing, processing, and analyzing
|
||||
optical tactile sensor data with support for stitching, optical flow tracking,
|
||||
and Helmholtz-Hodge decomposition.
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from .Stitcher import Stitcher
|
||||
from .tracker import Tracker
|
||||
from .nHHD import *
|
||||
import time
|
||||
from scipy.interpolate import griddata
|
||||
import time
|
||||
from enum import Enum, auto
|
||||
from typing import Optional, Tuple, Union, Any
|
||||
from .getDensity import getFilteredDensity
|
||||
from . import configs
|
||||
from typing import Tuple
|
||||
from threading import Thread, Event
|
||||
from queue import Queue, Empty
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from .Stitcher import Stitcher
|
||||
from .tracker import Tracker
|
||||
from .nHHD import nHHD
|
||||
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().
|
||||
Optical tactile sensor interface with deformation analysis.
|
||||
|
||||
Provides image capture, stitching, optical flow tracking, and force decomposition.
|
||||
Supports both synchronous and asynchronous (pipelined) processing modes.
|
||||
"""
|
||||
def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True, config_name="prototype", cal_path=None,override=False, box=None):
|
||||
"""
|
||||
Initialize Sensor with video source and configuration.
|
||||
|
||||
Args:
|
||||
vid_src: Video source (int for camera ID, str for file path)
|
||||
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)
|
||||
config_name: Name of the mdoel (default: "prototype")
|
||||
"""
|
||||
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 None:
|
||||
raise Exception("User must specify a video source.")
|
||||
|
||||
self.vid_src, self.isstitch, self.cuda, self.verbose, self.config_name, self.cal_path = vid_src, isstitch, cuda, verbose, config_name, cal_path
|
||||
|
||||
# Initialize state variables
|
||||
self.stitch_valid = self.imgupdate = self.vupdate = False
|
||||
self.stitch_calibration = []
|
||||
self.img = self.flow = self.normal = self.shear = self.force = self.depth = self.n_boxes = self.img_filtered = None
|
||||
self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0
|
||||
|
||||
# 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
|
||||
|
||||
# Setup video capture
|
||||
self.__open_video_source(override, box)
|
||||
|
||||
# Logging state
|
||||
self.last_log_time = time.perf_counter()
|
||||
self.fps = 0
|
||||
|
||||
# Async pipeline state (3-stage: Capture → Flow → HHD)
|
||||
self._async_enabled = False
|
||||
self._capture_thread = None
|
||||
self._flow_thread = None
|
||||
self._compute_thread = None
|
||||
self._capture_queue = None
|
||||
self._flow_queue = None
|
||||
self._result_queue = None
|
||||
self._stop_event = None
|
||||
|
||||
if self.verbose:
|
||||
print(f" Video source: {vid_src}")
|
||||
print(f" Stitch mode: {self.stitch_valid}")
|
||||
print(f" CUDA acceleration: {self.cuda}")
|
||||
print(f" Cupy available: {self.CUPY_AVAILABLE}")
|
||||
print(f" OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")
|
||||
|
||||
# =========================================================================
|
||||
# Info Enum (for read_info API)
|
||||
# =========================================================================
|
||||
|
||||
class info(Enum):
|
||||
"""Data types available through read_info()."""
|
||||
VRAW = auto()
|
||||
VNORMAL = auto()
|
||||
VSHEAR = auto()
|
||||
@@ -92,234 +46,273 @@ class Sensor:
|
||||
FSHEARY = auto()
|
||||
DEPTH = auto()
|
||||
|
||||
def __open_video_source(self, override, box):
|
||||
# =========================================================================
|
||||
# Initialization
|
||||
# =========================================================================
|
||||
|
||||
def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True,
|
||||
config_name="prototype", cal_path=None, override=False, box=None):
|
||||
"""
|
||||
Initialize Sensor with video source and configuration.
|
||||
|
||||
Args:
|
||||
vid_src: Video source (int for camera ID, str for file path)
|
||||
isstitch: Enable stitching mode for multi-tile sensor (default: False)
|
||||
verbose: Print initialization logs (default: True)
|
||||
cuda: Enable CUDA acceleration (default: True)
|
||||
config_name: Configuration name (default: "prototype")
|
||||
cal_path: Path to calibration file (default: None)
|
||||
override: Skip auto-detection, use provided box (default: False)
|
||||
box: Manual bounding boxes when override=True (default: None)
|
||||
"""
|
||||
# Check CUDA availability
|
||||
try:
|
||||
import cupy as cp
|
||||
self.CUPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
self.CUPY_AVAILABLE = False
|
||||
|
||||
if vid_src is None:
|
||||
raise Exception("User must specify a video source.")
|
||||
|
||||
# Store settings
|
||||
self.vid_src = vid_src
|
||||
self.isstitch = isstitch
|
||||
self.cuda = cuda
|
||||
self.verbose = verbose
|
||||
self.config_name = config_name
|
||||
self.cal_path = cal_path
|
||||
|
||||
# State variables
|
||||
self.stitch_valid = False
|
||||
self.imgupdate = False
|
||||
self.vupdate = False
|
||||
self.stitch_calibration = []
|
||||
|
||||
# Data variables
|
||||
self.img = None
|
||||
self.img_filtered = None
|
||||
self.flow = None
|
||||
self.normal = None
|
||||
self.shear = None
|
||||
self.depth = None
|
||||
self.n_boxes = None
|
||||
self.fnormal = self.fshearx = self.fsheary = self.ftor = 0.0
|
||||
|
||||
# Timing
|
||||
self.last_log_time = time.perf_counter()
|
||||
self.fps = 0
|
||||
|
||||
# Async pipeline state
|
||||
self._async_enabled = False
|
||||
self._capture_thread = None
|
||||
self._flow_thread = None
|
||||
self._compute_thread = None
|
||||
self._capture_queue = None
|
||||
self._flow_queue = None
|
||||
self._result_queue = None
|
||||
self._stop_event = None
|
||||
|
||||
# Load configuration
|
||||
self._load_config()
|
||||
|
||||
# Initialize components
|
||||
self._stitcher = Stitcher(isstitch=isstitch, index=self.config_index)
|
||||
self.tracker = Tracker(adaptive=True, cuda=cuda)
|
||||
self.hhd = None
|
||||
|
||||
# Setup video capture
|
||||
self.cap = cv2.VideoCapture(self.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
|
||||
if self.cap.isOpened():
|
||||
while frame is None:
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
time.sleep(1)
|
||||
else:
|
||||
raise Exception("Failed to open video source.")
|
||||
|
||||
if override:
|
||||
if box is not None:
|
||||
self.__stitcher.isstitch = self.isstitch = False
|
||||
self.__stitcher.boxes = box
|
||||
else:
|
||||
print("Override is True but no box provided.")
|
||||
else:
|
||||
# Initialize and calibrate
|
||||
self.__find_stitch_boxes(frame, show=False)
|
||||
self.__get_tiles(frame)
|
||||
self.stitch_cal()
|
||||
|
||||
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"]
|
||||
self._open_video_source(override, box)
|
||||
|
||||
if self.verbose:
|
||||
print("Config loaded: ", str(self.config_name))
|
||||
print(" Resolution raw: ", str(self.config_resolution_raw))
|
||||
print(" Camera FPS: ", str(self.config_fps))
|
||||
print(" Number of VisionUnits: ", str(self.config_units))
|
||||
print(" Resolution after stitch: ", str(self.config_resolution_stitch))
|
||||
print(" Index for stitching: ", str(self.config_index))
|
||||
print(f" Video source: {vid_src}")
|
||||
print(f" Stitch mode: {self.stitch_valid}")
|
||||
print(f" CUDA acceleration: {self.cuda}")
|
||||
print(f" CuPy available: {self.CUPY_AVAILABLE}")
|
||||
print(f" OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")
|
||||
|
||||
# =========================================================================
|
||||
# Public API - Core Methods (Most Frequently Used)
|
||||
# =========================================================================
|
||||
|
||||
def get_img(self, log_timing=False):
|
||||
"""
|
||||
Capture frame from video source, convert to grayscale, and perform stitching.
|
||||
Capture and preprocess a frame.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Stitched grayscale image (contiguous array) or None if capture fails
|
||||
np.ndarray: Processed grayscale image
|
||||
"""
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
print("Failed to get_img from video source.")
|
||||
exit(1)
|
||||
return None
|
||||
|
||||
if log_timing:
|
||||
timings = {}
|
||||
t_start = time.perf_counter()
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
if log_timing:
|
||||
timings['cvt_color'] = (time.perf_counter() - t0) * 1000
|
||||
# Convert to grayscale
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
if log_timing: timings['cvt_color'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
self.__stitcher.get_stitch_unit(gray_frame)
|
||||
self.__get_tiles(gray_frame)
|
||||
if log_timing:
|
||||
timings['get_tiles'] = (time.perf_counter() - t0) * 1000
|
||||
# Extract and stitch tiles
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
self._stitcher.get_stitch_unit(gray)
|
||||
self._extract_tiles(gray)
|
||||
if log_timing: timings['get_tiles'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Step 1: Optical flow tracking
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
img = np.ascontiguousarray(self.__stitcher.stitch_rt_optimized())
|
||||
|
||||
if log_timing:
|
||||
timings['stitch'] = (time.perf_counter() - t0) * 1000
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
img = np.ascontiguousarray(self._stitcher.stitch_rt_optimized())
|
||||
if log_timing: timings['stitch'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
img_resized = cv2.resize(img, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
|
||||
# Resize
|
||||
img = cv2.resize(img, tuple(self.config_resolution_stitch), 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)
|
||||
# Temporal filtering (noise reduction)
|
||||
alpha = 0.7
|
||||
if self.img_filtered is None:
|
||||
# First frame: no previous frame to blend with
|
||||
self.img_filtered = img_resized.astype(np.float32)
|
||||
self.img_filtered = img.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
|
||||
self.img_filtered = alpha * img + (1 - alpha) * self.img_filtered
|
||||
|
||||
# Convert back to uint8
|
||||
self.img = self.img_filtered.astype(np.uint8)
|
||||
self.imgupdate = True
|
||||
|
||||
if log_timing:
|
||||
timings['get_img'] = (time.perf_counter() - t_start) * 1000
|
||||
print("\n[get_img timing]")
|
||||
print(f" cvt_color: {timings['cvt_color']:7.2f} ms")
|
||||
print(f" get_tiles: {timings['get_tiles']:7.2f} ms")
|
||||
print(f" stitch: {timings['stitch']:7.2f} ms")
|
||||
print(f" ─────────────────────────")
|
||||
print(f" TOTAL: {timings['get_img']:7.2f} ms ({1000/timings['get_img']:.1f} FPS)")
|
||||
timings['total'] = (time.perf_counter() - t_start) * 1000
|
||||
self._print_timings("get_img", timings)
|
||||
|
||||
# Store timings for external access
|
||||
self._last_timings = timings
|
||||
return self.img
|
||||
|
||||
def compute_deformation(self, log_timing=False):
|
||||
"""
|
||||
Track optical flow and decompose into force field.
|
||||
|
||||
Args:
|
||||
log_timing: If True, print timing information for each computation step (default: False)
|
||||
Compute optical flow and decompose into force components.
|
||||
|
||||
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.
|
||||
|
||||
np.ndarray: Optical flow field
|
||||
"""
|
||||
if log_timing:
|
||||
timings = {}
|
||||
t_start = time.perf_counter()
|
||||
|
||||
# Step 1: Optical flow tracking
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Optical flow
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
flow = self.tracker.track(self.img)
|
||||
|
||||
if log_timing:
|
||||
timings['optical_flow'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
if hasattr(flow, 'download'):
|
||||
flow = flow.download()
|
||||
if log_timing: timings['optical_flow'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Initialize HHD solver
|
||||
if self.hhd is None:
|
||||
self.hhd = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18))
|
||||
self.hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
# Use GPU-accelerated force computation if CUDA is enabled
|
||||
# Force computation
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
if self.cuda and self.CUPY_AVAILABLE:
|
||||
# Use the most optimized GPU method with custom kernel
|
||||
flag, [self.fnormal,self.fshearx,self.fsheary,self.ftor] = self.hhd.get_force_gpu_v2(flow)
|
||||
_, forces = self.hhd.get_force_gpu_v2(flow)
|
||||
else:
|
||||
flag, [self.fnormal,self.fshearx,self.fsheary,self.ftor] = self.hhd.get_force(flow)
|
||||
# self.hhd.decompose(flow)
|
||||
if log_timing:
|
||||
timings['hhd_decompose'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
self.imgupdate = False
|
||||
self.vupdate = True
|
||||
self.flow = flow
|
||||
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
_, forces = self.hhd.get_force(flow)
|
||||
self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
|
||||
if log_timing: timings['hhd_decompose'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Get components
|
||||
self.normal = self.hhd.get_normal()
|
||||
self.shear = self.hhd.get_shear()
|
||||
|
||||
# Step 8: Compute depth/density
|
||||
if log_timing:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
self.depth = getFilteredDensity(self.flow, use_cuda=self.cuda)
|
||||
# Depth computation
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
self.depth = getFilteredDensity(flow, use_cuda=self.cuda)
|
||||
if log_timing: timings['get_density'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
if log_timing:
|
||||
timings['get_density'] = (time.perf_counter() - t0) * 1000
|
||||
# Update state
|
||||
self.flow = flow
|
||||
self.imgupdate = False
|
||||
self.vupdate = True
|
||||
|
||||
# Calculate FPS
|
||||
self.fps = 1/((time.perf_counter() - self.last_log_time))
|
||||
# FPS calculation
|
||||
self.fps = 1.0 / (time.perf_counter() - self.last_log_time)
|
||||
self.last_log_time = time.perf_counter()
|
||||
|
||||
# Print timing summary
|
||||
if log_timing:
|
||||
timings['total'] = (time.perf_counter() - t_start) * 1000
|
||||
print("\n[compute_deformation timing]")
|
||||
print(f" optical_flow: {timings['optical_flow']:7.2f} ms")
|
||||
print(f" hhd_decompose: {timings['hhd_decompose']:7.2f} ms")
|
||||
print(f" get_density: {timings['get_density']:7.2f} ms")
|
||||
print(f" ─────────────────────────")
|
||||
print(f" TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")
|
||||
|
||||
# Store timings for external access
|
||||
self._last_timings = timings
|
||||
self._print_timings("compute_deformation", timings)
|
||||
|
||||
return self.flow
|
||||
|
||||
def update(self):
|
||||
"""
|
||||
Convenience method: get_img + compute_deformation in one call.
|
||||
|
||||
Returns:
|
||||
dict: All sensor data
|
||||
"""
|
||||
self.get_img()
|
||||
self.compute_deformation()
|
||||
return {
|
||||
'img': self.img, 'flow': self.flow,
|
||||
'normal': self.normal, 'shear': self.shear, 'depth': self.depth,
|
||||
'fps': self.fps, 'fnormal': self.fnormal,
|
||||
'fshearx': self.fshearx, 'fsheary': self.fsheary, 'ftor': self.ftor,
|
||||
}
|
||||
|
||||
def read_info(self, *args) -> Tuple:
|
||||
"""
|
||||
Get sensor data by info type.
|
||||
|
||||
Args:
|
||||
*args: info enum values (e.g., sensor.info.FPS, sensor.info.FNORMAL)
|
||||
|
||||
Returns:
|
||||
Tuple: Requested data in order
|
||||
"""
|
||||
if not self.vupdate:
|
||||
print("Warning: Call compute_deformation() first.")
|
||||
|
||||
mapping = {
|
||||
self.info.VRAW: lambda: self.flow if self.flow is not None else np.array([]),
|
||||
self.info.VNORMAL: lambda: self.normal if self.normal is not None else np.array([]),
|
||||
self.info.VSHEAR: lambda: self.shear if self.shear is not None else np.array([]),
|
||||
self.info.FNORMAL: lambda: self.fnormal,
|
||||
self.info.FSHEARX: lambda: self.fshearx,
|
||||
self.info.FSHEARY: lambda: self.fsheary,
|
||||
self.info.FPS: lambda: self.fps,
|
||||
self.info.TIMESTAMP: lambda: time.strftime('%H:%M:%S'),
|
||||
self.info.DEPTH: lambda: self.depth if self.depth is not None else np.array([]),
|
||||
}
|
||||
|
||||
return tuple(mapping.get(arg, lambda: None)() for arg in args if isinstance(arg, self.info))
|
||||
|
||||
def reset(self):
|
||||
"""Reset optical flow tracker."""
|
||||
self.tracker.reset()
|
||||
|
||||
def disconnect(self):
|
||||
"""Release resources and stop async threads."""
|
||||
if self._async_enabled:
|
||||
self.disable_async()
|
||||
if self.cap.isOpened():
|
||||
self.cap.release()
|
||||
if self.verbose:
|
||||
print("Sensor disconnected.")
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
# =========================================================================
|
||||
# Async/Pipelined Processing API (3-stage: Capture → Flow → HHD)
|
||||
# Public API - Async Pipeline
|
||||
# =========================================================================
|
||||
|
||||
def _queue_put(self, q, item):
|
||||
"""Helper to put item in queue, dropping old if full."""
|
||||
if q.full():
|
||||
try: q.get_nowait()
|
||||
except Empty: pass
|
||||
q.put(item)
|
||||
|
||||
def enable_async(self, queue_size=2):
|
||||
"""
|
||||
Enable 3-stage async pipeline for maximum throughput.
|
||||
Enable 3-stage async pipeline: Capture → Flow → HHD.
|
||||
|
||||
Pipeline: [Capture] → [Optical Flow] → [HHD Compute]
|
||||
Each stage runs in parallel, overlapping all operations.
|
||||
Each stage runs in parallel for maximum throughput.
|
||||
"""
|
||||
if self._async_enabled:
|
||||
return
|
||||
|
||||
self._capture_queue = Queue(maxsize=queue_size) # img
|
||||
self._flow_queue = Queue(maxsize=queue_size) # (img, flow)
|
||||
self._result_queue = Queue(maxsize=queue_size) # final result
|
||||
self._capture_queue = Queue(maxsize=queue_size)
|
||||
self._flow_queue = Queue(maxsize=queue_size)
|
||||
self._result_queue = Queue(maxsize=queue_size)
|
||||
self._stop_event = Event()
|
||||
self._async_enabled = True
|
||||
|
||||
@@ -332,8 +325,169 @@ class Sensor:
|
||||
self._compute_thread.start()
|
||||
|
||||
if self.verbose:
|
||||
print("Async 3-stage pipeline enabled: Capture → Flow → HHD")
|
||||
print("Async pipeline enabled: Capture → Flow → HHD")
|
||||
|
||||
def get_async(self, timeout=0.05):
|
||||
"""
|
||||
Get latest result from async pipeline (non-blocking).
|
||||
|
||||
Returns:
|
||||
dict or None: Sensor data if available
|
||||
"""
|
||||
if not self._async_enabled:
|
||||
return None
|
||||
try:
|
||||
result = self._result_queue.get(timeout=timeout)
|
||||
t = time.perf_counter()
|
||||
result['fps'] = 1.0 / (t - self.last_log_time) if self.last_log_time else 0
|
||||
self.last_log_time = t
|
||||
|
||||
# Sync instance variables
|
||||
for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
setattr(self, k, result.get(k))
|
||||
self.vupdate = True
|
||||
return result
|
||||
except Empty:
|
||||
return None
|
||||
|
||||
def disable_async(self):
|
||||
"""Stop async pipeline."""
|
||||
if not self._async_enabled:
|
||||
return
|
||||
self._stop_event.set()
|
||||
for t in [self._capture_thread, self._flow_thread, self._compute_thread]:
|
||||
if t: t.join(timeout=2.0)
|
||||
self._async_enabled = False
|
||||
self._capture_thread = self._flow_thread = self._compute_thread = None
|
||||
self._capture_queue = self._flow_queue = self._result_queue = self._stop_event = None
|
||||
if self.verbose:
|
||||
print("Async pipeline disabled")
|
||||
|
||||
# =========================================================================
|
||||
# Public API - Configuration & Calibration
|
||||
# =========================================================================
|
||||
|
||||
def get_config(self):
|
||||
"""Return current configuration dict."""
|
||||
return self.config
|
||||
|
||||
def stitch_cal(self):
|
||||
"""
|
||||
Load or create stitch calibration.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Stitched calibration image
|
||||
"""
|
||||
if self.cal_path is None:
|
||||
self.cal_path = "./config/cal.npy"
|
||||
|
||||
if os.path.exists(self.cal_path):
|
||||
print("Loading calibration from", self.cal_path)
|
||||
self.stitch_calibration = np.load(self.cal_path, allow_pickle=True)
|
||||
self._stitcher.stitch_calibration = self.stitch_calibration
|
||||
return self._stitcher.stitch_rt_optimized()
|
||||
else:
|
||||
print("Creating new calibration...")
|
||||
img, self.stitch_calibration = self._stitcher.stitch_cal()
|
||||
os.makedirs("./config", exist_ok=True)
|
||||
np.save(self.cal_path, self.stitch_calibration)
|
||||
print("Calibration saved to", self.cal_path)
|
||||
return img
|
||||
|
||||
# =========================================================================
|
||||
# Private - Initialization Helpers
|
||||
# =========================================================================
|
||||
|
||||
def _load_config(self):
|
||||
"""Load sensor configuration from JSON file."""
|
||||
self.config = configs.load_config(self.config_name)
|
||||
cam = self.config["camera"]
|
||||
self.config_resolution_raw = cam["resolution_raw"]
|
||||
self.config_fps = cam["fps"]
|
||||
self.config_units = cam["units"]
|
||||
self.config_resolution_stitch = cam["resolution_stitch"]
|
||||
self.config_index = cam["index"]
|
||||
self.config_pos_default = cam["pos_default"]
|
||||
|
||||
if self.verbose:
|
||||
print(f"Config: {self.config_name}")
|
||||
print(f" Resolution: {self.config_resolution_raw} → {self.config_resolution_stitch}")
|
||||
print(f" FPS: {self.config_fps}, Units: {self.config_units}")
|
||||
|
||||
def _open_video_source(self, override, box):
|
||||
"""Initialize video capture and calibration."""
|
||||
self.cap = cv2.VideoCapture(self.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])
|
||||
|
||||
if self.verbose:
|
||||
print("Connecting to video source...")
|
||||
|
||||
if not self.cap.isOpened():
|
||||
raise Exception("Failed to open video source.")
|
||||
|
||||
# Wait for first valid frame
|
||||
frame = None
|
||||
while frame is None:
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
time.sleep(0.1)
|
||||
|
||||
if override:
|
||||
if box is not None:
|
||||
self._stitcher.isstitch = self.isstitch = False
|
||||
self._stitcher.boxes = box
|
||||
else:
|
||||
print("Override=True but no box provided.")
|
||||
else:
|
||||
self._find_stitch_boxes(frame)
|
||||
self._extract_tiles(frame)
|
||||
self.stitch_cal()
|
||||
|
||||
def _find_stitch_boxes(self, img_src, show=False):
|
||||
"""Detect vision tile boxes with camera stabilization."""
|
||||
n_boxes, v_mean0, timeout = 0, 0, 5
|
||||
start_time = cv2.getTickCount()
|
||||
|
||||
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 finding vision units.")
|
||||
self._stitcher.isstitch = self.isstitch = False
|
||||
break
|
||||
v_mean0 = v_mean
|
||||
|
||||
self.n_boxes = n_boxes
|
||||
self.stitch_valid = self.isstitch and n_boxes == self.config_units
|
||||
|
||||
status = "valid for stitching" if self.stitch_valid else "non-stitch mode"
|
||||
print(f"Sensor: {status}")
|
||||
return n_boxes
|
||||
|
||||
def _extract_tiles(self, img_src):
|
||||
"""Extract tiles from source image."""
|
||||
self._stitcher.get_stitch_unit(img_src)
|
||||
self._stitcher.extract_tiles()
|
||||
return True
|
||||
|
||||
# =========================================================================
|
||||
# Private - Async Workers
|
||||
# =========================================================================
|
||||
|
||||
def _queue_put(self, q, item):
|
||||
"""Put item in queue, dropping old if full."""
|
||||
if q.full():
|
||||
try: q.get_nowait()
|
||||
except Empty: pass
|
||||
q.put(item)
|
||||
|
||||
def _capture_worker(self):
|
||||
"""Stage 1: Capture and preprocess frames."""
|
||||
alpha = 0.7
|
||||
@@ -344,12 +498,11 @@ class Sensor:
|
||||
continue
|
||||
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
self._Sensor__stitcher.get_stitch_unit(gray)
|
||||
self._Sensor__get_tiles(gray)
|
||||
img = np.ascontiguousarray(self._Sensor__stitcher.stitch_rt_optimized())
|
||||
self._stitcher.get_stitch_unit(gray)
|
||||
self._extract_tiles(gray)
|
||||
img = np.ascontiguousarray(self._stitcher.stitch_rt_optimized())
|
||||
img = cv2.resize(img, tuple(self.config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
# Temporal filtering
|
||||
if self.img_filtered is None:
|
||||
self.img_filtered = img.astype(np.float32)
|
||||
else:
|
||||
@@ -358,7 +511,7 @@ class Sensor:
|
||||
self._queue_put(self._capture_queue, self.img_filtered.astype(np.uint8))
|
||||
except Exception as e:
|
||||
if self.verbose: print(f"Capture error: {e}")
|
||||
|
||||
|
||||
def _flow_worker(self):
|
||||
"""Stage 2: Compute optical flow."""
|
||||
while not self._stop_event.is_set():
|
||||
@@ -367,15 +520,14 @@ class Sensor:
|
||||
flow = self.tracker.track(img)
|
||||
if hasattr(flow, 'download'):
|
||||
flow = flow.download()
|
||||
|
||||
self._queue_put(self._flow_queue, (img, flow))
|
||||
except Empty:
|
||||
continue
|
||||
except Exception as e:
|
||||
if self.verbose: print(f"Flow error: {e}")
|
||||
|
||||
|
||||
def _hhd_worker(self):
|
||||
"""Stage 3: HHD decomposition and force computation."""
|
||||
"""Stage 3: HHD decomposition."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
img, flow = self._flow_queue.get(timeout=0.1)
|
||||
@@ -392,213 +544,23 @@ class Sensor:
|
||||
'img': img, 'flow': flow,
|
||||
'normal': self.hhd.get_normal(), 'shear': self.hhd.get_shear(),
|
||||
'depth': getFilteredDensity(flow, use_cuda=self.cuda),
|
||||
'fnormal': forces[0], 'fshearx': forces[1],
|
||||
'fnormal': forces[0], 'fshearx': forces[1],
|
||||
'fsheary': forces[2], 'ftor': forces[3],
|
||||
})
|
||||
except Empty:
|
||||
continue
|
||||
except Exception as e:
|
||||
if self.verbose: print(f"HHD error: {e}")
|
||||
|
||||
def get_async(self, timeout=0.05):
|
||||
"""Get latest result (non-blocking). Returns dict or None."""
|
||||
if not self._async_enabled:
|
||||
return None
|
||||
try:
|
||||
result = self._result_queue.get(timeout=timeout)
|
||||
t = time.perf_counter()
|
||||
result['fps'] = 1.0 / (t - self.last_log_time) if self.last_log_time else 0
|
||||
self.last_log_time = t
|
||||
|
||||
# Update instance vars
|
||||
for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
setattr(self, k, result.get(k))
|
||||
self.vupdate = True
|
||||
return result
|
||||
except Empty:
|
||||
return None
|
||||
|
||||
def disable_async(self):
|
||||
"""Stop async pipeline."""
|
||||
if not self._async_enabled:
|
||||
return
|
||||
self._stop_event.set()
|
||||
for t in [self._capture_thread, self._flow_thread, self._compute_thread]:
|
||||
if t: t.join(timeout=2.0)
|
||||
self._async_enabled = False
|
||||
self._capture_thread = self._flow_thread = self._compute_thread = None
|
||||
self._capture_queue = self._flow_queue = self._result_queue = self._stop_event = None
|
||||
if self.verbose:
|
||||
print("Async pipeline disabled")
|
||||
|
||||
def update(self):
|
||||
"""Convenience: get_img + compute_deformation in one call."""
|
||||
self.get_img()
|
||||
self.compute_deformation()
|
||||
return {
|
||||
'img': self.img,
|
||||
'flow': self.flow,
|
||||
'normal': self.normal,
|
||||
'shear': self.shear,
|
||||
'depth': self.depth,
|
||||
'fps': self.fps,
|
||||
'fnormal': self.fnormal,
|
||||
'fshearx': self.fshearx,
|
||||
'fsheary': self.fsheary,
|
||||
'ftor': self.ftor,
|
||||
}
|
||||
|
||||
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 and stop async threads.
|
||||
"""
|
||||
# Stop async pipeline if running
|
||||
if self._async_enabled:
|
||||
self.disable_async()
|
||||
|
||||
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, 5
|
||||
start_time = cv2.getTickCount()
|
||||
|
||||
# Wait 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.__stitcher.isstitch = self.isstitch = False
|
||||
break
|
||||
v_mean0 = v_mean
|
||||
|
||||
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
|
||||
"""
|
||||
if self.cal_path is None:
|
||||
self.cal_path = "./config/cal.npy"
|
||||
cal_path = self.cal_path
|
||||
|
||||
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)
|
||||
|
||||
return img_stitched_height
|
||||
# =========================================================================
|
||||
# Private - Utilities
|
||||
# =========================================================================
|
||||
|
||||
def _print_timings(self, name, timings):
|
||||
"""Print timing breakdown."""
|
||||
print(f"\n[{name} timing]")
|
||||
for k, v in timings.items():
|
||||
if k != 'total':
|
||||
print(f" {k}: {v:7.2f} ms")
|
||||
print(f" ─────────────────────────")
|
||||
print(f" TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")
|
||||
|
||||
Reference in New Issue
Block a user