add multi-sensor support

This commit is contained in:
bluepulsar
2025-12-27 13:42:47 +08:00
parent 3f231c2c42
commit 3f4f42ee4a
7 changed files with 418 additions and 42 deletions
+269
View File
@@ -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()
+2 -1
View File
@@ -9,6 +9,7 @@ __version__ = "0.1.0"
__author__ = "Orisys"
from .sensor import Sensor
from .MultiSensor import MultiSensor
from . import util
__all__ = ['Sensor', 'util']
__all__ = ['Sensor', 'MultiSensor', 'util']
+45 -19
View File
@@ -16,7 +16,7 @@ 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", override=False, box=None):
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.
@@ -38,7 +38,7 @@ class Sensor:
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 = vid_src, isstitch, cuda, verbose, config_name
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
@@ -58,7 +58,7 @@ class Sensor:
self.__open_video_source(override, box)
# Logging state
self.last_log_time = time.time()
self.last_log_time = time.perf_counter()
self.fps = 0
if self.verbose:
@@ -102,9 +102,7 @@ class Sensor:
if override:
if box is not None:
self.isstitch = False
self.__stitcher.isstitch = False
self.__stitcher.isstitch = self.isstitch = False
self.__stitcher.boxes = box
else:
print("Override is True but no box provided.")
@@ -137,7 +135,7 @@ class Sensor:
print(" Resolution after stitch: ", str(self.config_resolution_stitch))
print(" Index for stitching: ", str(self.config_index))
def get_img(self):
def get_img(self, log_timing=False):
"""
Capture frame from video source, convert to grayscale, and perform stitching.
@@ -147,14 +145,32 @@ class Sensor:
ret, frame = self.cap.read()
if not ret:
print("Failed to get_img from video source.")
self.cap.release()
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
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
# 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
img_resized = cv2.resize(img, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
@@ -171,6 +187,18 @@ class Sensor:
# 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)")
# Store timings for external access
self._last_timings = timings
return self.img
def compute_deformation(self, log_timing=False):
@@ -199,8 +227,6 @@ class Sensor:
if log_timing:
timings['optical_flow'] = (time.perf_counter() - t0) * 1000
if hasattr(flow, 'download'):
flow = flow.download()
@@ -239,8 +265,8 @@ class Sensor:
timings['get_density'] = (time.perf_counter() - t0) * 1000
# Calculate FPS
self.fps = 1/(time.time() - self.last_log_time)
self.last_log_time = time.time()
self.fps = 1/((time.perf_counter() - self.last_log_time))
self.last_log_time = time.perf_counter()
# Print timing summary
if log_timing:
@@ -351,7 +377,7 @@ class Sensor:
n_boxes, v_mean0, timeout = 0, 0, 5
start_time = cv2.getTickCount()
# print("Waiting for camera stabilization...")
# Wait for camera stabilization
while n_boxes < self.config_units:
ret, frame = self.cap.read()
v_mean = np.sum(cv2.mean(frame))
@@ -361,8 +387,7 @@ class Sensor:
if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout:
print("Timeout for finding valid vision units.")
self.isstitch = False
self.__stitcher.isstitch = False
self.__stitcher.isstitch = self.isstitch = False
break
v_mean0 = v_mean
@@ -390,13 +415,16 @@ class Sensor:
- 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 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()
@@ -404,6 +432,4 @@ class Sensor:
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
-20
View File
@@ -87,23 +87,6 @@ def draw_magnitude_map(image, flow_field, threshold=3.0, colormap=cv2.COLORMAP_J
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)
@@ -128,9 +111,6 @@ def draw_magnitude_map(image, flow_field, threshold=3.0, colormap=cv2.COLORMAP_J
# 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 magnitude_colored
def draw_fps(image, fps):