#!/usr/bin/env python3 """Interactive desktop window for the Dropbear raw-RGB live stream. Frames arrive on stdin. Camera drags and robot-velocity joystick motion are emitted as one-line commands on stdout for the simulator process. """ from __future__ import annotations import argparse import queue import sys import threading import time import tkinter as tk import cv2 import numpy as np from PIL import Image, ImageTk def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--width", type=int, required=True) parser.add_argument("--height", type=int, required=True) parser.add_argument("--display-width", type=int, default=1200) parser.add_argument("--joystick-max-speed", type=float, default=1.0) parser.add_argument("--joystick-max-yaw-rate", type=float, default=1.0) parser.add_argument("--source-fps", type=float, default=5.0) parser.add_argument("--display-fps", type=float, default=50.0) parser.add_argument("--training-live-enabled", action="store_true") parser.add_argument("--title", required=True) args = parser.parse_args() if not 0.25 <= args.source_fps <= 60.0: parser.error("--source-fps must be between 0.25 and 60") if not 10.0 <= args.display_fps <= 60.0: parser.error("--display-fps must be between 10 and 60") display_height = round(args.height * args.display_width / args.width) expected_bytes = args.width * args.height * 3 frames: queue.Queue[bytes] = queue.Queue(maxsize=2) stopped = threading.Event() root = tk.Tk() root.title(args.title) root.geometry(f"{args.display_width}x{display_height}") root.minsize(640, 360) root.resizable(True, True) canvas = tk.Canvas( root, width=args.display_width, height=display_height, borderwidth=0, highlightthickness=0, background="#050910", ) canvas.pack(fill=tk.BOTH, expand=True) image_item = canvas.create_image(0, 0, anchor=tk.NW) pad_radius = 46 knob_radius = 8 pad_center = {"x": 64, "y": display_height - 62} panel_item = canvas.create_rectangle( 0, 0, 1, 1, fill="#050910", outline="#26384f", width=1, ) label_item = canvas.create_text( 0, 0, anchor=tk.W, text=( f"PLANAR ±{args.joystick_max_speed:.2f} m/s " f"YAW ±{args.joystick_max_yaw_rate:.2f}" ), fill="#dce6f5", font=("Sans", 8), ) value_item = canvas.create_text( 0, 0, anchor=tk.W, text="vx +0.00 vy +0.00 wz +0.00", fill="#9fe8ff", font=("Sans", 7), ) pad_item = canvas.create_oval( 0, 0, 1, 1, fill="#17273a", outline="#bed2eb", width=1, ) horizontal_axis = canvas.create_line(0, 0, 1, 1, fill="#52708f", width=1) vertical_axis = canvas.create_line(0, 0, 1, 1, fill="#52708f", width=1) direction_items = { direction: canvas.create_text( 0, 0, text=label, fill="#dcecff", font=("Sans", 7, "bold"), ) for direction, label in ( ("forward", "F"), ("left", "L"), ("right", "R"), ("back", "B"), ) } knob_item = canvas.create_oval( 0, 0, 1, 1, fill="#65d9ff", outline="#effaff", width=1, ) mode_bounds = {"left": 126, "top": 0, "right": 254, "bottom": 0} mode_state = {"training": False} mode_button_item = canvas.create_rectangle( 0, 0, 1, 1, fill="#162b42", outline="#8fcff3", width=1, ) mode_button_text = canvas.create_text( 0, 0, text="MANUAL CONTROL\nclick: TRAIN LIVE", fill="#e7f6ff", font=("Sans", 8, "bold"), justify=tk.CENTER, ) def refresh_mode_controls() -> None: if not args.training_live_enabled: canvas.itemconfigure( mode_button_item, fill="#20252d", outline="#66707c", ) canvas.itemconfigure( mode_button_text, text="TRAIN LIVE\nbridge offline", fill="#8f99a5", ) elif mode_state["training"]: canvas.itemconfigure( mode_button_item, fill="#124b38", outline="#77f0b5", ) canvas.itemconfigure( mode_button_text, text="TRAIN LIVE • DR\nclick: MANUAL", fill="#dffff0", ) else: canvas.itemconfigure( mode_button_item, fill="#162b42", outline="#8fcff3", ) canvas.itemconfigure( mode_button_text, text="MANUAL CONTROL\nclick: TRAIN LIVE", fill="#e7f6ff", ) joystick_enabled = not mode_state["training"] canvas.itemconfigure( pad_item, fill="#17273a" if joystick_enabled else "#141a20", outline="#bed2eb" if joystick_enabled else "#58626d", ) canvas.itemconfigure( knob_item, fill="#65d9ff" if joystick_enabled else "#56616c", outline="#effaff" if joystick_enabled else "#737d87", ) for item in direction_items.values(): canvas.itemconfigure( item, fill="#dcecff" if joystick_enabled else "#64707c", ) def layout_controls(canvas_height: int) -> None: center_x = 64 center_y = max(76, canvas_height - 62) pad_center["x"] = center_x pad_center["y"] = center_y panel_top = center_y - pad_radius - 37 canvas.coords( panel_item, 8, panel_top, 265, center_y + pad_radius + 8, ) canvas.coords(label_item, 14, panel_top + 11) canvas.coords(value_item, 14, panel_top + 25) canvas.coords( pad_item, center_x - pad_radius, center_y - pad_radius, center_x + pad_radius, center_y + pad_radius, ) canvas.coords( horizontal_axis, center_x - pad_radius + 8, center_y, center_x + pad_radius - 8, center_y, ) canvas.coords( vertical_axis, center_x, center_y - pad_radius + 8, center_x, center_y + pad_radius - 8, ) canvas.coords(direction_items["forward"], center_x, center_y - 34) canvas.coords(direction_items["left"], center_x - 34, center_y) canvas.coords(direction_items["right"], center_x + 34, center_y) canvas.coords(direction_items["back"], center_x, center_y + 34) canvas.coords( knob_item, center_x - knob_radius, center_y - knob_radius, center_x + knob_radius, center_y + knob_radius, ) mode_bounds["top"] = center_y - 32 mode_bounds["bottom"] = center_y + 32 canvas.coords( mode_button_item, mode_bounds["left"], mode_bounds["top"], mode_bounds["right"], mode_bounds["bottom"], ) canvas.coords( mode_button_text, (mode_bounds["left"] + mode_bounds["right"]) / 2, center_y, ) layout_controls(display_height) refresh_mode_controls() drag = { "mode": None, "last_x": 0, "last_y": 0, } def in_joystick(x: int, y: int) -> bool: if mode_state["training"]: return False return ( (x - pad_center["x"]) ** 2 + (y - pad_center["y"]) ** 2 <= (pad_radius + 6) ** 2 ) def in_mode_button(x: int, y: int) -> bool: return ( mode_bounds["left"] <= x <= mode_bounds["right"] and mode_bounds["top"] <= y <= mode_bounds["bottom"] ) def send_command(command: str) -> None: sys.stdout.write(f"{command}\n") sys.stdout.flush() def update_joystick(x: int, y: int) -> None: delta_x = float(x - pad_center["x"]) delta_y = float(y - pad_center["y"]) magnitude = (delta_x**2 + delta_y**2) ** 0.5 if magnitude > pad_radius: scale = pad_radius / magnitude delta_x *= scale delta_y *= scale forward = -delta_y / pad_radius left = -delta_x / pad_radius radius_fraction = min(1.0, magnitude / pad_radius) side_quadrant = abs(delta_x) > abs(delta_y) rim_turn = max(0.0, min(1.0, (radius_fraction - 0.65) / 0.35)) yaw_left = ( max(-1.0, min(1.0, left)) * rim_turn if side_quadrant else 0.0 ) canvas.coords( knob_item, pad_center["x"] + delta_x - knob_radius, pad_center["y"] + delta_y - knob_radius, pad_center["x"] + delta_x + knob_radius, pad_center["y"] + delta_y + knob_radius, ) canvas.itemconfigure( value_item, text=( f"vx {forward * args.joystick_max_speed:+.2f} " f"vy {left * args.joystick_max_speed:+.2f} " f"wz {yaw_left * args.joystick_max_yaw_rate:+.2f}" ), ) send_command( f"robot_velocity {forward:.4f} {left:.4f} {yaw_left:.4f}" ) def on_motion(event: tk.Event) -> None: if drag["mode"] == "joystick": update_joystick(int(event.x), int(event.y)) canvas.configure(cursor="hand2") return if drag["mode"] == "camera": delta_x = int(event.x) - int(drag["last_x"]) delta_y = int(event.y) - int(drag["last_y"]) drag["last_x"] = int(event.x) drag["last_y"] = int(event.y) if delta_x != 0 or delta_y != 0: # Horizontal orbit follows the user's preferred drag direction. send_command( "camera_orbit_delta " f"{-0.28 * delta_x:.3f} {0.22 * delta_y:.3f}" ) canvas.configure(cursor="fleur") return canvas.configure( cursor=( "hand2" if in_joystick(event.x, event.y) or in_mode_button(event.x, event.y) else "" ) ) def on_press(event: tk.Event) -> None: if in_mode_button(event.x, event.y): drag["mode"] = "mode_button" if args.training_live_enabled: mode_state["training"] = not mode_state["training"] send_command( "view_mode " f"{'training' if mode_state['training'] else 'manual'}" ) refresh_mode_controls() canvas.configure(cursor="hand2") return if in_joystick(event.x, event.y): drag["mode"] = "joystick" update_joystick(int(event.x), int(event.y)) canvas.configure(cursor="hand2") return drag["mode"] = "camera" drag["last_x"] = int(event.x) drag["last_y"] = int(event.y) canvas.configure(cursor="fleur") def on_release(_event: tk.Event) -> None: if drag["mode"] == "joystick": send_command("robot_velocity 0.0 0.0 0.0") canvas.coords( knob_item, pad_center["x"] - knob_radius, pad_center["y"] - knob_radius, pad_center["x"] + knob_radius, pad_center["y"] + knob_radius, ) canvas.itemconfigure( value_item, text="vx +0.00 vy +0.00 wz +0.00", ) drag["mode"] = None canvas.configure(cursor="") def close_window(_event: tk.Event | None = None) -> None: stopped.set() root.destroy() canvas.bind("", on_motion) canvas.bind("", lambda _event: canvas.configure(cursor="")) canvas.bind("", on_press) canvas.bind("", on_motion) canvas.bind("", on_release) root.bind("", close_window) root.bind("q", close_window) root.protocol("WM_DELETE_WINDOW", close_window) def read_exact(stream, buffer: bytearray) -> bool: buffer_view = memoryview(buffer) offset = 0 while offset < len(buffer): chunk = stream.readinto(buffer_view[offset:]) if not chunk: return False offset += chunk return True def enqueue_frame(frame_bytes: bytes) -> None: try: frames.put_nowait(frame_bytes) except queue.Full: try: frames.get_nowait() except queue.Empty: pass try: frames.put_nowait(frame_bytes) except queue.Full: pass def read_source_frames() -> None: frame_buffer = bytearray(expected_bytes) while not stopped.is_set(): if not read_exact(sys.stdin.buffer, frame_buffer): stopped.set() return enqueue_frame(bytes(frame_buffer)) threading.Thread(target=read_source_frames, daemon=True).start() photo_reference: dict[str, ImageTk.PhotoImage | None] = {"image": None} transition: dict[str, np.ndarray | float | bool | None] = { "previous": None, "current": None, "displayed": None, "forward_flow": None, "backward_flow": None, "started_at": 0.0, "scene_cut": False, } transition_duration_s = 1.0 / args.source_fps grid_x, grid_y = np.meshgrid( np.arange(args.width, dtype=np.float32), np.arange(args.height, dtype=np.float32), ) redraw = {"needed": True} def estimate_bidirectional_flow( previous: np.ndarray, current: np.ndarray, ) -> tuple[np.ndarray | None, np.ndarray | None, bool]: """Estimate scene motion at half scale and detect discontinuous resets.""" small_width = max(160, args.width // 2) small_height = max(90, args.height // 2) previous_gray = cv2.resize( cv2.cvtColor(previous, cv2.COLOR_RGB2GRAY), (small_width, small_height), interpolation=cv2.INTER_AREA, ) current_gray = cv2.resize( cv2.cvtColor(current, cv2.COLOR_RGB2GRAY), (small_width, small_height), interpolation=cv2.INTER_AREA, ) mean_difference = float( cv2.absdiff(previous_gray, current_gray).mean() ) try: estimator = cv2.DISOpticalFlow_create( cv2.DISOPTICAL_FLOW_PRESET_FAST ) forward_small = estimator.calc( previous_gray, current_gray, None, ) backward_small = estimator.calc( current_gray, previous_gray, None, ) except (AttributeError, cv2.error): return None, None, mean_difference > 42.0 scale_x = args.width / small_width scale_y = args.height / small_height forward = cv2.resize( forward_small, (args.width, args.height), interpolation=cv2.INTER_LINEAR, ) backward = cv2.resize( backward_small, (args.width, args.height), interpolation=cv2.INTER_LINEAR, ) forward[..., 0] *= scale_x forward[..., 1] *= scale_y backward[..., 0] *= scale_x backward[..., 1] *= scale_y magnitude_95 = float( np.percentile( np.linalg.norm(forward, axis=2), 95, ) ) scene_cut = mean_difference > 42.0 or magnitude_95 > 120.0 return forward, backward, scene_cut def begin_transition(frame_bytes: bytes) -> None: current = np.frombuffer(frame_bytes, dtype=np.uint8).reshape( args.height, args.width, 3, ).copy() # Always interpolate between two genuine RTX frames. Feeding a # synthetic in-between image back into the next flow estimate causes # recursive trails and eventually multiple translucent robots. previous = transition["current"] if not isinstance(previous, np.ndarray): transition["previous"] = current transition["current"] = current transition["displayed"] = current transition["forward_flow"] = None transition["backward_flow"] = None transition["scene_cut"] = True transition["started_at"] = time.perf_counter() return forward, backward, scene_cut = estimate_bidirectional_flow( previous, current, ) transition["previous"] = previous transition["current"] = current transition["forward_flow"] = forward transition["backward_flow"] = backward transition["scene_cut"] = scene_cut transition["started_at"] = time.perf_counter() def interpolated_frame(now: float) -> np.ndarray | None: previous = transition["previous"] current = transition["current"] if not isinstance(current, np.ndarray): return None if not isinstance(previous, np.ndarray) or bool( transition["scene_cut"] ): return current alpha = min( 1.0, max( 0.0, (now - float(transition["started_at"])) / transition_duration_s, ), ) forward = transition["forward_flow"] backward = transition["backward_flow"] if ( alpha >= 1.0 or not isinstance(forward, np.ndarray) or not isinstance(backward, np.ndarray) ): return current previous_map_x = grid_x - alpha * forward[..., 0] previous_map_y = grid_y - alpha * forward[..., 1] current_map_x = grid_x - (1.0 - alpha) * backward[..., 0] current_map_y = grid_y - (1.0 - alpha) * backward[..., 1] previous_warped = cv2.remap( previous, previous_map_x, previous_map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE, ) current_warped = cv2.remap( current, current_map_x, current_map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE, ) blended = cv2.addWeighted( previous_warped, 1.0 - alpha, current_warped, alpha, 0.0, ) # HUD text and the loss plot are semantic overlays, not scene # geometry. Keep them crisp instead of optical-flow warping glyphs. left_width = min(args.width, round(args.width * 0.38)) left_height = min(args.height, round(args.height * 0.72)) right_start = max(0, round(args.width * 0.60)) top_height = min(args.height, round(args.height * 0.22)) graph_start_y = max(0, round(args.height * 0.76)) blended[:left_height, :left_width] = current[ :left_height, :left_width, ] blended[:top_height, right_start:] = current[ :top_height, right_start:, ] blended[graph_start_y:, right_start:] = current[ graph_start_y:, right_start:, ] return blended def on_resize(event: tk.Event) -> None: if event.widget is not canvas: return layout_controls(max(int(event.height), 1)) redraw["needed"] = True canvas.bind("", on_resize) def refresh() -> None: if stopped.is_set(): close_window() return newest: bytes | None = None while True: try: newest = frames.get_nowait() except queue.Empty: break if newest is not None: begin_transition(newest) redraw["needed"] = True current = transition["current"] if isinstance(current, np.ndarray): transition_alpha = min( 1.0, max( 0.0, ( time.perf_counter() - float(transition["started_at"]) ) / transition_duration_s, ), ) if transition_alpha < 1.0: redraw["needed"] = True if redraw["needed"] and isinstance(current, np.ndarray): canvas_width = max(canvas.winfo_width(), 1) canvas_height = max(canvas.winfo_height(), 1) display_frame = interpolated_frame(time.perf_counter()) if display_frame is None: root.after( max(1, round(1000.0 / args.display_fps)), refresh, ) return transition["displayed"] = display_frame image = Image.fromarray(display_frame, mode="RGB") scale = min( canvas_width / args.width, canvas_height / args.height, ) fitted_width = max(1, round(args.width * scale)) fitted_height = max(1, round(args.height * scale)) if fitted_width != args.width or fitted_height != args.height: image = image.resize( (fitted_width, fitted_height), Image.Resampling.BILINEAR, ) photo_reference["image"] = ImageTk.PhotoImage(image) image_x = (canvas_width - fitted_width) // 2 image_y = (canvas_height - fitted_height) // 2 canvas.coords(image_item, image_x, image_y) canvas.itemconfigure( image_item, image=photo_reference["image"], ) canvas.tag_lower(image_item) redraw["needed"] = False root.after(max(1, round(1000.0 / args.display_fps)), refresh) root.after(0, refresh) root.mainloop() return 0 if __name__ == "__main__": raise SystemExit(main())