add multi-sensor support
This commit is contained in:
bluepulsar
@@ -0,0 +1,269 @@
|
||||
"""
|
||||
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):
|
||||
"""
|
||||
Worker function that runs in a separate PROCESS for each sensor.
|
||||
Each process has its own Python interpreter and CUDA context.
|
||||
|
||||
Note: This function must be at module level for multiprocessing to work.
|
||||
"""
|
||||
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
||||
|
||||
# Import inside process to get fresh CUDA context
|
||||
from .sensor import Sensor
|
||||
|
||||
print(f"[Sensor {sensor_id}] Process started (PID: {os.getpid()})")
|
||||
|
||||
# Initialize sensor in this process
|
||||
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)
|
||||
)
|
||||
|
||||
frame_count = 0
|
||||
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
# Get image and compute deformation
|
||||
sensor.get_img()
|
||||
sensor.compute_deformation()
|
||||
|
||||
# Package all data
|
||||
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.normal.copy() if sensor.normal 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,
|
||||
}
|
||||
|
||||
# Non-blocking put - drop old frames if queue is full
|
||||
try:
|
||||
while not data_queue.empty():
|
||||
try:
|
||||
data_queue.get_nowait()
|
||||
except:
|
||||
break
|
||||
data_queue.put_nowait(data)
|
||||
except:
|
||||
pass
|
||||
|
||||
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.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):
|
||||
"""Start all sensor worker processes."""
|
||||
if self._started:
|
||||
print("MultiSensor already started")
|
||||
return
|
||||
|
||||
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]),
|
||||
daemon=True
|
||||
)
|
||||
process.start()
|
||||
self.processes.append(process)
|
||||
print(f"Started sensor {i} (PID: {process.pid})")
|
||||
|
||||
self._started = True
|
||||
print(f"All {len(self.processes)} sensors started")
|
||||
|
||||
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."""
|
||||
self.start()
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user