refractor for HHD and depth backend

This commit is contained in:
bluepulsar
2026-02-18 10:56:30 +08:00
parent ef3a25a158
commit 54baf1ecde
4 changed files with 494 additions and 609 deletions
+18 -219
View File
@@ -20,6 +20,7 @@ from .tracker import Tracker
from .HHD import nHHD
from . import configs
from .get_depth import Contact3D
from .backend_availability import (check_cupy_available, check_opencl_available)
class Sensor:
"""
@@ -67,12 +68,9 @@ class Sensor:
override: Skip auto-detection, use provided box (default: False)
box: Manual bounding boxes when override=True (default: None)
"""
# Check CUDA availability
try:
import cupy as cp
self.CUPY_AVAILABLE = True
except ImportError:
self.CUPY_AVAILABLE = False
self.CUPY_AVAILABLE = check_cupy_available()
self.OPENCL_AVAILABLE = check_opencl_available()
if vid_src is None:
raise Exception("User must specify a video source.")
@@ -124,18 +122,6 @@ class Sensor:
self._last_capture_time = 0.0
self._last_processing_time = 0.0
# Async pipeline state
# 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
# Initialize components
self.stitcher = Stitcher(isstitch=isstitch, config_stitcher=self.__config_stitcher, verbose=self.verbose)
self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker)
@@ -149,9 +135,13 @@ class Sensor:
if self.verbose:
print(f" Video source: {vid_src}")
print(f" Stitch mode: {self.stitch_valid}")
print(f" CUDA acceleration: {self.cuda}")
print(f" CuPy available: {self.CUPY_AVAILABLE}")
print(f" Optical flow algorithm: {self.__config_tracker['algorithm']}")
print(f" Poisson backend: {self.poisson_backend}")
# print(f" CUDA acceleration: {self.cuda}")
print(f" CuPy available: {self.CUPY_AVAILABLE}")
print(f" OpenCL available: {self.OPENCL_AVAILABLE}")
# =========================================================================
# Public API - Core Methods (Most Frequently Used)
@@ -247,14 +237,15 @@ class Sensor:
self.__hhd = nHHD(cuda=self.cuda, grid=flow.shape[:2], spacings=(18, 18), poisson_backend=self.poisson_backend)
# Force computation
# When OpenCL backend is selected, use CPU force path to ensure
# When OpenCL backend is selected, use host force path to ensure
# Poisson solve goes through PoissonSolver.solve_opencl().
if self.poisson_backend == 'opencl':
_, forces = self.__hhd.get_force(flow)
elif self.cuda and self.CUPY_AVAILABLE:
_, forces = self.__hhd.get_force_gpu_v2(flow)
else:
_, forces = self.__hhd.get_force(flow)
_, forces = self.__hhd.get_force(flow)
# if self.poisson_backend == 'opencl':
# _, forces = self.__hhd.get_force_host(flow)
# elif self.cuda and self.CUPY_AVAILABLE:
# _, forces = self.__hhd.get_force_gpu_v2(flow)
# else:
# _, forces = self.__hhd.get_force_host(flow)
self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
# Get components
@@ -359,8 +350,6 @@ class Sensor:
def disconnect(self):
"""Release resources and stop async threads."""
# if self._async_enabled:
# self.disable_async()
if self.cap.isOpened():
self.cap.release()
if self.verbose:
@@ -521,193 +510,3 @@ class Sensor:
print(f" ─────────────────────────")
print(f" TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")
# =========================================================================
# Public API - Async Pipeline
# =========================================================================
# def enable_async(self, queue_size=2, compute_contact=False):
# """
# Enable 3-stage async pipeline: Capture → Flow → HHD.
# Each stage runs in parallel for maximum throughput.
# """
# if self._async_enabled:
# return
# self._compute_contact_enabled = compute_contact
# self._capture_queue = Queue(maxsize=queue_size)
# self._flow_queue = Queue(maxsize=queue_size)
# self._result_queue = Queue(maxsize=queue_size)
# if compute_contact:
# self._contact_queue = Queue(maxsize=queue_size)
# 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)
# if compute_contact:
# self._contact_thread = Thread(target=self._contact_worker, daemon=True)
# self._capture_thread.start()
# self._flow_thread.start()
# self._compute_thread.start()
# if compute_contact:
# self._contact_thread.start()
# if self.verbose:
# if compute_contact:
# print("Async pipeline enabled: Capture → Flow → HHD → Contact")
# else:
# print("Async pipeline enabled: Capture → Flow → HHD")
# def get_async(self, timeout=0.05):
# """
# Get latest result from async pipeline (non-blocking).
# Returns:
# dict or None: Sensor data if available
# """
# compute_contact = self._compute_contact_enabled
# 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
# # Sync instance variables
# # for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
# if compute_contact:
# for k in ['img', 'flow', 'normal', 'shear', 'depth', 'centroid', 'centroid_found', 'contour', 'fps', 'fnormal', 'fshearx', 'fsheary', 'ftor']:
# setattr(self, k, result.get(k))
# else:
# for k in ['img', 'flow', 'normal', 'shear', '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")
# =========================================================================
# Private - Async Workers
# =========================================================================
# def _queue_put(self, q, item):
# """Put item in queue, dropping old if full."""
# if q.full():
# try: q.get_nowait()
# except Empty: pass
# q.put(item)
# def _capture_worker(self):
# """Stage 1: Capture and preprocess frames by reusing get_img()."""
# while not self._stop_event.is_set():
# try:
# # Reuse get_img() to avoid code duplication
# if self.get_img() is None:
# continue
# # Put the processed image in the queue
# # get_img() updates self.img, so we use that
# if self.img is not None:
# self._queue_put(self._capture_queue, self.img.copy())
# except Exception as e:
# if self.verbose: print(f"Capture error: {e}")
# def _flow_worker(self):
# """Stage 2: Compute optical flow by reusing __compute_flow()."""
# while not self._stop_event.is_set():
# try:
# img = self._capture_queue.get(timeout=0.1)
# # Set __img to reuse __compute_flow()
# self.img = img
# self.__compute_flow()
# flow = self.flow.copy() if self.flow is not None else None
# 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 by reusing __compute_hhd()."""
# while not self._stop_event.is_set():
# try:
# img, flow = self._flow_queue.get(timeout=0.1)
# # Set __flow to reuse __compute_hhd()
# self.flow = flow
# self.__compute_hhd()
# # Get results from internal variables set by __compute_hhd()
# result = {
# 'img': img, 'flow': flow,
# 'normal': self.vnormal.copy() if self.vnormal is not None else None,
# 'shear': self.vshear.copy() if self.vshear is not None else None,
# # 'depth': getFilteredDensity(flow, use_cuda=self.cuda),
# 'fnormal': self.fnormal, 'fshearx': self.fshearx,
# 'fsheary': self.fsheary, 'ftor': self.ftor,
# }
# # Always include contact fields to ensure consistent result structure
# # When compute_contact is True, contact worker will update these values
# # When compute_contact is False, they remain None
# result.update({
# 'centroid': None,
# 'centroid_found': False,
# 'contour': None,
# 'depth_map': None,
# 'depth': None, # Alias for depth_map for compatibility
# })
# self._queue_put(self._result_queue, result)
# except Empty:
# continue
# except Exception as e:
# if self.verbose: print(f"HHD error: {e}")
# def _contact_worker(self):
# """Stage 4: Contact computation using vnormal from HHD results."""
# while not self._stop_event.is_set():
# try:
# # Get results from HHD worker which includes vnormal
# hhd_result = self._result_queue.get(timeout=0.1)
# vnormal = hhd_result.get('normal') # This is vnormal from HHD
# if vnormal is not None:
# # Compute contact using vnormal directly
# contact_result = self.compute_contact(vnormal=vnormal)
# # Add contact results to the existing HHD result
# depth_map = contact_result.get('divergence_map')
# hhd_result.update({
# 'centroid': contact_result.get('centroid'),
# 'centroid_found': contact_result.get('centroid_found'),
# 'contour': contact_result.get('contour'),
# 'depth_map': depth_map,
# 'depth': depth_map, # Alias for depth_map for compatibility
# })
# # Put updated result back to queue
# self._queue_put(self._result_queue, hhd_result)
# except Empty:
# continue
# except Exception as e:
# if self.verbose: print(f"Contact error: {e}")