add async to speed up 2x

This commit is contained in:
bluepulsar
2025-12-27 15:05:07 +08:00
parent 3f4f42ee4a
commit a8b02105fc
3 changed files with 230 additions and 67 deletions
+170 -1
View File
@@ -10,6 +10,9 @@ from enum import Enum, auto
from typing import Optional, Tuple, Union, Any
from .getDensity import getFilteredDensity
from . import configs
from threading import Thread, Event
from queue import Queue, Empty
from concurrent.futures import ThreadPoolExecutor
class Sensor:
"""
@@ -60,6 +63,16 @@ class Sensor:
# Logging state
self.last_log_time = time.perf_counter()
self.fps = 0
# Async pipeline state (3-stage: Capture → Flow → HHD)
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
if self.verbose:
print(f" Video source: {vid_src}")
@@ -282,6 +295,158 @@ class Sensor:
self._last_timings = timings
return self.flow
# =========================================================================
# Async/Pipelined Processing API (3-stage: Capture → Flow → HHD)
# =========================================================================
def _queue_put(self, q, item):
"""Helper to put item in queue, dropping old if full."""
if q.full():
try: q.get_nowait()
except Empty: pass
q.put(item)
def enable_async(self, queue_size=2):
"""
Enable 3-stage async pipeline for maximum throughput.
Pipeline: [Capture] → [Optical Flow] → [HHD Compute]
Each stage runs in parallel, overlapping all operations.
"""
if self._async_enabled:
return
self._capture_queue = Queue(maxsize=queue_size) # img
self._flow_queue = Queue(maxsize=queue_size) # (img, flow)
self._result_queue = Queue(maxsize=queue_size) # final result
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)
self._capture_thread.start()
self._flow_thread.start()
self._compute_thread.start()
if self.verbose:
print("Async 3-stage pipeline enabled: Capture → Flow → HHD")
def _capture_worker(self):
"""Stage 1: Capture and preprocess frames."""
alpha = 0.7
while not self._stop_event.is_set():
try:
ret, frame = self.cap.read()
if not ret:
continue
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
self._Sensor__stitcher.get_stitch_unit(gray)
self._Sensor__get_tiles(gray)
img = np.ascontiguousarray(self._Sensor__stitcher.stitch_rt_optimized())
img = cv2.resize(img, tuple(self.config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
# Temporal filtering
if self.img_filtered is None:
self.img_filtered = img.astype(np.float32)
else:
self.img_filtered = alpha * img + (1 - alpha) * self.img_filtered
self._queue_put(self._capture_queue, self.img_filtered.astype(np.uint8))
except Exception as e:
if self.verbose: print(f"Capture error: {e}")
def _flow_worker(self):
"""Stage 2: Compute optical flow."""
while not self._stop_event.is_set():
try:
img = self._capture_queue.get(timeout=0.1)
flow = self.tracker.track(img)
if hasattr(flow, 'download'):
flow = flow.download()
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 and force computation."""
while not self._stop_event.is_set():
try:
img, flow = self._flow_queue.get(timeout=0.1)
if self.hhd is None:
self.hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))
if self.cuda and self.CUPY_AVAILABLE:
_, forces = self.hhd.get_force_gpu_v2(flow)
else:
_, forces = self.hhd.get_force(flow)
self._queue_put(self._result_queue, {
'img': img, 'flow': flow,
'normal': self.hhd.get_normal(), 'shear': self.hhd.get_shear(),
'depth': getFilteredDensity(flow, use_cuda=self.cuda),
'fnormal': forces[0], 'fshearx': forces[1],
'fsheary': forces[2], 'ftor': forces[3],
})
except Empty:
continue
except Exception as e:
if self.verbose: print(f"HHD error: {e}")
def get_async(self, timeout=0.05):
"""Get latest result (non-blocking). Returns dict or None."""
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
# Update instance vars
for k in ['img', 'flow', 'normal', 'shear', 'depth', '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")
def update(self):
"""Convenience: get_img + compute_deformation in one call."""
self.get_img()
self.compute_deformation()
return {
'img': self.img,
'flow': self.flow,
'normal': self.normal,
'shear': self.shear,
'depth': self.depth,
'fps': self.fps,
'fnormal': self.fnormal,
'fshearx': self.fshearx,
'fsheary': self.fsheary,
'ftor': self.ftor,
}
def read_info(self, *args: info) -> Tuple:
"""
@@ -341,8 +506,12 @@ class Sensor:
def disconnect(self):
"""
Release video capture resources.
Release video capture resources and stop async threads.
"""
# Stop async pipeline if running
if self._async_enabled:
self.disable_async()
if self.cap.isOpened():
self.cap.release()
if self.verbose: