add optical flow debug script
This commit is contained in:
ykcheungab
+449
@@ -0,0 +1,449 @@
|
||||
|
||||
"""
|
||||
Prototype tactile 2.5D motion analysis on top of the existing dense flow output.
|
||||
|
||||
This script compares several motion representations derived from optical flow:
|
||||
1. Raw flow arrows / magnitude
|
||||
2. Radial and tangential flow around an estimated contact center
|
||||
3. Differential flow features: divergence, curl, strain magnitude
|
||||
4. Pseudo-height reconstructed from divergence (qualitative only)
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import time
|
||||
|
||||
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
WINDOW_TITLE = "Tactile 2.5D Motion Prototype"
|
||||
|
||||
_local_src = os.path.join(os.path.dirname(__file__), "src")
|
||||
if _local_src not in sys.path:
|
||||
sys.path.insert(0, _local_src)
|
||||
|
||||
import orisys
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--video", "-v", type=str, default="0", help="Camera ID / video file path")
|
||||
parser.add_argument("--config", "-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
|
||||
parser.add_argument("--cal", "-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
|
||||
parser.add_argument("--verbose", action="store_true", help="Print detailed sensor logs")
|
||||
parser.add_argument("--no-display", action="store_true", help="Disable OpenCV windows")
|
||||
parser.add_argument("--save", type=str, default="", help="Optional output video path for the comparison panel")
|
||||
parser.add_argument("--max-frames", type=int, default=0, help="Stop after N frames (0 = unlimited)")
|
||||
parser.add_argument("--grid-spacing", type=int, default=20, help="Arrow visualization grid spacing")
|
||||
parser.add_argument("--flow-threshold", type=float, default=2.0, help="Arrow visualization magnitude threshold")
|
||||
parser.add_argument("--print-every", type=int, default=1, help="Print scalar metrics every N frames")
|
||||
parser.add_argument("--panel-size", type=int, default=320, help="Width/height of each panel tile")
|
||||
parser.add_argument("--plane-width-mm", type=float, default=20.0, help="Physical tactile imaging plane width in mm")
|
||||
parser.add_argument("--plane-height-mm", type=float, default=20.0, help="Physical tactile imaging plane height in mm")
|
||||
parser.add_argument("--camera-distance-mm", type=float, default=10.0, help="Distance from camera plane to imaging plane in mm")
|
||||
parser.add_argument("--display-max-width", type=int, default=1800, help="Maximum preview window width in pixels")
|
||||
parser.add_argument("--display-max-height", type=int, default=1000, help="Maximum preview window height in pixels")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def create_sensor(args):
|
||||
try:
|
||||
video_input = int(args.video)
|
||||
except ValueError:
|
||||
video_input = args.video
|
||||
|
||||
return orisys.Sensor(
|
||||
video_input,
|
||||
config_name=args.config,
|
||||
cal_path=args.cal,
|
||||
verbose=args.verbose,
|
||||
)
|
||||
|
||||
|
||||
def ensure_bgr(image):
|
||||
if image.ndim == 2:
|
||||
return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||||
return image
|
||||
|
||||
|
||||
def resize_panel(image, panel_size):
|
||||
return cv2.resize(ensure_bgr(image), (panel_size, panel_size), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
|
||||
def fit_display_image(image, max_width, max_height):
|
||||
h, w = image.shape[:2]
|
||||
scale = min(max_width / max(float(w), 1.0), max_height / max(float(h), 1.0), 1.0)
|
||||
if scale >= 0.999:
|
||||
return image
|
||||
new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
|
||||
return cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def signed_heatmap(field, colormap=cv2.COLORMAP_TURBO):
|
||||
max_abs = float(np.max(np.abs(field)))
|
||||
if max_abs <= 1e-9:
|
||||
norm = np.full(field.shape, 127, dtype=np.uint8)
|
||||
else:
|
||||
norm = np.clip((field / max_abs + 1.0) * 127.5, 0, 255).astype(np.uint8)
|
||||
return cv2.applyColorMap(norm, colormap)
|
||||
|
||||
|
||||
def magnitude_heatmap(field, colormap=cv2.COLORMAP_TURBO):
|
||||
max_val = float(np.max(field))
|
||||
if max_val <= 1e-9:
|
||||
norm = np.zeros(field.shape, dtype=np.uint8)
|
||||
else:
|
||||
norm = np.clip(field / max_val * 255.0, 0, 255).astype(np.uint8)
|
||||
return cv2.applyColorMap(norm, colormap)
|
||||
|
||||
|
||||
def add_title(image, title):
|
||||
panel = image.copy()
|
||||
cv2.rectangle(panel, (0, 0), (panel.shape[1], 28), (0, 0, 0), -1)
|
||||
cv2.putText(panel, title, (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
|
||||
return panel
|
||||
|
||||
|
||||
def make_panel(image, title, panel_size):
|
||||
return add_title(resize_panel(image, panel_size), title)
|
||||
|
||||
|
||||
def stack_panels(panels, columns=4):
|
||||
filler = np.zeros_like(panels[0])
|
||||
rows = []
|
||||
for idx in range(0, len(panels), columns):
|
||||
row = panels[idx:idx + columns]
|
||||
if len(row) < columns:
|
||||
row = row + [filler] * (columns - len(row))
|
||||
rows.append(cv2.hconcat(row))
|
||||
return cv2.vconcat(rows)
|
||||
|
||||
|
||||
def estimate_contact_center(divergence, flow_magnitude):
|
||||
positive = np.maximum(divergence, 0.0)
|
||||
weights = positive * (flow_magnitude + 1e-6)
|
||||
total = float(weights.sum())
|
||||
h, w = divergence.shape
|
||||
|
||||
if total <= 1e-9:
|
||||
return (w // 2, h // 2)
|
||||
|
||||
ys, xs = np.indices(divergence.shape, dtype=np.float32)
|
||||
cx = int(np.clip((weights * xs).sum() / total, 0, w - 1))
|
||||
cy = int(np.clip((weights * ys).sum() / total, 0, h - 1))
|
||||
return (cx, cy)
|
||||
|
||||
|
||||
def prepare_support_mask(support_map):
|
||||
if support_map is None:
|
||||
return None
|
||||
|
||||
support = np.abs(support_map).astype(np.float32)
|
||||
max_support = float(np.max(support))
|
||||
if max_support > 0.0:
|
||||
support /= max_support
|
||||
support = (support > 0.05).astype(np.float32)
|
||||
if float(np.mean(support)) < 0.002:
|
||||
return None
|
||||
return support
|
||||
|
||||
|
||||
def decompose_radial_tangential(flow, center):
|
||||
h, w = flow.shape[:2]
|
||||
ys, xs = np.indices((h, w), dtype=np.float32)
|
||||
cx, cy = center
|
||||
|
||||
rx = xs - float(cx)
|
||||
ry = ys - float(cy)
|
||||
r_norm = np.sqrt(rx * rx + ry * ry)
|
||||
r_norm = np.where(r_norm < 1e-6, 1.0, r_norm)
|
||||
|
||||
rhat_x = rx / r_norm
|
||||
rhat_y = ry / r_norm
|
||||
that_x = -rhat_y
|
||||
that_y = rhat_x
|
||||
|
||||
u = flow[:, :, 0]
|
||||
v = flow[:, :, 1]
|
||||
radial = u * rhat_x + v * rhat_y
|
||||
tangential = u * that_x + v * that_y
|
||||
return radial, tangential
|
||||
|
||||
|
||||
def compute_flow_jacobian_features(flow):
|
||||
u = flow[:, :, 0]
|
||||
v = flow[:, :, 1]
|
||||
|
||||
du_dy, du_dx = np.gradient(u)
|
||||
dv_dy, dv_dx = np.gradient(v)
|
||||
|
||||
divergence = du_dx + dv_dy
|
||||
curl = dv_dx - du_dy
|
||||
strain_xx = du_dx
|
||||
strain_yy = dv_dy
|
||||
strain_xy = 0.5 * (du_dy + dv_dx)
|
||||
strain_mag = np.sqrt(strain_xx * strain_xx + strain_yy * strain_yy + 2.0 * strain_xy * strain_xy)
|
||||
|
||||
return {
|
||||
"divergence": divergence,
|
||||
"curl": curl,
|
||||
"strain_mag": strain_mag,
|
||||
}
|
||||
|
||||
|
||||
def reconstruct_pseudo_height(divergence, flow_magnitude=None, support_map=None):
|
||||
rhs = divergence.astype(np.float32).copy()
|
||||
|
||||
# No-contact guard: tiny flow noise becomes large low-frequency drift after Poisson solve.
|
||||
if flow_magnitude is not None:
|
||||
mean_mag = float(np.mean(flow_magnitude))
|
||||
max_mag = float(np.max(flow_magnitude))
|
||||
if mean_mag < 0.03 and max_mag < 0.15:
|
||||
return np.zeros_like(rhs, dtype=np.float32)
|
||||
|
||||
support = prepare_support_mask(support_map)
|
||||
if support_map is not None:
|
||||
if support is None:
|
||||
return np.zeros_like(rhs, dtype=np.float32)
|
||||
rhs *= support
|
||||
|
||||
noise_floor = max(float(np.std(rhs)) * 1.5, 1e-4)
|
||||
rhs[np.abs(rhs) < noise_floor] = 0.0
|
||||
if not np.any(rhs):
|
||||
return np.zeros_like(rhs, dtype=np.float32)
|
||||
|
||||
# Zero-pad before FFT solve to reduce periodic wrap-around ghosts, then crop back.
|
||||
h, w = rhs.shape
|
||||
pad_y = max(8, h // 2)
|
||||
pad_x = max(8, w // 2)
|
||||
rhs_padded = np.pad(rhs, ((pad_y, pad_y), (pad_x, pad_x)), mode="constant")
|
||||
rhs_padded = rhs_padded - float(rhs_padded.mean())
|
||||
|
||||
hp, wp = rhs_padded.shape
|
||||
fy = np.fft.fftfreq(hp).reshape(-1, 1)
|
||||
fx = np.fft.fftfreq(wp).reshape(1, -1)
|
||||
denom = (2.0 * np.pi * fx) ** 2 + (2.0 * np.pi * fy) ** 2
|
||||
denom[0, 0] = 1.0
|
||||
|
||||
rhs_hat = np.fft.fft2(rhs_padded)
|
||||
height_hat = -rhs_hat / denom
|
||||
height_hat[0, 0] = 0.0
|
||||
height_full = np.real(np.fft.ifft2(height_hat)).astype(np.float32)
|
||||
height = height_full[pad_y:pad_y + h, pad_x:pad_x + w]
|
||||
height = cv2.GaussianBlur(height, (0, 0), 1.0)
|
||||
return height
|
||||
|
||||
|
||||
def build_visualization(img, flow, center, features, radial, tangential, pseudo_height, depth_map, args):
|
||||
arrows = orisys.util.draw_arrows(
|
||||
img,
|
||||
flow,
|
||||
threshold=args.flow_threshold,
|
||||
grid_spacing=args.grid_spacing,
|
||||
arrow_scale=1.0,
|
||||
)
|
||||
|
||||
contact_view = ensure_bgr(img)
|
||||
cv2.circle(contact_view, center, 6, (0, 255, 0), -1)
|
||||
cv2.circle(contact_view, center, 12, (0, 255, 0), 2)
|
||||
|
||||
panel_images = [
|
||||
(contact_view, "Image / center"),
|
||||
(arrows, "Raw flow"),
|
||||
(signed_heatmap(radial), "Radial flow"),
|
||||
(signed_heatmap(tangential), "Tangential flow"),
|
||||
(signed_heatmap(features["divergence"]), "Divergence"),
|
||||
(signed_heatmap(features["curl"]), "Curl"),
|
||||
(magnitude_heatmap(features["strain_mag"]), "Strain mag"),
|
||||
(signed_heatmap(depth_map), "Depth map"),
|
||||
(signed_heatmap(pseudo_height), "Pseudo height"),
|
||||
]
|
||||
panels = [make_panel(image, title, args.panel_size) for image, title in panel_images]
|
||||
return stack_panels(panels)
|
||||
|
||||
|
||||
def make_writer(path, frame_size, fps):
|
||||
if not path:
|
||||
return None
|
||||
parent = Path(path).parent
|
||||
if parent and not parent.exists():
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), fps, frame_size)
|
||||
|
||||
|
||||
def approximate_raw_flow_to_mm(flow_px, plane_width_mm, plane_height_mm, camera_distance_mm, image_shape):
|
||||
h, w = image_shape[:2]
|
||||
fx = camera_distance_mm * w / max(plane_width_mm, 1e-6)
|
||||
fy = camera_distance_mm * h / max(plane_height_mm, 1e-6)
|
||||
flow_mm = flow_px.copy().astype(np.float32)
|
||||
flow_mm[:, :, 0] *= camera_distance_mm / max(fx, 1e-6)
|
||||
flow_mm[:, :, 1] *= camera_distance_mm / max(fy, 1e-6)
|
||||
return flow_mm
|
||||
|
||||
|
||||
def compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height):
|
||||
return {
|
||||
"flow_mean": float(flow_magnitude.mean()),
|
||||
"flow_mean_mm": float(np.linalg.norm(raw_flow_mm, axis=2).mean()),
|
||||
"radial_abs_mean": float(np.mean(np.abs(radial))),
|
||||
"tangential_abs_mean": float(np.mean(np.abs(tangential))),
|
||||
"divergence_mean": float(features["divergence"].mean()),
|
||||
"curl_abs_mean": float(np.mean(np.abs(features["curl"]))),
|
||||
"strain_mean": float(features["strain_mag"].mean()),
|
||||
"depth_peak": float(np.max(np.abs(depth_map))),
|
||||
"pseudo_height_peak": float(np.max(np.abs(pseudo_height))),
|
||||
}
|
||||
|
||||
|
||||
def format_metrics(frame_count, fps, metrics):
|
||||
line = (
|
||||
f"frame={frame_count:04d} fps={fps:6.2f} "
|
||||
f"flow_mean={metrics['flow_mean']:.4f} "
|
||||
f"flow_mean_mm={metrics['flow_mean_mm']:.4f} "
|
||||
f"radial_abs={metrics['radial_abs_mean']:.4f} "
|
||||
f"tangential_abs={metrics['tangential_abs_mean']:.4f} "
|
||||
f"div_mean={metrics['divergence_mean']:.4f} "
|
||||
f"curl_abs={metrics['curl_abs_mean']:.4f} "
|
||||
f"strain={metrics['strain_mean']:.4f} "
|
||||
f"depth_peak={metrics['depth_peak']:.4f} "
|
||||
f"height_peak={metrics['pseudo_height_peak']:.4f}"
|
||||
)
|
||||
t = metrics.get("_timings")
|
||||
if t:
|
||||
line += (
|
||||
f" | analyze {t['total']:.2f}ms"
|
||||
f" (mag:{t['magnitude']:.2f} jac:{t['jacobian']:.2f}"
|
||||
f" rad_tan:{t['rad_tan']:.2f} pseudo_h:{t['pseudo_h']:.2f}"
|
||||
f" metrics:{t['metrics']:.2f} vis:{t['visualize']:.2f})"
|
||||
)
|
||||
return line
|
||||
|
||||
|
||||
def read_sensor_frame(sensor):
|
||||
sensor.compute_flow()
|
||||
sensor.compute_hhd()
|
||||
sensor.compute_contact()
|
||||
return sensor.read_info(
|
||||
sensor.info.FPS,
|
||||
sensor.info.VRAW,
|
||||
sensor.info.IMG,
|
||||
sensor.info.DEPTH,
|
||||
)
|
||||
|
||||
|
||||
def analyze_frame(img, flow, depth_map, args):
|
||||
t = time.perf_counter()
|
||||
|
||||
flow_magnitude = np.linalg.norm(flow, axis=2)
|
||||
t_mag = time.perf_counter()
|
||||
|
||||
features = compute_flow_jacobian_features(flow)
|
||||
t_jac = time.perf_counter()
|
||||
|
||||
center = estimate_contact_center(features["divergence"], flow_magnitude)
|
||||
radial, tangential = decompose_radial_tangential(flow, center)
|
||||
t_decomp = time.perf_counter()
|
||||
|
||||
pseudo_height = reconstruct_pseudo_height(
|
||||
features["divergence"],
|
||||
flow_magnitude=flow_magnitude,
|
||||
support_map=depth_map,
|
||||
)
|
||||
t_height = time.perf_counter()
|
||||
|
||||
raw_flow_mm = approximate_raw_flow_to_mm(
|
||||
flow,
|
||||
args.plane_width_mm,
|
||||
args.plane_height_mm,
|
||||
args.camera_distance_mm,
|
||||
img.shape,
|
||||
)
|
||||
metrics = compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height)
|
||||
t_metrics = time.perf_counter()
|
||||
|
||||
canvas = build_visualization(
|
||||
img,
|
||||
flow,
|
||||
center,
|
||||
features,
|
||||
radial,
|
||||
tangential,
|
||||
pseudo_height,
|
||||
depth_map,
|
||||
args,
|
||||
)
|
||||
t_vis = time.perf_counter()
|
||||
|
||||
metrics["_timings"] = {
|
||||
"magnitude": (t_mag - t) * 1000,
|
||||
"jacobian": (t_jac - t_mag) * 1000,
|
||||
"rad_tan": (t_decomp - t_jac) * 1000,
|
||||
"pseudo_h": (t_height - t_decomp) * 1000,
|
||||
"metrics": (t_metrics - t_height) * 1000,
|
||||
"visualize": (t_vis - t_metrics) * 1000,
|
||||
"total": (t_vis - t) * 1000,
|
||||
}
|
||||
return canvas, metrics
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
sensor = create_sensor(args)
|
||||
writer = None
|
||||
frame_count = 0
|
||||
print_every = max(args.print_every, 1)
|
||||
|
||||
try:
|
||||
print("\nPress 'q' to quit, 'r' to reset tracking.")
|
||||
|
||||
while True:
|
||||
img = sensor.get_img()
|
||||
if img is None:
|
||||
print("No more frames returned from sensor.")
|
||||
break
|
||||
|
||||
fps, flow, img, depth_map = read_sensor_frame(sensor)
|
||||
canvas, metrics = analyze_frame(img, flow, depth_map, args)
|
||||
|
||||
if writer is None and args.save:
|
||||
writer = make_writer(args.save, (canvas.shape[1], canvas.shape[0]), max(float(fps), 1.0))
|
||||
|
||||
if writer is not None:
|
||||
writer.write(canvas)
|
||||
|
||||
frame_count += 1
|
||||
|
||||
if frame_count % print_every == 0:
|
||||
print(format_metrics(frame_count, fps, metrics))
|
||||
|
||||
if not args.no_display:
|
||||
display_canvas = fit_display_image(
|
||||
canvas,
|
||||
max_width=args.display_max_width,
|
||||
max_height=args.display_max_height,
|
||||
)
|
||||
cv2.imshow(WINDOW_TITLE, display_canvas)
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
if key == ord("q"):
|
||||
print("\nExiting...")
|
||||
break
|
||||
if key == ord("r"):
|
||||
sensor.reset()
|
||||
print("Sensor reset.")
|
||||
|
||||
if args.max_frames > 0 and frame_count >= args.max_frames:
|
||||
print(f"Reached max_frames={args.max_frames}.")
|
||||
break
|
||||
finally:
|
||||
if writer is not None:
|
||||
writer.release()
|
||||
sensor.disconnect()
|
||||
if not args.no_display:
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user