Files
orisys-sdk/src/orisys/sensor.py
T
2026-01-14 15:22:31 +08:00

689 lines
26 KiB
Python

"""
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
import time
from enum import Enum, auto
from typing import Tuple
from threading import Thread, Event
from queue import Queue, Empty
from .Stitcher import Stitcher
from .tracker import Tracker
from .HHD import nHHD
from . import configs
from .get_depth import Contact3D
class Sensor:
"""
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.
"""
# =========================================================================
# Info Enum (for read_info API)
# =========================================================================
class info(Enum):
"""Data types available through read_info()."""
VRAW = auto()
VNORMAL = auto()
VSHEAR = auto()
FPS = auto()
TIMESTAMP = auto()
FNORMAL = auto()
FSHEARX = auto()
FSHEARY = auto()
CENTROID = auto()
CONTOUR = auto()
DEPTH = auto()
IMG = auto()
# =========================================================================
# Initialization
# =========================================================================
def __init__(self, vid_src, cuda=True, isstitch=True, verbose=True,
config_name="prototype", cal_path=None, override=False, box=None, mirror=True):
"""
Initialize Sensor with video source and configuration.
Args:
vid_src: Video source (int for camera ID, str for file path)
cuda: Enable CUDA acceleration (default: True)
isstitch: Enable stitching mode for multi-tile sensor (default: False)
verbose: Print initialization logs (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
self.mirror = mirror
# Load configuration
self.__config_cam = None
self.__config_stitcher = None
self.__config_contact = None
self._load_config()
# State variables
self.stitch_valid = False
self.stitch_calibration = []
self.compute_deformation_valid = False
self.compute_contact_valid = False
# Data Variables (public)
shape = (self.__config_cam["resolution_stitch"][1], self.__config_cam["resolution_stitch"][0])
self.img = None
self.img_size_raw = None
self.img_filtered = None
self.initial_img = None
self.flow = np.zeros(shape + (2,), dtype=np.float32)
self.vnormal = np.zeros_like(self.flow)
self.vshear = np.zeros_like(self.flow)
self.n_boxes = None
self.fnormal = self.fshearx = self.fsheary = self.ftor = 0.0
# Contact Data Variables (public)
self.centroid = (0, 0)
self.centroid_found = False
self.contour = np.empty(0)
self.depth_map = np.zeros_like(self.img)
# 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
# Initialize components
self.stitcher = Stitcher(isstitch=isstitch, config_stitcher=self.__config_stitcher, verbose=self.verbose)
self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker)
self.__contact3d = Contact3D(config_contact=self.__config_contact, verbose=self.verbose)
self.__hhd = None
# Setup video capture
self._open_video_source(override, box)
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()}")
# =========================================================================
# Public API - Core Methods (Most Frequently Used)
# =========================================================================
def get_img(self):
"""
Capture and preprocess a frame.
Returns:
np.ndarray: Processed grayscale image
"""
ret, frame = self.cap.read()
if not ret:
print("Failed to get_img from video source.")
return None
# Convert to grayscale
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# Extract and stitch tiles
self.stitcher.get_stitch_unit(gray)
self._extract_tiles(gray)
img = np.ascontiguousarray(self.stitcher.stitch_rt_optimized())
self.img_size_raw = img.shape
# Resize
img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
if self.mirror:
img = cv2.flip(img, 1)
self.img = img.astype(np.uint8)
# FPS calculation
self._fps = 1.0 / (time.perf_counter() - self.last_log_time)
self.last_log_time = time.perf_counter()
return self.img.copy()
def detect_motion(self, threshold=4):
if self.initial_img is None:
return False
diff = cv2.absdiff(self.img, self.initial_img)
return diff.mean() > threshold
def compute_deformation(self):
"""
Compute optical flow and decompose into force components.
Returns:
np.ndarray: Optical flow field
"""
if self.img is None:
raise Exception("Warning: Get image first")
self.__compute_flow()
self.__compute_hhd()
return True
def __compute_flow(self):
if self.img is None:
raise Exception("Warning: Get image first")
flow = self.__tracker.track(self.img)
if hasattr(flow, 'download'):
flow = flow.download()
self.flow = flow
def __compute_hhd(self):
flow = self.flow
if flow is None:
raise Exception("Warning: Compute flow first")
# Initialize HHD solver
if self.__hhd is None:
self.__hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))
# Force computation
if self.cuda and self.CUPY_AVAILABLE:
_, forces = self.__hhd.get_force_gpu_v2(flow)
else:
_, forces = self.__hhd.get_force(flow)
self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
# Get components
self.vnormal = self.__hhd.get_normal()
self.vshear = self.__hhd.get_shear()
# Update state
self.compute_deformation_valid = True
return True
def compute_contact(self, vnormal=None):
"""
Compute contact information from vnormal field.
Args:
vnormal: Optional vnormal field. If None, uses self.vnormal.
Returns:
dict: Contact result with centroid, contour, and depth_map
"""
multiple_foe = False
# Use provided vnormal or fall back to internal state
vnormal_to_use = vnormal if vnormal is not None else self.vnormal
if vnormal_to_use is None:
raise Exception("Warning: vnormal is required for contact computation")
if multiple_foe:
result = self.__contact3d.compute_multiple_foe(vnormal_to_use)
else:
if self.cuda:
result = self.__contact3d.compute_foe_gpu(vnormal_to_use)
else:
result = self.__contact3d.compute_foe(vnormal_to_use)
self.centroid = result['centroid']
self.centroid_found = result['centroid_found']
if 'contour' in result:
self.contour = result['contour']
self.depth_map = result['divergence_map']
self.compute_contact_valid = True
return result
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 any(arg in [self.info.IMG] for arg in args):
if self.img is None:
raise Exception("Warning: Get image first")
if any(arg in [self.info.VRAW, self.info.VNORMAL, self.info.VSHEAR, self.info.FNORMAL, self.info.FSHEARX, self.info.FSHEARY] for arg in args):
if not self.compute_deformation_valid:
print("Warning: Compute deformation first")
# raise Exception("Warning: Compute deformation first")
if any(arg in [self.info.CENTROID, self.info.CONTOUR, self.info.DEPTH] for arg in args):
if not self.compute_contact_valid:
print("Warning: Compute contact first")
# raise Exception("Warning: Compute contact first")
mapping = {
# Return copies of arrays to prevent external modification of internal state
self.info.VRAW: lambda: self.flow.copy() if self.flow is not None else np.array([]),
self.info.VNORMAL: lambda: self.vnormal.copy() if self.vnormal is not None else np.array([]),
self.info.VSHEAR: lambda: self.vshear.copy() if self.vshear is not None else np.array([]),
# Scalar values are safe to return directly
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'),
# Return copies of arrays to prevent external modification
self.info.DEPTH: lambda: self.depth_map.copy() if self.depth_map is not None else np.array([]),
# Centroid: can be tuple, list, or array - return a copy to prevent modification
# Return None if centroid is not found or empty
self.info.CENTROID: lambda: (
np.array(self.centroid).copy()
if self.centroid is not None and len(self.centroid) > 0
else None
),
self.info.CONTOUR: lambda: self.contour.copy() if self.contour is not None else np.array([]),
self.info.IMG: lambda: self.img.copy() if self.img 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()
self.initial_img = self.img.copy()
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()
# =========================================================================
# Public API - Configuration & Calibration
# =========================================================================
# def export_config(self, path="./config/exported_config.json"):
# """Return current configuration dict."""
# configs.export_config(self._config, path)
# if self.verbose:
# print(f"Config exported to {path}")
# return True
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):
if self.verbose:
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:
if self.verbose:
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)
if self.verbose:
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)
self.__config_cam = self._config["camera"]
self._config_resolution_raw = self.__config_cam["resolution_raw"]
self._config_fps = self.__config_cam["fps"]
self._config_units = self.__config_cam["units"]
self._config_resolution_stitch = self.__config_cam["resolution_stitch"]
self._config_dis_preset = self._config["tracker"]["dis_preset"]
self.__config_tracker = self._config["tracker"]
self.__config_stitcher = self._config["stitcher"]
self.__config_contact = self._config["contact"]
if self.verbose:
print(f"Config: {self.config_name}")
return True
def _open_video_source(self, override, box):
"""Initialize video capture and calibration."""
self.cap = cv2.VideoCapture(self.vid_src)
# Set buffer size to 1 to reduce latency and avoid frame accumulation
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
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)
img = self.stitch_cal()
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
self.initial_img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
if self.mirror:
self.initial_img = cv2.flip(self.initial_img, 1)
self.initial_img = self.initial_img.astype(np.uint8)
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, 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.boxes = 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 - 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)")
# =========================================================================
# Public API - Async Pipeline
# =========================================================================
# def enable_async(self, queue_size=2, compute_contact=False):
# """
# Enable 3-stage async pipeline: Capture → Flow → HHD.
# Each stage runs in parallel for maximum throughput.
# """
# if self._async_enabled:
# return
# self._compute_contact_enabled = compute_contact
# self._capture_queue = Queue(maxsize=queue_size)
# self._flow_queue = Queue(maxsize=queue_size)
# self._result_queue = Queue(maxsize=queue_size)
# if compute_contact:
# self._contact_queue = Queue(maxsize=queue_size)
# self._stop_event = Event()
# self._async_enabled = True
# self._capture_thread = Thread(target=self._capture_worker, daemon=True)
# self._flow_thread = Thread(target=self._flow_worker, daemon=True)
# self._compute_thread = Thread(target=self._hhd_worker, daemon=True)
# if compute_contact:
# self._contact_thread = Thread(target=self._contact_worker, daemon=True)
# self._capture_thread.start()
# self._flow_thread.start()
# self._compute_thread.start()
# if compute_contact:
# self._contact_thread.start()
# if self.verbose:
# if compute_contact:
# print("Async pipeline enabled: Capture → Flow → HHD → Contact")
# else:
# 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
# """
# compute_contact = self._compute_contact_enabled
# 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']:
# if compute_contact:
# for k in ['img', 'flow', 'normal', 'shear', 'depth', 'centroid', 'centroid_found', 'contour', 'fps', 'fnormal', 'fshearx', 'fsheary', 'ftor']:
# setattr(self, k, result.get(k))
# else:
# for k in ['img', 'flow', 'normal', 'shear', '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")
# =========================================================================
# 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 by reusing get_img()."""
# while not self._stop_event.is_set():
# try:
# # Reuse get_img() to avoid code duplication
# if self.get_img() is None:
# continue
# # Put the processed image in the queue
# # get_img() updates self.img, so we use that
# if self.img is not None:
# self._queue_put(self._capture_queue, self.img.copy())
# except Exception as e:
# if self.verbose: print(f"Capture error: {e}")
# def _flow_worker(self):
# """Stage 2: Compute optical flow by reusing __compute_flow()."""
# while not self._stop_event.is_set():
# try:
# img = self._capture_queue.get(timeout=0.1)
# # Set __img to reuse __compute_flow()
# self.img = img
# self.__compute_flow()
# flow = self.flow.copy() if self.flow is not None else None
# 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 by reusing __compute_hhd()."""
# while not self._stop_event.is_set():
# try:
# img, flow = self._flow_queue.get(timeout=0.1)
# # Set __flow to reuse __compute_hhd()
# self.flow = flow
# self.__compute_hhd()
# # Get results from internal variables set by __compute_hhd()
# result = {
# 'img': img, 'flow': flow,
# 'normal': self.vnormal.copy() if self.vnormal is not None else None,
# 'shear': self.vshear.copy() if self.vshear is not None else None,
# # 'depth': getFilteredDensity(flow, use_cuda=self.cuda),
# 'fnormal': self.fnormal, 'fshearx': self.fshearx,
# 'fsheary': self.fsheary, 'ftor': self.ftor,
# }
# # Always include contact fields to ensure consistent result structure
# # When compute_contact is True, contact worker will update these values
# # When compute_contact is False, they remain None
# result.update({
# 'centroid': None,
# 'centroid_found': False,
# 'contour': None,
# 'depth_map': None,
# 'depth': None, # Alias for depth_map for compatibility
# })
# self._queue_put(self._result_queue, result)
# except Empty:
# continue
# except Exception as e:
# if self.verbose: print(f"HHD error: {e}")
# def _contact_worker(self):
# """Stage 4: Contact computation using vnormal from HHD results."""
# while not self._stop_event.is_set():
# try:
# # Get results from HHD worker which includes vnormal
# hhd_result = self._result_queue.get(timeout=0.1)
# vnormal = hhd_result.get('normal') # This is vnormal from HHD
# if vnormal is not None:
# # Compute contact using vnormal directly
# contact_result = self.compute_contact(vnormal=vnormal)
# # Add contact results to the existing HHD result
# depth_map = contact_result.get('divergence_map')
# hhd_result.update({
# 'centroid': contact_result.get('centroid'),
# 'centroid_found': contact_result.get('centroid_found'),
# 'contour': contact_result.get('contour'),
# 'depth_map': depth_map,
# 'depth': depth_map, # Alias for depth_map for compatibility
# })
# # Put updated result back to queue
# self._queue_put(self._result_queue, hhd_result)
# except Empty:
# continue
# except Exception as e:
# if self.verbose: print(f"Contact error: {e}")