qt第一版可视化

gripper_control/force_reader.py

@@ -7,6 +7,9 @@ import queue
 import time
 from multiprocessing import Process, Queue
 
+import cv2
+import numpy as np
+
 from .paths import ensure_project_paths
 
 ensure_project_paths()
@@ -38,6 +41,69 @@ def _put_latest(result_queue, data):
         pass
 
 
+def _draw_magnitude_map(flow, visual_size, threshold=1.5):
+    magnitude = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)
+    h, w = magnitude.shape
+    if h <= 0 or w <= 0:
+        return None
+
+    downsample_factor = 10
+    magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
+    magnitude_smoothed = cv2.resize(
+        magnitude_downsampled,
+        (w, h),
+        interpolation=cv2.INTER_CUBIC,
+    )
+
+    magnitude_thresholded = magnitude_smoothed.copy()
+    magnitude_thresholded[magnitude_smoothed < threshold] = 0
+
+    mag_min = float(magnitude_thresholded.min())
+    mag_max = float(magnitude_thresholded.max())
+    if mag_max > mag_min:
+        magnitude_norm = (
+            (magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255
+        ).astype(np.uint8)
+    else:
+        magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)
+
+    magnitude_colored = cv2.applyColorMap(magnitude_norm, cv2.COLORMAP_JET)
+    return cv2.resize(
+        magnitude_colored,
+        (visual_size, visual_size),
+        interpolation=cv2.INTER_CUBIC,
+    )
+
+
+def _make_visuals(orisys_module, img, flow, visual_size):
+    if img is None or flow is None or np.asarray(flow).size == 0:
+        return None, None, 0.0, 0.0
+
+    if len(img.shape) == 2:
+        img_bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    else:
+        img_bgr = img
+
+    h, w = img_bgr.shape[:2]
+    bg = np.empty((h, w, 3), dtype=img_bgr.dtype)
+    bg[:, :, 0] = 22
+    bg[:, :, 1] = 16
+    bg[:, :, 2] = 11
+
+    arrows = orisys_module.util.draw_arrows(
+        bg,
+        flow,
+        threshold=2,
+        grid_spacing=20,
+        arrow_scale=0.5,
+        below_threshold_color=(160, 160, 160),
+    )
+    arrows = cv2.resize(arrows, (visual_size, visual_size), interpolation=cv2.INTER_CUBIC)
+    magnitude = _draw_magnitude_map(flow, visual_size)
+    flow_mag = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)
+    return arrows, magnitude, float(flow_mag.mean()), float(flow_mag.max())
+
+
 def _sensor_force_worker(
     sensor_id,
     vid_src,
@@ -50,6 +116,8 @@ def _sensor_force_worker(
     isstitch,
     backend,
     motion_threshold,
+    include_visuals,
+    visual_size,
 ):
     ensure_project_paths()
     import orisys
@@ -79,25 +147,49 @@ def _sensor_force_worker(
                 check_motion=True,
                 threshold=motion_threshold,
             )
-            fps, fnormal, fshearx, fsheary = sensor.read_info(
-                sensor.info.FPS,
-                sensor.info.FNORMAL,
-                sensor.info.FSHEARX,
-                sensor.info.FSHEARY,
-            )
+            if include_visuals:
+                fps, fnormal, fshearx, fsheary, flow, img_view = sensor.read_info(
+                    sensor.info.FPS,
+                    sensor.info.FNORMAL,
+                    sensor.info.FSHEARX,
+                    sensor.info.FSHEARY,
+                    sensor.info.VRAW,
+                    sensor.info.IMG,
+                )
+                flow_view, magnitude_view, flow_mean, flow_max = _make_visuals(
+                    orisys,
+                    img_view,
+                    flow,
+                    int(visual_size),
+                )
+            else:
+                fps, fnormal, fshearx, fsheary = sensor.read_info(
+                    sensor.info.FPS,
+                    sensor.info.FNORMAL,
+                    sensor.info.FSHEARX,
+                    sensor.info.FSHEARY,
+                )
+                flow_view = None
+                magnitude_view = None
+                flow_mean = 0.0
+                flow_max = 0.0
 
-            _put_latest(
-                result_queue,
-                {
-                    "sensor_id": sensor_id,
-                    "fnormal": float(fnormal),
-                    "fshearx": float(fshearx),
-                    "fsheary": float(fsheary),
-                    "fps": float(fps),
-                    "is_contact": bool(is_contact),
-                    "timestamp": time.time(),
-                },
-            )
+            data = {
+                "sensor_id": sensor_id,
+                "fnormal": float(fnormal),
+                "fshearx": float(fshearx),
+                "fsheary": float(fsheary),
+                "fps": float(fps),
+                "is_contact": bool(is_contact),
+                "timestamp": time.time(),
+                "flow_mean": flow_mean,
+                "flow_max": flow_max,
+            }
+            if include_visuals:
+                data["flow_view"] = flow_view
+                data["magnitude_view"] = magnitude_view
+
+            _put_latest(result_queue, data)
 
             if frame_interval > 0:
                 sleep_time = frame_interval - (time.perf_counter() - t_start)
@@ -136,6 +228,8 @@ class TwoSensorForceReader:
         backend="auto",
         motion_threshold=0,
         queue_size=5,
+        include_visuals=False,
+        visual_size=320,
     ):
         self.sensor_configs = [
             {
@@ -157,6 +251,8 @@ class TwoSensorForceReader:
         self.backend = backend
         self.motion_threshold = motion_threshold
         self.queue_size = queue_size
+        self.include_visuals = include_visuals
+        self.visual_size = visual_size
 
         self.result_queues = []
         self.stop_events = []
@@ -186,6 +282,8 @@ class TwoSensorForceReader:
                     self.isstitch,
                     self.backend,
                     self.motion_threshold,
+                    self.include_visuals,
+                    self.visual_size,
                 ),
             )
             process.start()