cheng
406 lines
12 KiB
Python
406 lines
12 KiB
Python
import os
|
|
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
|
|
|
import argparse
|
|
import multiprocessing as mp
|
|
import queue
|
|
import time
|
|
from multiprocessing import Process, Queue
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .paths import ensure_project_paths
|
|
|
|
ensure_project_paths()
|
|
|
|
|
|
def parse_video_source(value):
|
|
if isinstance(value, str):
|
|
value = value.strip()
|
|
if value.isdigit():
|
|
return int(value)
|
|
return value
|
|
|
|
|
|
def _put_latest(result_queue, data):
|
|
try:
|
|
result_queue.put_nowait(data)
|
|
return
|
|
except queue.Full:
|
|
pass
|
|
|
|
try:
|
|
result_queue.get_nowait()
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
result_queue.put_nowait(data)
|
|
except queue.Full:
|
|
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,
|
|
config_name,
|
|
rotation_config_path,
|
|
result_queue,
|
|
stop_event,
|
|
target_fps,
|
|
cuda,
|
|
isstitch,
|
|
backend,
|
|
motion_threshold,
|
|
include_visuals,
|
|
visual_size,
|
|
):
|
|
ensure_project_paths()
|
|
import orisys
|
|
|
|
sensor = None
|
|
try:
|
|
sensor = orisys.Sensor(
|
|
vid_src,
|
|
isstitch=isstitch,
|
|
cuda=cuda,
|
|
verbose=False,
|
|
config_name=config_name,
|
|
backend=backend,
|
|
rotation_config_path=rotation_config_path,
|
|
)
|
|
|
|
frame_interval = 1.0 / target_fps if target_fps and target_fps > 0 else 0.0
|
|
|
|
while not stop_event.is_set():
|
|
t_start = time.perf_counter()
|
|
img = sensor.get_img()
|
|
if img is None:
|
|
time.sleep(0.01)
|
|
continue
|
|
|
|
is_contact = sensor.compute_deformation(
|
|
check_motion=True,
|
|
threshold=motion_threshold,
|
|
)
|
|
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
|
|
|
|
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)
|
|
if sleep_time > 0.001:
|
|
time.sleep(sleep_time)
|
|
|
|
except Exception as exc:
|
|
_put_latest(
|
|
result_queue,
|
|
{
|
|
"sensor_id": sensor_id,
|
|
"error": str(exc),
|
|
"timestamp": time.time(),
|
|
},
|
|
)
|
|
finally:
|
|
if sensor is not None:
|
|
try:
|
|
sensor.disconnect()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class TwoSensorForceReader:
|
|
def __init__(
|
|
self,
|
|
video1=0,
|
|
video2=1,
|
|
config1="./config/ddjx01.json",
|
|
config2="./config/ddjx01.json",
|
|
rotation1="./config/rotation_config_0.json",
|
|
rotation2="./config/rotation_config_1.json",
|
|
target_fps=30,
|
|
cuda=True,
|
|
isstitch=True,
|
|
backend="auto",
|
|
motion_threshold=0,
|
|
queue_size=5,
|
|
include_visuals=False,
|
|
visual_size=320,
|
|
):
|
|
self.sensor_configs = [
|
|
{
|
|
"sensor_id": 0,
|
|
"vid_src": parse_video_source(video1),
|
|
"config_name": config1,
|
|
"rotation_config_path": rotation1,
|
|
},
|
|
{
|
|
"sensor_id": 1,
|
|
"vid_src": parse_video_source(video2),
|
|
"config_name": config2,
|
|
"rotation_config_path": rotation2,
|
|
},
|
|
]
|
|
self.target_fps = target_fps
|
|
self.cuda = cuda
|
|
self.isstitch = isstitch
|
|
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 = []
|
|
self.processes = []
|
|
self.latest = [None, None]
|
|
|
|
def start(self):
|
|
if self.processes:
|
|
return
|
|
|
|
self.result_queues = [Queue(maxsize=self.queue_size) for _ in self.sensor_configs]
|
|
self.stop_events = [mp.Event() for _ in self.sensor_configs]
|
|
self.latest = [None] * len(self.sensor_configs)
|
|
|
|
for idx, config in enumerate(self.sensor_configs):
|
|
process = Process(
|
|
target=_sensor_force_worker,
|
|
args=(
|
|
config["sensor_id"],
|
|
config["vid_src"],
|
|
config["config_name"],
|
|
config["rotation_config_path"],
|
|
self.result_queues[idx],
|
|
self.stop_events[idx],
|
|
self.target_fps,
|
|
self.cuda,
|
|
self.isstitch,
|
|
self.backend,
|
|
self.motion_threshold,
|
|
self.include_visuals,
|
|
self.visual_size,
|
|
),
|
|
)
|
|
process.start()
|
|
self.processes.append(process)
|
|
|
|
def update(self):
|
|
for idx, result_queue in enumerate(self.result_queues):
|
|
data = None
|
|
while True:
|
|
try:
|
|
data = result_queue.get_nowait()
|
|
except queue.Empty:
|
|
break
|
|
except Exception:
|
|
break
|
|
|
|
if data is not None:
|
|
if "error" in data:
|
|
raise RuntimeError(
|
|
f"sensor {data.get('sensor_id', idx)} failed: {data['error']}"
|
|
)
|
|
self.latest[idx] = data
|
|
|
|
return self.latest
|
|
|
|
def get_samples(self, timeout=0.0):
|
|
deadline = time.perf_counter() + max(0.0, float(timeout))
|
|
while True:
|
|
self.update()
|
|
if all(sample is not None for sample in self.latest):
|
|
return self.latest[0], self.latest[1]
|
|
if time.perf_counter() >= deadline:
|
|
return self.latest[0], self.latest[1]
|
|
time.sleep(0.005)
|
|
|
|
def get_forces(self, timeout=0.0):
|
|
left, right = self.get_samples(timeout=timeout)
|
|
left_force = None if left is None else left["fnormal"]
|
|
right_force = None if right is None else right["fnormal"]
|
|
return left_force, right_force
|
|
|
|
def stop(self):
|
|
for event in self.stop_events:
|
|
event.set()
|
|
|
|
for process in self.processes:
|
|
process.join(timeout=2.0)
|
|
if process.is_alive():
|
|
process.terminate()
|
|
process.join(timeout=1.0)
|
|
|
|
self.processes = []
|
|
self.stop_events = []
|
|
self.result_queues = []
|
|
|
|
def __enter__(self):
|
|
self.start()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, traceback):
|
|
self.stop()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--video1", "-v1", type=str, default="0")
|
|
parser.add_argument("--video2", "-v2", type=str, default="1")
|
|
parser.add_argument("--config1", type=str, default="./config/ddjx01.json")
|
|
parser.add_argument("--config2", type=str, default="./config/ddjx01.json")
|
|
parser.add_argument("--rotation1", "-r1", type=str, default="./config/rotation_config_0.json")
|
|
parser.add_argument("--rotation2", "-r2", type=str, default="./config/rotation_config_1.json")
|
|
parser.add_argument("--target-fps", type=float, default=30)
|
|
parser.add_argument("--motion-threshold", type=float, default=0)
|
|
parser.add_argument("--cpu", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
reader = TwoSensorForceReader(
|
|
video1=args.video1,
|
|
video2=args.video2,
|
|
config1=args.config1,
|
|
config2=args.config2,
|
|
rotation1=args.rotation1,
|
|
rotation2=args.rotation2,
|
|
target_fps=args.target_fps,
|
|
cuda=not args.cpu,
|
|
motion_threshold=args.motion_threshold,
|
|
)
|
|
|
|
print("Starting two-sensor force reader. Press Ctrl+C to stop.")
|
|
reader.start()
|
|
try:
|
|
while True:
|
|
left_sample, right_sample = reader.get_samples(timeout=1.0)
|
|
left_force, right_force = reader.get_forces()
|
|
left_fps = 0.0 if left_sample is None else left_sample["fps"]
|
|
right_fps = 0.0 if right_sample is None else right_sample["fps"]
|
|
left_contact = False if left_sample is None else left_sample["is_contact"]
|
|
right_contact = False if right_sample is None else right_sample["is_contact"]
|
|
|
|
print(
|
|
"left={:>10} right={:>10} fps=({:5.1f}, {:5.1f}) contact=({}, {})".format(
|
|
"None" if left_force is None else f"{left_force:.4f}",
|
|
"None" if right_force is None else f"{right_force:.4f}",
|
|
left_fps,
|
|
right_fps,
|
|
int(left_contact),
|
|
int(right_contact),
|
|
)
|
|
)
|
|
time.sleep(0.05)
|
|
except KeyboardInterrupt:
|
|
print("\nStopping.")
|
|
finally:
|
|
reader.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mp.freeze_support()
|
|
main()
|