add multi-sensor support
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user