scripts/dropbear_live_window.py

#!/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("<Motion>", on_motion)
    canvas.bind("<Leave>", lambda _event: canvas.configure(cursor=""))
    canvas.bind("<ButtonPress-1>", on_press)
    canvas.bind("<B1-Motion>", on_motion)
    canvas.bind("<ButtonRelease-1>", on_release)
    root.bind("<Escape>", 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("<Configure>", 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())