fixed async pipeline
This commit is contained in:
unknown
+118
-77
@@ -202,7 +202,7 @@ class Sensor:
|
||||
|
||||
return True
|
||||
|
||||
def compute_deformation(self, log_timing=False):
|
||||
def compute_deformation(self):
|
||||
"""
|
||||
Compute optical flow and decompose into force components.
|
||||
|
||||
@@ -211,66 +211,76 @@ class Sensor:
|
||||
"""
|
||||
if self.__img is None:
|
||||
raise Exception("Warning: Get image first")
|
||||
if log_timing:
|
||||
timings = {}
|
||||
t_start = time.perf_counter()
|
||||
|
||||
# Optical flow
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
self.__compute_flow()
|
||||
self.__compute_hhd()
|
||||
# FPS calculation
|
||||
self._fps = 1.0 / (time.perf_counter() - self.last_log_time)
|
||||
self.last_log_time = time.perf_counter()
|
||||
|
||||
return True
|
||||
|
||||
def __compute_flow(self):
|
||||
if self.__img is None:
|
||||
raise Exception("Warning: Get image first")
|
||||
|
||||
flow = self.__tracker.track(self.__img)
|
||||
if hasattr(flow, 'download'):
|
||||
flow = flow.download()
|
||||
if log_timing: timings['optical_flow'] = (time.perf_counter() - t0) * 1000
|
||||
self.__flow = flow
|
||||
return True
|
||||
|
||||
def __compute_hhd(self):
|
||||
|
||||
flow = self.__flow
|
||||
if flow is None:
|
||||
raise Exception("Warning: Compute flow first")
|
||||
|
||||
# Initialize HHD solver
|
||||
if self.__hhd is None:
|
||||
self.__hhd = nHHD(grid=flow.shape[:2], spacings=(18, 18))
|
||||
|
||||
# Force computation
|
||||
if log_timing: t0 = time.perf_counter()
|
||||
if self.cuda and self.CUPY_AVAILABLE:
|
||||
_, forces = self.__hhd.get_force_gpu_v2(flow)
|
||||
else:
|
||||
_, forces = self.__hhd.get_force(flow)
|
||||
self.__fnormal, self.__fshearx, self.__fsheary, self.__ftor = forces
|
||||
if log_timing: timings['hhd_decompose'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
# Get components
|
||||
self.__vnormal = self.__hhd.get_normal()
|
||||
self.__vshear = self.__hhd.get_shear()
|
||||
|
||||
# # Depth computation
|
||||
# if log_timing: t0 = time.perf_counter()
|
||||
# self.depth = getFilteredDensity(self.vnormal, use_cuda=self.cuda)
|
||||
# if log_timing: timings['get_density'] = (time.perf_counter() - t0) * 1000
|
||||
|
||||
|
||||
# Update state
|
||||
self.__flow = flow
|
||||
self.compute_deformation_valid = True
|
||||
# FPS calculation
|
||||
self._fps = 1.0 / (time.perf_counter() - self.last_log_time)
|
||||
self.last_log_time = time.perf_counter()
|
||||
|
||||
if log_timing:
|
||||
timings['total'] = (time.perf_counter() - t_start) * 1000
|
||||
self._print_timings("compute_deformation", timings)
|
||||
|
||||
return True
|
||||
|
||||
def compute_contact(self):
|
||||
def compute_contact(self, vnormal=None):
|
||||
"""
|
||||
Compute contact information from vnormal field.
|
||||
|
||||
multiple_foe=False
|
||||
|
||||
if not self.compute_deformation_valid:
|
||||
raise Exception("Warning: Compute deformation first")
|
||||
Args:
|
||||
vnormal: Optional vnormal field. If None, uses self.__vnormal.
|
||||
|
||||
Returns:
|
||||
dict: Contact result with centroid, contour, and depth_map
|
||||
"""
|
||||
multiple_foe = False
|
||||
|
||||
# Use provided vnormal or fall back to internal state
|
||||
vnormal_to_use = vnormal if vnormal is not None else self.__vnormal
|
||||
|
||||
if vnormal_to_use is None:
|
||||
raise Exception("Warning: vnormal is required for contact computation")
|
||||
|
||||
if multiple_foe:
|
||||
result = self.__contact3d.compute_multiple_foe(self.__vnormal)
|
||||
result = self.__contact3d.compute_multiple_foe(vnormal_to_use)
|
||||
else:
|
||||
if self.cuda:
|
||||
result = self.__contact3d.compute_foe_gpu(self.__vnormal)
|
||||
result = self.__contact3d.compute_foe_gpu(vnormal_to_use)
|
||||
else:
|
||||
result = self.__contact3d.compute_foe(self.__vnormal)
|
||||
result = self.__contact3d.compute_foe(vnormal_to_use)
|
||||
|
||||
self.__centroid = result['centroid']
|
||||
self.__centroid_found = result['centroid_found']
|
||||
@@ -279,7 +289,7 @@ class Sensor:
|
||||
self.__depth_map = result['divergence_map']
|
||||
|
||||
self.compute_contact_valid = True
|
||||
return True
|
||||
return result
|
||||
|
||||
|
||||
# def get_depth(self, denoise_sigma=15):
|
||||
@@ -368,7 +378,7 @@ class Sensor:
|
||||
# Public API - Async Pipeline
|
||||
# =========================================================================
|
||||
|
||||
def enable_async(self, queue_size=2):
|
||||
def enable_async(self, queue_size=2, compute_contact=False):
|
||||
"""
|
||||
Enable 3-stage async pipeline: Capture → Flow → HHD.
|
||||
|
||||
@@ -376,25 +386,33 @@ class Sensor:
|
||||
"""
|
||||
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)
|
||||
self._contact_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)
|
||||
self._contact_thread = Thread(target=self._contact_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:
|
||||
print("Async pipeline enabled: Capture → Flow → HHD")
|
||||
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):
|
||||
"""
|
||||
@@ -403,6 +421,7 @@ class Sensor:
|
||||
Returns:
|
||||
dict or None: Sensor data if available
|
||||
"""
|
||||
compute_contact = self._compute_contact_enabled
|
||||
if not self._async_enabled:
|
||||
return None
|
||||
try:
|
||||
@@ -413,8 +432,12 @@ class Sensor:
|
||||
|
||||
# Sync instance variables
|
||||
# for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
for k in ['img', 'flow', 'normal', 'shear', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
setattr(self, k, result.get(k))
|
||||
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:
|
||||
@@ -564,37 +587,29 @@ class Sensor:
|
||||
q.put(item)
|
||||
|
||||
def _capture_worker(self):
|
||||
"""Stage 1: Capture and preprocess frames."""
|
||||
alpha = 0.7
|
||||
"""Stage 1: Capture and preprocess frames by reusing get_img()."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
ret, frame = self.cap.read()
|
||||
if not ret:
|
||||
# Reuse get_img() to avoid code duplication
|
||||
if self.get_img() is None:
|
||||
continue
|
||||
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
self.__stitcher.get_stitch_unit(gray)
|
||||
self._extract_tiles(gray)
|
||||
img = np.ascontiguousarray(self.__stitcher.stitch_rt_optimized())
|
||||
img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
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))
|
||||
# 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."""
|
||||
"""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)
|
||||
flow = self.tracker.track(img)
|
||||
if hasattr(flow, 'download'):
|
||||
flow = flow.download()
|
||||
# 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
|
||||
@@ -602,38 +617,64 @@ class Sensor:
|
||||
if self.verbose: print(f"Flow error: {e}")
|
||||
|
||||
def _hhd_worker(self):
|
||||
"""Stage 3: HHD decomposition."""
|
||||
"""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)
|
||||
|
||||
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, {
|
||||
# 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.hhd.get_normal(), 'shear': self.hhd.get_shear(),
|
||||
'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': forces[0], 'fshearx': forces[1],
|
||||
'fsheary': forces[2], 'ftor': forces[3],
|
||||
'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."""
|
||||
"""Stage 4: Contact computation using vnormal from HHD results."""
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
img, flow = self._flow_queue.get(timeout=0.1)
|
||||
contact_result = self.compute_contact(flow)
|
||||
self._queue_put(self._result_queue, contact_result)
|
||||
# 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:
|
||||
|
||||
Reference in New Issue
Block a user