318 lines
11 KiB
Python
318 lines
11 KiB
Python
"""
|
|
MultiSensor - Multi-process sensor management for parallel processing.
|
|
|
|
This module provides the MultiSensor class for running multiple Sensor
|
|
instances in parallel processes, enabling true parallel processing with
|
|
independent CUDA contexts.
|
|
"""
|
|
|
|
import os
|
|
|
|
|
|
def _sensor_worker(sensor_id, config, data_queue, stop_event, use_async=False):
|
|
"""
|
|
Worker function that runs in a separate PROCESS for each sensor.
|
|
Each process has its own Python interpreter and CUDA context.
|
|
|
|
Args:
|
|
use_async: If True, use 3-stage async pipeline within the process
|
|
"""
|
|
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
|
|
|
from .sensor import Sensor
|
|
|
|
print(f"[Sensor {sensor_id}] Process started (PID: {os.getpid()}), async={use_async}")
|
|
|
|
sensor = Sensor(
|
|
vid_src=config.get('vid_src'),
|
|
isstitch=config.get('isstitch', True),
|
|
cuda=config.get('cuda', True),
|
|
verbose=config.get('verbose', False),
|
|
config_name=config.get('config_name', 'prototype'),
|
|
cal_path=config.get('cal_path', None)
|
|
)
|
|
|
|
# Enable async pipeline if requested
|
|
if use_async:
|
|
sensor.enable_async(queue_size=2)
|
|
|
|
frame_count = 0
|
|
|
|
def queue_put(data):
|
|
try:
|
|
while not data_queue.empty():
|
|
try: data_queue.get_nowait()
|
|
except: break
|
|
data_queue.put_nowait(data)
|
|
except: pass
|
|
|
|
while not stop_event.is_set():
|
|
try:
|
|
if use_async:
|
|
# Use async get
|
|
result = sensor.get_async(timeout=0.05)
|
|
if result is None:
|
|
continue
|
|
data = {
|
|
'sensor_id': sensor_id,
|
|
'img': result['img'].copy() if result['img'] is not None else None,
|
|
'flow': result['flow'].copy() if result['flow'] is not None else None,
|
|
'normal': result['normal'].copy() if result['normal'] is not None else None,
|
|
'shear': result['shear'].copy() if result['shear'] is not None else None,
|
|
'depth': result['depth'].copy() if result['depth'] is not None else None,
|
|
'fps': result['fps'],
|
|
'fnormal': result['fnormal'],
|
|
'fshearx': result['fshearx'],
|
|
'fsheary': result['fsheary'],
|
|
'ftor': result['ftor'],
|
|
'frame_count': frame_count,
|
|
}
|
|
else:
|
|
# Sync mode
|
|
sensor.get_img()
|
|
sensor.compute_deformation()
|
|
data = {
|
|
'sensor_id': sensor_id,
|
|
'img': sensor.img.copy() if sensor.img is not None else None,
|
|
'flow': sensor.flow.copy() if sensor.flow is not None else None,
|
|
'normal': sensor.vnormal.copy() if sensor.vnormal is not None else None,
|
|
'shear': sensor.shear.copy() if sensor.shear is not None else None,
|
|
'depth': sensor.depth.copy() if sensor.depth is not None else None,
|
|
'fps': sensor.fps,
|
|
'fnormal': sensor.fnormal,
|
|
'fshearx': sensor.fshearx,
|
|
'fsheary': sensor.fsheary,
|
|
'ftor': sensor.ftor,
|
|
'frame_count': frame_count,
|
|
}
|
|
|
|
queue_put(data)
|
|
frame_count += 1
|
|
|
|
except Exception as e:
|
|
print(f"[Sensor {sensor_id}] Error: {e}")
|
|
break
|
|
|
|
sensor.disconnect()
|
|
print(f"[Sensor {sensor_id}] Process stopped")
|
|
|
|
|
|
class MultiSensor:
|
|
"""
|
|
Manages multiple Sensor instances running in parallel processes.
|
|
|
|
Each sensor runs in its own process with independent CUDA context,
|
|
enabling true parallel processing without GIL limitations.
|
|
|
|
Example:
|
|
# Simple usage with context manager
|
|
configs = [
|
|
{'vid_src': './s1.mp4', 'isstitch': True},
|
|
{'vid_src': './s2.mp4', 'isstitch': True},
|
|
]
|
|
|
|
with MultiSensor(configs) as multi:
|
|
for all_data in multi:
|
|
for data in all_data:
|
|
if data:
|
|
print(f"Sensor {data['sensor_id']}: FPS={data['fps']}")
|
|
cv2.imshow(f"Sensor {data['sensor_id']}", data['img'])
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
# Or manual control
|
|
multi = MultiSensor(configs)
|
|
multi.start()
|
|
|
|
data = multi.get(0) # Get from sensor 0
|
|
all_data = multi.get_all() # Get from all sensors
|
|
|
|
multi.stop()
|
|
"""
|
|
|
|
def __init__(self, sensor_configs):
|
|
"""
|
|
Initialize MultiSensor manager.
|
|
|
|
Args:
|
|
sensor_configs: List of dicts, each containing Sensor constructor args:
|
|
- vid_src: Video source (required) - int for camera ID, str for file path
|
|
- isstitch: Whether to stitch (default: True)
|
|
- cuda: Use CUDA acceleration (default: True)
|
|
- verbose: Print logs (default: False)
|
|
- config_name: Config name (default: 'prototype')
|
|
- cal_path: Calibration file path (default: None)
|
|
|
|
Example:
|
|
configs = [
|
|
{'vid_src': 0, 'isstitch': True, 'cuda': True},
|
|
{'vid_src': 1, 'isstitch': True, 'cuda': True},
|
|
]
|
|
multi = MultiSensor(configs)
|
|
"""
|
|
import multiprocessing as mp
|
|
|
|
self.sensor_configs = sensor_configs
|
|
self.num_sensors = len(sensor_configs)
|
|
self._mp = mp
|
|
self._use_async = False
|
|
|
|
self.processes = []
|
|
self.data_queues = []
|
|
self.stop_events = []
|
|
self._started = False
|
|
|
|
# Create queues and events for each sensor
|
|
for i in range(self.num_sensors):
|
|
self.data_queues.append(mp.Queue(maxsize=2))
|
|
self.stop_events.append(mp.Event())
|
|
|
|
def start(self, use_async=False):
|
|
"""
|
|
Start all sensor worker processes.
|
|
|
|
Args:
|
|
use_async: If True, each sensor process uses 3-stage async pipeline
|
|
for maximum throughput (Capture → Flow → HHD in parallel)
|
|
"""
|
|
if self._started:
|
|
print("MultiSensor already started")
|
|
return
|
|
|
|
self._use_async = use_async
|
|
|
|
for i, config in enumerate(self.sensor_configs):
|
|
process = self._mp.Process(
|
|
target=_sensor_worker,
|
|
args=(i, config, self.data_queues[i], self.stop_events[i], use_async),
|
|
daemon=True
|
|
)
|
|
process.start()
|
|
self.processes.append(process)
|
|
print(f"Started sensor {i} (PID: {process.pid})")
|
|
|
|
mode = "async (3-stage pipeline)" if use_async else "sync"
|
|
self._started = True
|
|
print(f"All {len(self.processes)} sensors started in {mode} mode")
|
|
|
|
def get(self, sensor_idx, timeout=0.01):
|
|
"""
|
|
Get the latest data from a specific sensor.
|
|
|
|
Args:
|
|
sensor_idx: Index of the sensor (0-based)
|
|
timeout: How long to wait for data (seconds)
|
|
|
|
Returns:
|
|
dict: Sensor data or None if no data available
|
|
- sensor_id: int
|
|
- img: np.ndarray (H, W, 3)
|
|
- flow: np.ndarray (H, W, 2)
|
|
- normal: np.ndarray
|
|
- shear: np.ndarray
|
|
- depth: np.ndarray
|
|
- fps: float
|
|
- fnormal, fshearx, fsheary, ftor: float
|
|
- frame_count: int
|
|
"""
|
|
try:
|
|
return self.data_queues[sensor_idx].get(timeout=timeout)
|
|
except:
|
|
return None
|
|
|
|
def get_all(self, timeout=0.01):
|
|
"""
|
|
Get the latest data from all sensors.
|
|
|
|
Returns:
|
|
list: List of sensor data dicts (None for sensors with no new data)
|
|
"""
|
|
return [self.get(i, timeout) for i in range(self.num_sensors)]
|
|
|
|
def __iter__(self):
|
|
"""
|
|
Iterate over sensor data. Yields all sensor data on each iteration.
|
|
|
|
Example:
|
|
for all_data in multi:
|
|
for data in all_data:
|
|
if data:
|
|
process(data)
|
|
"""
|
|
return self
|
|
|
|
def __next__(self):
|
|
"""Get next batch of sensor data."""
|
|
if not self._started:
|
|
raise StopIteration("MultiSensor not started")
|
|
|
|
# Check if any process is still alive
|
|
if not any(p.is_alive() for p in self.processes):
|
|
raise StopIteration("All sensor processes have stopped")
|
|
|
|
return self.get_all()
|
|
|
|
def __len__(self):
|
|
"""Return number of sensors."""
|
|
return self.num_sensors
|
|
|
|
def __getitem__(self, idx):
|
|
"""Get latest data from sensor by index."""
|
|
return self.get(idx)
|
|
|
|
def stop(self):
|
|
"""Stop all sensor processes gracefully."""
|
|
if not self._started:
|
|
return
|
|
|
|
print("Stopping all sensors...")
|
|
|
|
# Signal all processes to stop
|
|
for event in self.stop_events:
|
|
event.set()
|
|
|
|
# Wait for processes to finish
|
|
for i, process in enumerate(self.processes):
|
|
process.join(timeout=3.0)
|
|
if process.is_alive():
|
|
print(f"Force terminating sensor {i}")
|
|
process.terminate()
|
|
|
|
self.processes.clear()
|
|
self._started = False
|
|
print("All sensors stopped.")
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry - starts sensors (sync mode by default)."""
|
|
self.start(use_async=False)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit - stops sensors."""
|
|
self.stop()
|
|
return False
|
|
|
|
def __del__(self):
|
|
"""Cleanup on deletion."""
|
|
self.stop()
|
|
|
|
@classmethod
|
|
def async_context(cls, sensor_configs):
|
|
"""
|
|
Create a MultiSensor that uses async mode when used as context manager.
|
|
|
|
Example:
|
|
with MultiSensor.async_context(configs) as multi:
|
|
for all_data in multi:
|
|
...
|
|
"""
|
|
return _AsyncMultiSensorContext(sensor_configs)
|
|
|
|
|
|
class _AsyncMultiSensorContext(MultiSensor):
|
|
"""Helper class for async context manager."""
|
|
def __enter__(self):
|
|
self.start(use_async=True)
|
|
return self
|
|
|